auto magic_ponies()

in src/simulator/simulator_bindings.cpp [145:215]
63 lines of code
12 McCabe index (conditional complexity)

auto magic_ponies(const std::vector<unsigned char> &serialized_task,
                  const UserInput &user_input, bool keep_space_around_bodies,
                  int steps, int stride, bool need_images,
                  bool need_featurized_objects, int get_random_images) {
  SimpleTimer timer;
  Task task = deserialize<Task>(serialized_task);
  addUserInputToScene(user_input, keep_space_around_bodies,
                      /*allow_occlusions=*/false, &task.scene);
  auto simulation = simulateTask(task, steps, stride);

  const double simulation_seconds = timer.GetSeconds();
  const bool isSolved = simulation.isSolution;
  const bool hadOcclusions = hadSimulationOcclusions(simulation);

  const int kSampleSize = get_random_images > 0 ? get_random_images : 1000000;

  const int numScenes = simulation.sceneList.size();
  const int numImagesTotal =
      need_images ? std::min<int>(numScenes, 1 + kSampleSize) : 0;
  const int numScenesTotal =
      need_featurized_objects ? simulation.sceneList.size() : 0;

  const int imageSize = task.scene.width * task.scene.height;
  uint8_t *packedImages = new uint8_t[imageSize * numImagesTotal];
  if (numImagesTotal > 0) {
    int writeIndex = 0;
    int sceneIndex = 0;
    const int selectedStart = (numScenes >= 1 + kSampleSize)
        ? 1 + rand() % (numScenes - kSampleSize) : 1;
    for (const Scene &scene : simulation.sceneList) {
      if (sceneIndex == 0 || (sceneIndex >= selectedStart &&
                              sceneIndex < selectedStart + kSampleSize)) {
        renderTo(scene, packedImages + writeIndex);
        writeIndex += imageSize;
      }
      sceneIndex += 1;
    }
  }

  const int numSceneObjects = getNumObjects(simulation);
  float *packedVectorizedBodies =
      new float[numSceneObjects * kObjectFeatureSize * numScenesTotal];
  if (numScenesTotal > 0) {
    int writeIndex = 0;
    for (const Scene &scene : simulation.sceneList) {
      featurizeScene(scene, packedVectorizedBodies + writeIndex);
      writeIndex += kObjectFeatureSize * numSceneObjects;
    }
  }

  // Create a Python object that will free the allocated
  // memory when destroyed:
  py::capsule freeImagesWhenDone(packedImages, [](void *f) {
    auto *foo = reinterpret_cast<uint8_t *>(f);
    delete[] foo;
  });
  py::capsule freeObjectsWhenDone(packedVectorizedBodies, [](void *f) {
    auto *foo = reinterpret_cast<float *>(f);
    delete[] foo;
  });
  auto packedImagesArray =
      py::array_t<uint8_t>({numImagesTotal * imageSize},  // shape
                           {sizeof(uint8_t)}, packedImages, freeImagesWhenDone);
  auto packedObjectsArray = py::array_t<float>(
      {numScenesTotal * numSceneObjects * kObjectFeatureSize},  // shape
      {sizeof(float)}, packedVectorizedBodies, freeObjectsWhenDone);
  const double pack_seconds = timer.GetSeconds();
  return std::make_tuple(isSolved, hadOcclusions, packedImagesArray,
                         packedObjectsArray, numSceneObjects,
                         simulation_seconds, pack_seconds);
}