#include "application_integrator.h"
#include "algorithms/parallel_for.h"
#include "imgui.h"
#include "math/vec3fa.h"
#include "random_sampler.hpp"
#include "random_sampler_wrapper.hpp"
#include "sampler.h"
#include "tasking/taskschedulerinternal.h"
#include <atomic>
#include <mutex>

ApplicationIntegrator::ApplicationIntegrator(int argc, char **argv,
                                             const std::string &name)
    : Application(argc, argv, name) {
  resetRender();
}

void ApplicationIntegrator::drawGUI() {

  bool bDirty = false;


  if (ImGui::SliderInt("chain lengths", &chain_lengths, 1, 200000))
    resetRender();

  if (ImGui::SliderFloat("small step size", &small_step_size, 0.0, 0.1))
    resetRender();

  if (ImGui::SliderFloat("large step probability", &large_step_probability, 0.0,
                         1.0))
    resetRender();

  if (ImGui::Checkbox("Metropolis", &bMetropolis)) {
    resetRender();
  }

  if (bDirty) {
    resetRender();
  }
}

inline float luminance(Vec3fa v) {
  return 0.2126f * v.x + 0.7152f * v.y + 0.0722f * v.z;
}

void ApplicationIntegrator::resetRender() {
  Application::resetRender();

  accepted.store(0, std::memory_order_relaxed);
  luminance_count.store(0, std::memory_order_relaxed);
  std::lock_guard<std::mutex> guard(large_step_global_luminance_mutex);
  large_step_global_luminance = 0.0;

  for (size_t i = 0; i < NUM_CHAINS; i++) {
    chains[i] = MLTRandomSampler(small_step_size, large_step_probability);
    chains[i].init(i);
  }
  if (bMetropolis) {
    data.film.count = false;
  } else {
    data.film.count = true;
    data.film.scalar = 1.0;
  }
}

void ApplicationIntegrator::render(int *pixels, int width, int height,
                                   float time, const ISPCCamera &camera) {
  deviceRender(camera);

  if (!bMetropolis) {
    mcRender(pixels, width, height, time, camera);
  } else {
    mltRender(pixels, width, height, time, camera);
  }
}


void ApplicationIntegrator::mltRender(int *pixels, int width, int height,
                                      float time, const ISPCCamera &camera) {

  // data.film.scalar = ... use it for setting up the correct normalization
  // coefficient
  //
  //
  // you may want to use Distribution1D for the bootstrap
  // d = Distribution1D(float* bis_values, num_bins)
  // float integral = d.funcInt;
  // int index_of_the_sampled_bin = d.SampleDiscrete(rng.get1D());
  
  
  parallel_for(size_t(0), size_t(NUM_CHAINS), [&](const range<size_t> &range) {
    const int threadIndex = (int)TaskScheduler::threadIndex();
    float large_luminance_outer_sum = 0.0;
    for (size_t i = range.begin(); i < range.end(); i++) {

      float large_luminance_sum = 0.0;

      for (size_t j = 0; j < chain_lengths; j++) {

        chains[i].new_ray(RandomSampler_get1D(chains[i].sampler) < large_step_probability);

        float x = chains[i].get1D() * width;
        float y = chains[i].get1D() * height;

        int x_pixel = x;
        int y_pixel = y;

        Vec3f f = renderPixel(x, y, camera, g_stats[threadIndex], chains[i]);
        float l = luminance(f);

        if (chains[i].is_large_step()) {
          large_luminance_sum += l;
          luminance_count.fetch_add(1, std::memory_order_relaxed);
        }

        if ((last_l[i] == 0.0 && l > 0.0) || (last_l[i] > 0.0 && RandomSampler_get1D(chains[i].sampler) < l / last_l[i])) {
          data.film.addSplat(x_pixel, y_pixel, f / l);
          // printf("accept %d\n", chains[i].is_large_step());
          chains[i].accept();
          accepted.fetch_add(1, std::memory_order_relaxed);
          last_l[i] = l;
        }

      }
      large_luminance_outer_sum += large_luminance_sum;
    }
    std::lock_guard<std::mutex> guard(large_step_global_luminance_mutex);
    large_step_global_luminance += large_luminance_outer_sum;
  });


  
  
  std::lock_guard<std::mutex> guard(large_step_global_luminance_mutex);
  data.film.scalar = (float)(width * height) * large_step_global_luminance / (accepted.load(std::memory_order_relaxed) * luminance_count.load(std::memory_order_relaxed));
  

  printf("%f\n", data.film.scalar);
}

void ApplicationIntegrator::mcRender(int *pixels, int width, int height,
                                     float time, const ISPCCamera &camera) {
  const int numTilesX = (width + TILE_SIZE_X - 1) / TILE_SIZE_X;
  const int numTilesY = (height + TILE_SIZE_Y - 1) / TILE_SIZE_Y;
  parallel_for(size_t(0), size_t(numTilesX * numTilesY),
               [&](const range<size_t> &range) {
                 const int threadIndex = (int)TaskScheduler::threadIndex();
                 for (size_t i = range.begin(); i < range.end(); i++)
                   renderTile((int)i, threadIndex, pixels, width, height, time,
                              camera, numTilesX, numTilesY);
               });

}

/* renders a single screen tile */
void ApplicationIntegrator::mcRenderTile(int taskIndex, int threadIndex,
                                         int *pixels, const unsigned int width,
                                         const unsigned int height, const float time,
                                         const ISPCCamera &camera,
                                         const int numTilesX,
                                         const int numTilesY) {
  const unsigned int tileY = taskIndex / numTilesX;
  const unsigned int tileX = taskIndex - tileY * numTilesX;
  const unsigned int x0 = tileX * TILE_SIZE_X;
  const unsigned int x1 = min(x0 + TILE_SIZE_X, width);
  const unsigned int y0 = tileY * TILE_SIZE_Y;
  const unsigned int y1 = min(y0 + TILE_SIZE_Y, height);

  for (unsigned int y = y0; y < y1; y++)
    for (unsigned int x = x0; x < x1; x++) {
      RandomSamplerWrapper sampler;
      Vec3fa L = Vec3fa(0.0f);

      for (int i = 0; i < data.spp; i++) {
        sampler.init(x, y, (data.frame_count) * data.spp + i);

        /* calculate pixel color */
        float fx = x + sampler.get1D();
        float fy = y + sampler.get1D();
        L = L + renderPixel(fx, fy, camera, g_stats[threadIndex], sampler);
      }
      L = L / (float)data.spp;

      /* write color to framebuffer */
      data.film.addSplat(x, y, L);
    }
}