#include "Object.h"
#include "Ray.h"
#include "Camera.h"
#include "Sphere.h"
#include "Light.h"
#include "Plane.h"
#include "read_json.h"
#include "viewing_ray.h"
#include "first_hit.h"
#include "raycolor.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:   ./raytracing <path to json file> \n"
      "  Windows: raytracing.exe <path to json file> \n"
      "For example: \n"
      "  Linux:   ./raytracing ../data/sphere-and-plane.json \n"
      "  Windows: raytracing.exe ../../../data/sphere-and-plane.json \n"
  	);
  	return -1;
  }
  
  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;
  std::vector< std::shared_ptr<Light> > lights;
  // Read a camera and scene description from given .json file
  read_json(path_to_json, camera, objects, lights);

  int width =  640;
  int height = 360;
  std::vector<unsigned char> depth_image(1 * width * height);
  std::vector<unsigned char> rgb_image(3*width*height);

  // print a loading bar
  printf("[");
  for (int i = 0; i < 36; i++) printf("-");
  printf("]\r[");

  // for depth map
  const double zNear = camera.d;
  double t;
  Eigen::Vector3d n;
  int hit_id;

  // For each pixel (i,j)
  for(unsigned i=0; i<height; ++i) 
  {
    for(unsigned j=0; j<width; ++j)
    {
      // Set background color
	  depth_image[j + width * i] = 0;
      Eigen::Vector3d rgb(0,0,0);

      // Compute viewing ray
      Ray ray;
      viewing_ray(camera,i,j,width,height,ray);
      
      // Shoot ray and collect color
      raycolor(ray,1.0,objects,lights,0,rgb);

      // Write double precision color into image
      auto clamp = [](double s){ return std::max(std::min(s,1.0),0.0);};
      rgb_image[0+3*(j+width*i)] = 255.0*clamp(rgb(0));
      rgb_image[1+3*(j+width*i)] = 255.0*clamp(rgb(1));
      rgb_image[2+3*(j+width*i)] = 255.0*clamp(rgb(2));

	  // depth image
	  if (first_hit(ray, 1.0, objects, hit_id, t, n))
	  {
        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));
	  }

    }
	if (i % 10 == 0) printf("="); // fill the loading bar
  }

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