#include "application_integrator.h"
#include "algorithms/parallel_for.h"
#include "imgui.h"
#include "math/vec2.h"
#include "math/vec3fa.h"
#include "random_sampler.hpp"
#include "random_sampler_wrapper.hpp"
#include "tasking/taskschedulerinternal.h"
#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::SliderInt("num chains", &num_chains, 1, 2000))
    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();

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

class MLTRandomSampler : public RandomSamplerWrapper {
private:
  size_t index;
  std::vector<float> data;
  std::vector<float> new_data;
  std::vector<size_t> last_changed;
  size_t time;
  size_t last_large_step;
  float small_step_size;
  float large_step_probability;
  bool large_step;

  float normalize(float x) {
    if (x < 0.0) {
      return x + 1.0;
    } else if (x >= 1.0) {
      return x - 1.0;
    }
    return x;
  }

public:
  MLTRandomSampler(float small_step_size, float large_step_probability)
      : index(0), data({}), last_changed({}), time(0), last_large_step(0),
        small_step_size(small_step_size),
        large_step_probability(large_step_probability) {}

  void init(int id) override { RandomSampler_init(sampler, id); }

  void accept() {
    time++;
    for (size_t i = 0; i < new_data.size(); i++) {
      if (i >= data.size()) {
        data.push_back(new_data[i]);
        last_changed.push_back(time);
      } else {
        data.at(i) = new_data.at(i);
        last_changed.at(i) = time;
      }
    }
    new_data.clear();
  }
  
  void new_ray(bool l) {
    large_step = l;
  }

  bool is_large_step() { return large_step; }

  float get1D() override {
    if (is_large_step()) {
      float r = RandomSampler_get1D(sampler);
      new_data.push_back(r);
      index++;
      return r;
    } else {
      if (index >= data.size()) {
        float r = RandomSampler_get1D(sampler);
        data.push_back(r);
        last_changed.push_back(time);
        new_data.push_back(r);
        index++;
        return r;
      } else if (last_changed.at(index) < last_large_step) {
        float r = RandomSampler_get1D(sampler);
        data.at(index) = r;
        last_changed.at(index) = time;
        new_data.push_back(r);
        index++;
        return r;
      } else {
        size_t steps = time - last_changed.at(index);
        float d = data.at(index);
        for (size_t i = 0; i < steps; i++) {
          float r = RandomSampler_get1D(sampler);
          float o = r * small_step_size - (small_step_size / 2.0);
          d = normalize(d + o);
        }
        data.at(index) = d;
        last_changed.at(index) = time;
        
        float r = RandomSampler_get1D(sampler);
        float o = r * small_step_size - (small_step_size / 2.0);
        d = normalize(d + o);

        new_data.push_back(d);
      
        return d;
      }
    }
  }

  Vec2f get2D() override { return Vec2f(get1D(), get1D()); }

  Vec3fa get3D() override { return Vec3fa(get1D(), get1D(), get1D()); }
};

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();
    for (size_t i = range.begin(); i < range.end(); i++) {
      MLTRandomSampler sampler(small_step_size, large_step_probability);
      sampler.init(data.frame_count * num_chains + i);

      float last_l = 0.0;

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

        sampler.new_ray(RandomSampler_get1D(sampler.sampler) < large_step_probability);

        float x = sampler.get1D() * width;
        float y = sampler.get1D() * height;

        int x_pixel = x;
        int y_pixel = y;

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

        if ((last_l == 0.0 && l > 0.0) || (last_l > 0.0 && RandomSampler_get1D(sampler.sampler) < l / last_l)) {
          data.film.addSplat(x_pixel, y_pixel, f / l);
          sampler.accept();
        }

      }
    }
  });

}

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