#include "application_integrator.h"
#include "algorithms/parallel_for.h"
#include "distribution.hpp"
#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 <cmath>
#include <mutex>
#include <vector>

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

void ApplicationIntegrator::drawGUI() {

  bool bDirty = false;

  if (ImGui::Checkbox("keep chains", &keep_chains))
    resetRender();

  if (ImGui::Checkbox("bootstrap", &bootstrap))
    resetRender();

  if (ImGui::SliderInt("Bootstrap amount", &bootstrap_amount, 1, 2000000))
    resetRender();

  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_count.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;
  frame_count = 0;

  for (size_t i = 0; i < NUM_CHAINS; i++) {
    chains[i] = MLTRandomSampler(small_step_size);
    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());

  if (bootstrap) {
    if (frame_count == 0 || !keep_chains) {

      std::vector<MLTRandomSampler> bootstrap_chains(bootstrap_amount);
      std::vector<float> bootstrap_l(bootstrap_amount);

      parallel_for(size_t(0), size_t(bootstrap_amount), [&](const range<size_t> &range) {
        const int threadIndex = (int)TaskScheduler::threadIndex();

        double luminance_sum = 0.0;
        
        for (size_t i = range.begin(); i < range.end(); i++) {
          bootstrap_chains.at(i) = MLTRandomSampler(small_step_size);
          bootstrap_chains.at(i).init(frame_count * bootstrap_amount + i);

          bootstrap_chains.at(i).new_ray(true);

          float x = bootstrap_chains.at(i).get1D() * width;
          float y = bootstrap_chains.at(i).get1D() * height;

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

          luminance_sum += l;

          bootstrap_l.at(i) = l;
          
        }                    
        
        std::lock_guard<std::mutex> guard(large_step_global_luminance_mutex);
        large_step_global_luminance += luminance_sum;
      });

      luminance_count.fetch_add(bootstrap_amount);

      auto d = Distribution1D(bootstrap_l.data(), bootstrap_amount);

      float integral = d.funcInt;

      for (size_t i = 0; i < NUM_CHAINS; i++) {
        int index = d.SampleDiscrete(RandomSampler_get1D(bootstrap_chains.at(0).sampler));
        chains[i] = bootstrap_chains.at(index);
      }
      
    }
    
  } else if (!keep_chains) {
    for (size_t i = 0; i < NUM_CHAINS; i++) {
      chains[i] = MLTRandomSampler(small_step_size);
      chains[i].init(NUM_CHAINS * frame_count + i);
    }
  }
  

  parallel_for(size_t(0), size_t(NUM_CHAINS), [&](const range<size_t> &range) {
    const int threadIndex = (int)TaskScheduler::threadIndex();
    double large_luminance_outer_sum = 0.0;
    for (size_t i = range.begin(); i < range.end(); i++) {

      double 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);
        }


        float a;
        if (last_l[i] > 0) {
          a = std::fmin(l / last_l[i], 1.0);
        } else {
          a = 1.0;
        }

        if (l > 0.0) {
          data.film.addSplat(x_pixel, y_pixel, (f / l) * a);
          if (RandomSampler_get1D(chains[i].sampler) < a) {
            accepted_count.fetch_add(1, std::memory_order_relaxed);
            // if (i == 0)
            //   printf("%d accept %d %f %f %f\n", i, chains[i].is_large_step(), last_l[i], l, l / last_l[i]);
            chains[i].accept();
            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;
  });


  

  frame_count++;
  std::lock_guard<std::mutex> guard(large_step_global_luminance_mutex);
  double mean_image_brightness = large_step_global_luminance / luminance_count.load(std::memory_order_relaxed);
  double normalization = (double)(width * height) / accepted_count.load(std::memory_order_relaxed);
  data.film.scalar = mean_image_brightness * normalization;
  

  printf("%zu, %f, %f, %zu, %f\n", frame_count, data.film.scalar, mean_image_brightness, luminance_count.load(std::memory_order_relaxed), normalization);
}

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);
    }
}