#include "Object.h"
#include "Ray.h"
#include "Camera.h"
#include "Sphere.h"
#include "Plane.h"
#include "read_json.h"
#include "viewing_ray.h"
#include "first_hit.h"
#include <Eigen/Core>
#include <vector>
#include <iostream>
#include <memory>
#include <limits>
#include <functional>

// stb is a library commonly used when one wants to save an image in C++,
// since only one or two headers need to be included.
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"


int main(int argc, char * argv[])
{
  // parse command line arguments
  std::string path_to_json;
  if (argc == 2) {
    std::string args = argv[1];
    // check if ends with .json
    if (args.size() > 5 && args.substr(args.size() - 4, 4) == "json") {
      path_to_json = args;
    }
  }
  if (argc != 2 || path_to_json.empty()) {
    printf(
      "Error: received unexpected command line arguments. \n"
      "Correct usage: \n"
      "  Linux:   ./raycasting <path to json file> \n"
      "  Windows: raycasting.exe <path to json file> \n"
      "For example: \n"
      "  Linux:   ./raycasting ../data/sphere-and-plane.json \n"
      "  Windows: raycasting.exe ../../../data/sphere-and-plane.json \n"
    );
    return -1;
  }

  // list of colors used for per-object id coloring
  const std::vector<unsigned char> color_map = {
    228,26,28,
    55,126,184,
    77,175,74,
    152,78,163,
    255,127,0,
    255,255,51,
    166,86,40,
    247,129,191,
    153,153,153
  };

  Camera camera;
  std::vector< std::shared_ptr<Object> > objects;
  // Read a camera and scene description from given .json file
  read_json(path_to_json,camera,objects);

  int width = 640;
  int height = 360;
  std::vector<unsigned char> id_image(3*width*height);
  std::vector<unsigned char> normal_image(3*width*height);
  std::vector<unsigned char> depth_image(1*width*height);
  
  // print a loading bar
  for (int i = 0; i < 36; i++) printf("%c", 176);
  printf("\r");
  
  // For each pixel (i,j)
  for(int i=0; i<height; ++i)
  {
    for(int j=0; j<width; ++j)
    {
      // Set background color
      normal_image[0+3*(j+width*i)] = 0;
      normal_image[1+3*(j+width*i)] = 0;
      normal_image[2+3*(j+width*i)] = 0;
      depth_image[j+width*i] = 0;
      id_image[0+3*(j+width*i)] = 0;
      id_image[1+3*(j+width*i)] = 0;
      id_image[2+3*(j+width*i)] = 0;

      // Compute viewing ray
      Ray ray;
      /* NOTE I flipped the vertical axis of the pixels, because in computer graphics we assume the bottom left pixel
       * is (0,0), while traditional jpg assumes (0,0) to be the upper left pixel. */
      viewing_ray(camera,height - i,j,width,height,ray);

      // Find first visible object hit by ray and its surface normal n
      double t;
      Eigen::Vector3d n;
      int hit_id;
      if(first_hit(ray,1.0,objects,hit_id,t,n))
      {
        // object-id image
        const int color_id = hit_id%(color_map.size()/3);
        id_image[0+3*(j+width*i)] = color_map[0+3*color_id];
        id_image[1+3*(j+width*i)] = color_map[1+3*color_id];
        id_image[2+3*(j+width*i)] = color_map[2+3*color_id];

        // depth image
        const double zNear = camera.d;
        double linearized_depth = zNear/(t*ray.direction.norm());
        linearized_depth = linearized_depth<1?linearized_depth:1;
        depth_image[(j + width * i)] =  (uint8_t)(255.0*(linearized_depth));

        // set pixel color to value computed from hit point, light, and n
        // normal image
        auto normal_to_rgb = [](const Eigen::Vector3d & n, unsigned char & r, unsigned char & g, unsigned char & b)
        {
          r = (unsigned char) (255.0*(n(0)*0.5+0.5));
          g = (unsigned char) (255.0*(n(1)*0.5+0.5));
          b = (unsigned char) (255.0*(n(2)*0.5+0.5));
        };
        normal_to_rgb(
          n,
          normal_image[0+3*(j+width*i)],
          normal_image[1+3*(j+width*i)],
          normal_image[2+3*(j+width*i)]);
      }
    }
    if (i % 10 == 0) printf("%c", 178); // fill the loading bar
  }

  bool ret0 = stbi_write_png("normal.png", width, height, 3, normal_image.data(), width * 3);
  bool ret1 = stbi_write_png("depth.png", width, height, 1, depth_image.data(), width);
  bool ret2 = stbi_write_png("id.png", width, height, 3, id_image.data(), width * 3);
  printf("\nDone! Wrote to normal.png (succes=%d), depth.png (succes=%d) and id.png (succes=%d) \n", ret0, ret1, ret2);
}