//============================================================================
//  Copyright (c) Kitware, Inc.
//  All rights reserved.
//  See LICENSE.txt for details.
//
//  This software is distributed WITHOUT ANY WARRANTY; without even
//  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
//  PURPOSE.  See the above copyright notice for more information.
//============================================================================
#include <vtkm/rendering/raytracing/RayTracer.h>

#include <iostream>
#include <math.h>
#include <stdio.h>
#include <vtkm/cont/ArrayHandleUniformPointCoordinates.h>
#include <vtkm/cont/ColorTable.h>
#include <vtkm/cont/Timer.h>

#include <vtkm/rendering/raytracing/Camera.h>
#include <vtkm/rendering/raytracing/Logger.h>
#include <vtkm/rendering/raytracing/RayTracingTypeDefs.h>
#include <vtkm/worklet/DispatcherMapField.h>
#include <vtkm/worklet/WorkletMapField.h>

#include <vtkm/io/ImageReaderPNG.h>

namespace vtkm
{
namespace rendering
{
namespace raytracing
{

namespace detail
{

class SurfaceColor
{
public:
  class Shade : public vtkm::worklet::WorkletMapField
  {
  private:
    vtkm::Vec3f_32 _LightPosition;
    vtkm::Vec3f_32 LightAbmient;
    vtkm::Vec3f_32 LightDiffuse;
    vtkm::Vec3f_32 LightSpecular;
    vtkm::Float32 SpecularExponent;
    vtkm::Vec3f_32 CameraPosition;
    vtkm::Vec3f_32 LookAt;

    std::vector<vtkm::Vec3f_32> Light_pos;
    std::vector<vtkm::Vec3f_32> Light_col;
    std::vector<vtkm::Float32> Light_int;

    vtkm::rendering::Shader this_shader;
    vtkm::cont::DataSet this_Texture;
    std::vector<vtkm::Vec4f_32> pixbuff;

    void (*RandomFunction)(void);

  public:
    VTKM_CONT
    Shade(const vtkm::Vec3f_32& lightPosition,
          const vtkm::Vec3f_32& cameraPosition,
          const vtkm::Vec3f_32& lookAt
    	  ,std::vector<vtkm::Vec3f_32> light_pos,
    	  std::vector<vtkm::Vec3f_32> light_col,
    	  std::vector<vtkm::Float32> light_int
        , void (*func)(void)
        , vtkm::rendering::Shader shader
	  )
      : _LightPosition(lightPosition)
      , CameraPosition(cameraPosition)
      , LookAt(lookAt)
      , Light_pos(light_pos)
      , Light_col(light_col)
      , Light_int(light_int)
      , RandomFunction(func)
      , this_shader(shader)
    {
      //Set up some default lighting parameters for now
      //0TTER
      //
	    std::cout << "THIS IS FINALLY LIGHT BEING PORTED: " << std::endl
		      << light_col[0][0] << " "
		      << light_col[0][1] << " "
		      << light_col[0][2] << " " << std::endl;
	    std::cout << "light pos? " << std::endl 
		      << Light_pos[0][0] << " "  
		      << Light_pos[0][1] << " "  
		      << Light_pos[0][2] << " " << std::endl; 
	    std::cout << "cam pos? " << std::endl 
		      << CameraPosition[0] << " " 
		      << CameraPosition[1] << " " 
		      << CameraPosition[2] << " " << std::endl; 
      LightAbmient[0] = .1f;
      LightAbmient[1] = .1f;
      LightAbmient[2] = .1f;
      LightDiffuse[0] = .7f;
      LightDiffuse[1] = .7f;
      LightDiffuse[2] = .7f;
      LightSpecular[0] = .2f;
      LightSpecular[1] = .2f;
      LightSpecular[2] = .2f;
      SpecularExponent = 32.f;

      //load texture here 
      std::cout <<"loading texture " << std::endl;
      /*
      vtkm::cont::ArrayHandle<vtkm::Vec4f_32> pixels;
      vtkm::cont::DataSet texture;
      vtkm::io::ImageReaderPNG textureReader("data/seamlessGalaxy.png");
      //vtkm::io::ImageReaderPNG textureReader("data/noiseTexture.png");
      texture = textureReader.ReadDataSet();
      texture.GetField("color").GetData().CopyTo(pixels);
      std::cout << pixels.GetNumberOfValues() << "\n"; 
      auto pixelportal = pixels.ReadPortal();
      auto pix = pixelportal.Get(0);
      std::cout << pix << "\n\n";

      for(int i = 0; i < pixels.GetNumberOfValues() - 1; i++){
             pixbuff.push_back(pixelportal.Get(i));
      }
      */
      std::cout << "pixel buffer generated " << std::endl; 
    }

