// Copyright 2009 Intel Corporation
// SPDX-License-Identifier: Apache-2.0

/* This is a small example tutorial how to use OSPRay and the
 * MPIDistributedDevice in a data-parallel application.
 * Each rank must specify the same render parameters, however the data
 * to render on each rank can differ for distributed rendering. In this
 * tutorial each rank renders a unique quad, which is colored by the rank.
 *
 * On Linux build it in the build_directory with:
 *   mpicxx ../modules/mpi/tutorials/ospMPIDistributedTutorial.cpp \
 *       -I ../ospray/include -I ../components \
 *       -L . -lospray -lospray_common -Wl,-rpath,. \
 *       -o ospMPIDistributedTutorialCpp
 *
 * Then run it with MPI on some number of processes
 *   mpirun -n <N> ./ospMPIDistributedTutorialCpp
 *
 * The output image should show a sequence of quads, from dark to light blue
 */

#include <errno.h>
#include <mpi.h>
#include <stdint.h>
#include <stdio.h>
#ifdef _WIN32
#define NOMINMAX
#include <malloc.h>
#else
#include <alloca.h>
#endif

#include "rkcommon/math/box.h"
#include "rkcommon/math/vec.h"
#include "rkcommon/utility/SaveImage.h"
#include "rkcommon/utility/getEnvVar.h"
// Note: we define OSPRAY_CPP_RKCOMMON_TYPES in CMAke to use rkcommon types
// natively through the C++ wrappers
#include "ospray/ospray_cpp.h"
#include "ospray/ospray_cpp/ext/rkcommon.h"

using namespace ospray;
using namespace rkcommon;
using namespace rkcommon::math;