    template <typename Precision>
    float DistributionGGX(vtkm::Vec<Precision, 3>& N,
		          vtkm::Vec<Precision, 3>& H,
			  float roughness) const{
	    float a  = roughness * roughness;
	    float a2 = a * a;
	    float NdotH = vtkm::Max(vtkm::dot(N, H), 0.0);
	    float NdotH2 = NdotH * NdotH;
	    float PI = 3.14159265359;

	    float num = a2;
	    float denom = (NdotH2* (a2 - 1.0) + 1.0);
	    denom = PI * denom * denom;

	    return num / denom;
    }
    float GeometrySchlickGGX(float NdotV, float roughness) const{
    	float r = (roughness + 1.0);
	float k = (r*r) / 8.0;

	float num = NdotV;
	float denom = NdotV * (1.0 - k) + k;

	return num / denom;
    }

    template <typename Precision>
    float GeometrySmith  (vtkm::Vec<Precision, 3>& N,
		          vtkm::Vec<Precision, 3>& V,
		          vtkm::Vec<Precision, 3>& L,
			  float roughness) const{
	    float NdotV = vtkm::Max(vtkm::dot(N, V), 0.0);
	    float NdotL = vtkm::Max(vtkm::dot(N, L), 0.0);
	    float ggx2  = GeometrySchlickGGX(NdotV, roughness);
	    float ggx1  = GeometrySchlickGGX(NdotL, roughness);
	    return ggx1 * ggx2;
    }

    template <typename Precision>
    void fresnelSchlick(float cosTheta,
		        vtkm::Vec<Precision, 3>& F0,
			vtkm::Vec<Precision, 3>& result) const{
	    //clamp
	    float a = 1.0-cosTheta;
	    if(a > 1.0){a = 1.0;}
	    else if( a < 0.0){a = 0.0;}

	    result = F0 + (1.0 - F0) * pow(a, 5.0);
    }

    vtkm::Vec3f_32 mapToColor(float u, float v) const {
	    int width = 1024; 
	    int height = 1024;
	    vtkm::Vec3f_32 fc(0.0,0.0,0.0); 
	    int x = std::floor(u * width); 
	    int y = std::floor(v * height); 
	    fc[0] = pixbuff[y * width + x][0]; 
	    fc[1] = pixbuff[y * width + x][1]; 
	    fc[2] = pixbuff[y * width + x][2]; 
	    return fc; 
    }

    template <typename Precision>
    void TriPlanarTexMapping(
      vtkm::Vec<Precision, 3>& Wpos,
      vtkm::Vec<Precision, 3>& Normal,
      vtkm::Vec3f_32& color) const{
	vtkm::Vec3f_32 blending(0.0, 0.0, 0.0); 
	blending = vtkm::Abs(Normal);

	float scale = 0.15; 

	float xC = Wpos[0] * scale - std::floor(Wpos[0] * scale);
	float yC = Wpos[1] *scale - std::floor(Wpos[1] * scale);
	float zC = Wpos[2] * scale - std::floor(Wpos[2] * scale);

	vtkm::Vec3f_32 xaxis = mapToColor(yC, zC);
	vtkm::Vec3f_32 yaxis = mapToColor(xC, zC);
	vtkm::Vec3f_32 zaxis = mapToColor(xC, yC);

	color = blending[0] * xaxis + blending[1] * yaxis + blending[2] * zaxis;

    }

    using ControlSignature =
      void(FieldIn, FieldIn, FieldIn, FieldIn, WholeArrayInOut, WholeArrayIn);
    using ExecutionSignature = void(_1, _2, _3, _4, _5, _6, WorkIndex);

    template <typename ColorPortalType, typename Precision, typename ColorMapPortalType>
    VTKM_EXEC void operator()(const vtkm::Id& hitIdx,
                              const Precision& scalar,
                              const vtkm::Vec<Precision, 3>& normal,
                              const vtkm::Vec<Precision, 3>& intersection,
                              ColorPortalType& colors,
                              ColorMapPortalType colorMap
                              ,const vtkm::Id& idx
			      ) const
    {
      RandomFunction();

      /*
      vtkm::Vec3f_32 intersection(0.0, 0.0, 0.0);
      intersection[0] = inter[0]/100000.0;
      intersection[1] = inter[1]/100000.0;
      intersection[2] = inter[2]/100000.0;
      */
      vtkm::Vec3f_32 pixelCol(0.0, 0.0, 0.0);
      vtkm::Vec3f_32 texCol(1.0, 1.0, 1.0);
      vtkm::Vec<Precision, 3> pos = intersection; 
      vtkm::Vec<Precision, 3> nor = normal; 

      //TriPlanarTexMapping(pos, nor, texCol); 
      vtkm::Vec<Precision, 4> color;
      vtkm::Vec<Precision, 3> albedo(1.0, 0.75, 0.75);
      vtkm::Vec<Precision, 3> Lo(0.0, 0.0, 0.0);
      vtkm::Vec<Precision, 3> F0 (0.04, 0.04, 0.04);
      vtkm::Id offset = idx * 4;
      Precision PI = 3.14159265359;


      if (hitIdx < 0)
      {
        return;
      }

      color[0] = colors.Get(offset + 0);
      color[1] = colors.Get(offset + 1);
      color[2] = colors.Get(offset + 2);
      color[3] = colors.Get(offset + 3);


      Precision metallic = 0.93;
      Precision roughness = 0.48;
      Precision ao = 1.0;

      // get albedo
      vtkm::Int32 colorMapSize = static_cast<vtkm::Int32>(colorMap.GetNumberOfValues());
      vtkm::Int32 colorIdx = vtkm::Int32(scalar * Precision(colorMapSize - 1));
      colorIdx = vtkm::Max(0, colorIdx);
      colorIdx = vtkm::Min(colorMapSize - 1, colorIdx);

      albedo[0] = colorMap.Get(colorIdx)[0];
      albedo[1] = colorMap.Get(colorIdx)[1];
      albedo[2] = colorMap.Get(colorIdx)[2];

      vtkm::Vec<Precision, 3> N = normal;
      vtkm::Normalize(N);
      //checked up to here
      vtkm::Vec<Precision, 3> V = CameraPosition - intersection;
      vtkm::Normalize(V);

      //linearly interpolate base color by metallic value
      F0[0] = F0[0] + (metallic) * albedo[0];
      F0[1] = F0[1] + (metallic) * albedo[1];
      F0[2] = F0[2] + (metallic) * albedo[2];

   //START LOOP LIGHT HERE
   for(int  li  = 0; li < Light_pos.size(); li++){
      vtkm::Vec<Precision, 3> LightPosition = Light_pos[li];
      vtkm::Vec<Precision, 3> L = LightPosition - intersection;
      vtkm::Normalize(L);
      vtkm::Vec<Precision, 3> H = (V + L);
      vtkm::Normalize(H);

      vtkm::Vec<Precision, 3> distance = (LightPosition - intersection);
      Precision x = distance[0] * distance[0];
      Precision y = distance[1] * distance[1];
      Precision z = distance[2] * distance[2];
      Precision dist = sqrt(x + y + z);

      //if we dont want to treat these lights as point lights, remove attenuation
      Precision attenuation = 1.0/(dist * dist);
      attenuation = 1.0; 
      vtkm::Vec<Precision, 3> radiance = Light_int[li] * attenuation* Light_col[li];

      Precision NDF = DistributionGGX(N, H, roughness);
      Precision G   = GeometrySmith(N, V, L, roughness);
      Precision F_Vec_Par = vtkm::Max(vtkm::dot(H, V), 0.0);
      vtkm::Vec<Precision, 3> F(.0, .0, .0);
      fresnelSchlick(F_Vec_Par, F0, F);

      vtkm::Vec<Precision, 3> kS = F;
      vtkm::Vec<Precision, 3> OneVector(1.0, 1.0, 1.0);
      vtkm::Vec<Precision, 3> kD = (OneVector - kS);
      kD = kD * (1.0 - metallic);

      //ERROR?
      vtkm::Vec<Precision, 3> numerator = NDF * G * F;
      Precision denominator = 4.0 * vtkm::Max(vtkm::dot(N, V), 0.0) *
	                  vtkm::Max(vtkm::dot(N, L), 0.0) + 0.0001;
      vtkm::Vec<Precision, 3> specular  = numerator/denominator;

      Precision NdotL = vtkm::Max(vtkm::dot(N, L), 0.0);


      //Lo += (ambient * albedo / PI + specular) * radiance * NdotL;
      Lo += (kD * albedo / PI + specular) * radiance * NdotL;
      pixelCol += Lo;
   //HERE IS THE END OF LIGHT LOOP
   }
      //this_shader.ShaderFunction(pixelCol, this_shader);
      //vtkm::Vec<Precision, 3> ambient(0.01, 0.01, 0.01);
      /*
      for(int q = 0; q < 3; q++){
	      pixelCol[q] = (texCol[q] * 1.0);  
      }*/
      
      
      vtkm::Vec<Precision, 3> ambient(0.1, 0.1, 0.1);

      ambient = ambient * albedo * ao;
      pixelCol += ambient;

      //color balance  / gamma correction
      for(int i  = 0; i<3; i++){
	      pixelCol[i] = pixelCol[i]/(pixelCol[i] + 1.0);
	      pixelCol[i] = pow(pixelCol[i], 1.0/2.2);
      }

      colors.Set(offset + 0, pixelCol[0]);
      colors.Set(offset + 1, pixelCol[1]);
      colors.Set(offset + 2, pixelCol[2]);
      colors.Set(offset + 3, 1.0);
    }

  }; //class Shade