int main(int argc, char **argv)
{
  int mpiThreadCapability = 0;
  MPI_Init_thread(&argc, &argv, MPI_THREAD_MULTIPLE, &mpiThreadCapability);
  if (mpiThreadCapability != MPI_THREAD_MULTIPLE
      && mpiThreadCapability != MPI_THREAD_SERIALIZED) {
    fprintf(stderr,
        "OSPRay requires the MPI runtime to support thread "
        "multiple or thread serialized.\n");
    return 1;
  }

  int mpiRank = 0;
  int mpiWorldSize = 0;
  MPI_Comm_rank(MPI_COMM_WORLD, &mpiRank);
  MPI_Comm_size(MPI_COMM_WORLD, &mpiWorldSize);

  // image size
  vec2i imgSize;
  imgSize.x = 1024; // width
  imgSize.y = 768; // height

  // camera
  vec3f cam_pos{(mpiWorldSize + 1.f) / 2.f, 0.5f, -mpiWorldSize * 0.5f};
  vec3f cam_up{0.f, 1.f, 0.f};
  vec3f cam_view{0.f, 0.f, 1.f};

  // all ranks specify the same rendering parameters, with the exception of
  // the data to be rendered, which is distributed among the ranks
  // triangle mesh data
  vec3f vertex[] = {vec3f(mpiRank, 0.0f, 3.5f),
      vec3f(mpiRank, 1.0f, 3.0f),
      vec3f(1.0f * (mpiRank + 1.f), 0.0f, 3.0f),
      vec3f(1.0f * (mpiRank + 1.f), 1.0f, 2.5f)};
  vec4f color[] = {vec4f(0.0f, 0.0f, (mpiRank + 1.f) / mpiWorldSize, 1.0f),
      vec4f(0.0f, 0.0f, (mpiRank + 1.f) / mpiWorldSize, 1.0f),
      vec4f(0.0f, 0.0f, (mpiRank + 1.f) / mpiWorldSize, 1.0f),
      vec4f(0.0f, 0.0f, (mpiRank + 1.f) / mpiWorldSize, 1.0f)};
  vec3ui index[] = {vec3ui(0, 1, 2), vec3ui(1, 2, 3)};

  // load the MPI module, and select the MPI distributed device. Here we
  // do not call ospInit, as we want to explicitly pick the distributed
  // device
  auto OSPRAY_MPI_DISTRIBUTED_GPU =
      utility::getEnvVar<int>("OSPRAY_MPI_DISTRIBUTED_GPU").value_or(0);
  if (OSPRAY_MPI_DISTRIBUTED_GPU) {
    ospLoadModule("mpi_distributed_gpu");
  } else {
    ospLoadModule("mpi_distributed_cpu");
  }

  // use scoped lifetimes of wrappers to release everything before ospShutdown()
  {
    cpp::Device mpiDevice("mpiDistributed");
    mpiDevice.commit();
    mpiDevice.setCurrent();

    // create and setup camera
    cpp::Camera camera("perspective");
    camera.setParam("aspect", imgSize.x / (float)imgSize.y);
    camera.setParam("position", cam_pos);
    camera.setParam("direction", cam_view);
    camera.setParam("up", cam_up);
    camera.commit(); // commit each object to indicate modifications are done

    // create and setup model and mesh
    cpp::Geometry mesh("mesh");
    cpp::CopiedData data(vertex, 4);
    data.commit();
    mesh.setParam("vertex.position", data);

    data = cpp::CopiedData(color, 4);
    data.commit();
    mesh.setParam("vertex.color", data);

    data = cpp::CopiedData(index, 2);
    data.commit();
    mesh.setParam("index", data);

    mesh.commit();

    // put the mesh into a model
    cpp::GeometricModel model(mesh);
    model.commit();

    // put the model into a group (collection of models)
    cpp::Group group;
    group.setParam("geometry", cpp::CopiedData(model));
    group.commit();

    // put the group into an instance (give the group a world transform)
    cpp::Instance instance(group);
    instance.commit();

    cpp::World world;
    world.setParam("instance", cpp::CopiedData(instance));

    // Specify the region of the world this rank owns
    box3f regionBounds(
        vec3f(mpiRank, 0.f, 2.5f), vec3f(1.f * (mpiRank + 1.f), 1.f, 3.5f));
    world.setParam("region", cpp::CopiedData(regionBounds));

    world.commit();

    // create the mpi_raycast renderer (required for distributed rendering)
    cpp::Renderer renderer("mpiRaycast");
    renderer.commit();

    // create and setup framebuffer
    cpp::FrameBuffer framebuffer(imgSize.x,
        imgSize.y,
        OSP_FB_SRGBA,
        OSP_FB_COLOR | /*OSP_FB_DEPTH |*/ OSP_FB_ACCUM);
    framebuffer.clear();

    ospray::cpp::PickResult res =
        framebuffer.pick(renderer, camera, world, 0.5f, 0.5f);

    if (res.hasHit) {
      std::cout << "Rank " << mpiRank
                << " picked geometry [instance: " << res.instance.handle()
                << ", model: " << res.model.handle()
                << ", primitive: " << res.primID << "]" << std::endl;
    } else {
      std::cout << "Rank " << mpiRank << " pick missed" << std::endl;
    }

    // render one frame
    framebuffer.renderFrame(renderer, camera, world);

    // on rank 0, access framebuffer and write its content as PPM file
    if (mpiRank == 0) {
      uint32_t *fb = (uint32_t *)framebuffer.map(OSP_FB_COLOR);
      rkcommon::utility::writePPM(
          "firstFrameCpp.ppm", imgSize.x, imgSize.y, fb);
      framebuffer.unmap(fb);
    }

    // render 10 more frames, which are accumulated to result in a better
    // converged image
    for (int frames = 0; frames < 10; frames++)
      framebuffer.renderFrame(renderer, camera, world);

    if (mpiRank == 0) {
      uint32_t *fb = (uint32_t *)framebuffer.map(OSP_FB_COLOR);
      rkcommon::utility::writePPM(
          "accumulatedFrameCpp.ppm", imgSize.x, imgSize.y, fb);
      framebuffer.unmap(fb);
    }
  }

  ospShutdown();

  MPI_Finalize();

  return 0;
}