  class MapScalarToColor : public vtkm::worklet::WorkletMapField
				//describes the type of data will be parallized on the GPU
{
  public:
    VTKM_CONT //Macro that is used to distinguish that function supposed to run on host (CPU)
    MapScalarToColor() {}

    //signature of how code is supposed to be invoked
    using ControlSignature = void(FieldIn, FieldIn, WholeArrayInOut, WholeArrayIn);
    //FieldIn ~> one value
    //WholeArrayInOut -> whole array of values (can be read and written to)
    //WholeArrayIn -> whole array of values (can be read)

    //signature of how code is supposed to be executed
    using ExecutionSignature = void(_1, _2, _3, _4, WorkIndex);

    template <typename ColorPortalType, typename Precision, typename ColorMapPortalType>
    VTKM_EXEC //Macro that is used to distinguish that function in supposed to run in execution environment (GPU)
    void operator() //operator overloading for parenthesis
                    //(overloads anything that looks like a function call that returns void)
    			     (const vtkm::Id& hitIdx,
                              const Precision& scalar,
                              ColorPortalType& colors,
                              ColorMapPortalType colorMap,
                              const vtkm::Id& idx) const
    {

      if (hitIdx < 0)
      {
        return;
      }

      vtkm::Vec<Precision, 4> color;
      vtkm::Id offset = idx * 4;

      vtkm::Int32 colorMapSize = static_cast<vtkm::Int32>(colorMap.GetNumberOfValues());
      vtkm::Int32 colorIdx = vtkm::Int32(scalar * Precision(colorMapSize - 1));

      // clamp color index
      colorIdx = vtkm::Max(0, colorIdx);
      colorIdx = vtkm::Min(colorMapSize - 1, colorIdx);
      color = colorMap.Get(colorIdx);

      colors.Set(offset + 0, color[0]);
      colors.Set(offset + 1, color[1]);
      colors.Set(offset + 2, color[2]);
      colors.Set(offset + 3, color[3]);
    }

  }; //class MapScalarToColor

  template <typename Precision>
  VTKM_CONT void run(Ray<Precision>& rays,
                     vtkm::cont::ArrayHandle<vtkm::Vec4f_32>& colorMap,
                     const vtkm::rendering::raytracing::Camera& camera,
                     bool shade
		     , const std::vector<vtkm::rendering::Light>& lights
         , void (*func)(void)
         , vtkm::rendering::Shader s
		     )
  {
    if (shade)
    {
      // TODO: support light positions
      vtkm::Vec3f_32 scale(2, 2, 2);
      vtkm::Vec3f_32 lightOffset(1, 2, 0);
      vtkm::Vec3f_32 lightPosition = camera.GetPosition();
      //0TTER
      //parse light junk
      std::vector<vtkm::Vec3f_32> light_pos;
      std::vector<vtkm::Vec3f_32> light_col;
      std::vector<vtkm::Float32> light_int;
      for (int i = 0; i < lights.size(); i++){
	      light_pos.push_back(lights[i].GetPosition());
	      std::cout << "light color is: " << lights[i].GetColor() << std::endl;
	      light_col.push_back(lights[i].GetColor());
	      light_int.push_back(lights[i].GetIntensity());
      }

      vtkm::Vec3f_32 lightColor(0.0, 0.0, 0.0);
      //std::cout << "WE ARE SHADING!!!!" << std::endl;
      vtkm::worklet::DispatcherMapField<Shade>(
        Shade(lightPosition, camera.GetPosition(), camera.GetLookAt(), light_pos, light_col, light_int, func, s))
        .Invoke(rays.HitIdx,
                rays.Scalar,
                rays.Normal,
                rays.Intersection,
                rays.Buffers.at(0).Buffer,
                colorMap
	);

    }
    else
    {
      vtkm::worklet::DispatcherMapField<MapScalarToColor>(MapScalarToColor())
        .Invoke(rays.HitIdx, rays.Scalar, rays.Buffers.at(0).Buffer, colorMap);
    }
  }
}; // class SurfaceColor

} // namespace detail

RayTracer::RayTracer()
  : NumberOfShapes(0)
  , Shade(true)
{
}

RayTracer::~RayTracer()
{
  Clear();
}

Camera& RayTracer::GetCamera()
{
  return camera;
}


void RayTracer::AddShapeIntersector(std::shared_ptr<ShapeIntersector> intersector)
{
  NumberOfShapes += intersector->GetNumberOfShapes();
  Intersectors.push_back(intersector);
}

void RayTracer::SetField(const vtkm::cont::Field& scalarField, const vtkm::Range& scalarRange)
{
  ScalarField = scalarField;
  ScalarRange = scalarRange;
}

//0TTER
void RayTracer::SetNormalField(const vtkm::cont::Field& normalField){
  std::cout << "\t eyyy actually got the normal field into the RayTracer :D \n\n\n\n";
  NormalField = normalField;
}

void RayTracer::SetShadingFunction(void (*func)(void)){
  ShaderFunc = func;
}

void RayTracer::SetColorMap(const vtkm::cont::ArrayHandle<vtkm::Vec4f_32>& colorMap)
{
  ColorMap = colorMap;
}

void RayTracer::SetLights(const std::vector<vtkm::rendering::Light>& lights)
{
  Lights = lights;

  std::cout << Lights[0].GetIntensity()  << ":intensity\n";
  std::cout << Lights[0].GetPosition()   << ":position\n";
  std::cout << Lights[0].GetColor()      << ":color\n";
  std::cout << Lights[0].GetLookDir()    << ":lookdir\n";
  std::cout << Lights[0].IsDirectional() << ":is directional\n";

}


void RayTracer::SetRayShader(vtkm::rendering::Shader s){
  s_obj = s;
  std::cout << "calling shader setting within raytracer\n";
}

void RayTracer::Render(Ray<vtkm::Float32>& rays)
{
  RenderOnDevice(rays);
}

void RayTracer::Render(Ray<vtkm::Float64>& rays)
{
  RenderOnDevice(rays);
}

void RayTracer::SetShadingOn(bool on)
{
  Shade = on;
}

vtkm::Id RayTracer::GetNumberOfShapes() const
{
  return NumberOfShapes;
}

void RayTracer::Clear()
{
  Intersectors.clear();
}

template <typename Precision>
void RayTracer::RenderOnDevice(Ray<Precision>& rays)
{
  using Timer = vtkm::cont::Timer;

  Logger* logger = Logger::GetInstance();
  Timer renderTimer;
  renderTimer.Start();
  vtkm::Float64 time = 0.;
  logger->OpenLogEntry("ray_tracer");
  logger->AddLogData("device", GetDeviceString());

  logger->AddLogData("shapes", NumberOfShapes);
  logger->AddLogData("num_rays", rays.NumRays);

  size_t numShapes = Intersectors.size();
  if (NumberOfShapes > 0)
  {
    Timer timer;
    timer.Start();

    for (size_t i = 0; i < numShapes; ++i)
    {
      Intersectors[i]->IntersectRays(rays);
      time = timer.GetElapsedTime();
      logger->AddLogData("intersect", time);

      timer.Start();
      //0TTER
      //Goes to TriangleIntersector.cxx
      //this is where we need to put in the normal field data
      Intersectors[i]->SetFieldNormals(NormalField);
      Intersectors[i]->IntersectionData(rays, ScalarField, ScalarRange);

      time = timer.GetElapsedTime();
      logger->AddLogData("intersection_data", time);
      timer.Start();

      // Calculate the color at the intersection  point
      detail::SurfaceColor surfaceColor;

      surfaceColor.run(rays, ColorMap, camera, this->Shade, Lights, ShaderFunc, s_obj);

      time = timer.GetElapsedTime();
      logger->AddLogData("shade", time);
    }
  }

  time = renderTimer.GetElapsedTime();
  logger->CloseLogEntry(time);
} // RenderOnDevice
}
}
} // namespace vtkm::rendering::raytracing
