//////////////////////////////////////////////////////////////////////////
//                                                                      //
//  This file contains the implementation of the CShaderGenerator       //
//  class.                                                              //
//                                                                      //
//////////////////////////////////////////////////////////////////////////
#include <iomanip>
#include <iostream>
#include "CShaderProgram.hpp"
#include "CPlane.hpp"
#include "CQuadric.hpp"
#include "CSphere.hpp"
#include "CTriangle.hpp"
#include "CLight.hpp"
#include "CScene.hpp"
#include "CShaderGenerator.hpp"
using std::cerr;
using std::endl;
using std::ios;
using std::ostream;
using std::string;




CShaderGenerator *CShaderGenerator::s_pInstance = NULL;




//////////////////////////////////////////////////////////////////////////
//                                                                      //
//  Class    : CShaderGenerator                                         //
//                                                                      //
//  Function : getInstance() (static)                                   //
//                                                                      //
//  Comments : This function returns the unique instance of the         //
//             singleton.                                               //
//                                                                      //
//////////////////////////////////////////////////////////////////////////
CShaderGenerator* CShaderGenerator::getInstance()
{
  if (s_pInstance == NULL)
    s_pInstance = new CShaderGenerator();

  return s_pInstance;
}




//////////////////////////////////////////////////////////////////////////
//                                                                      //
//  Class    : CShaderGenerator                                         //
//                                                                      //
//  Function : freeInstance() (static)                                  //
//                                                                      //
//  Comments : This function frees the unique instance of the           //
//             singleton.                                               //
//                                                                      //
//////////////////////////////////////////////////////////////////////////
void CShaderGenerator::freeInstance()
{
  delete s_pInstance;
  s_pInstance = NULL;
}




//////////////////////////////////////////////////////////////////////////
//                                                                      //
//  Class    : CShaderGenerator                                         //
//                                                                      //
//  Function : CShaderGenerator() (default constructor)                 //
//                                                                      //
//  Comments :                                                          //
//                                                                      //
//////////////////////////////////////////////////////////////////////////
CShaderGenerator::CShaderGenerator()
{
}




//////////////////////////////////////////////////////////////////////////
//                                                                      //
//  Class    : CShaderGenerator                                         //
//                                                                      //
//  Function : ~CShaderGenerator() (destructor)                         //
//                                                                      //
//  Comments :                                                          //
//                                                                      //
//////////////////////////////////////////////////////////////////////////
CShaderGenerator::~CShaderGenerator()
{
}




//////////////////////////////////////////////////////////////////////////
//                                                                      //
//  Class    : CShaderGenerator                                         //
//                                                                      //
//  Function : buildShader()                                            //
//                                                                      //
//  Comments : This function compiles and link a shader program.        //
//                                                                      //
//////////////////////////////////////////////////////////////////////////
CShaderProgram *CShaderGenerator::buildShader(string shaderString)
{
  CShaderProgram *pShader = CShaderProgram::getInstance();
  if (pShader == NULL) {
    cerr << "No shader support." << endl;
    return NULL;
  }

  if (!pShader->addShader(CShaderProgram::FRAGMENT_SHADER, 
                          shaderString)) {
    cerr << "Could not compile the shader." << endl;
    delete pShader;
    return NULL;
  }
  if (!pShader->link()) {
    cerr << "Could not link the shader." << endl;
    delete pShader;
    return NULL;
  }

  return pShader;
}




//////////////////////////////////////////////////////////////////////////
//                                                                      //
//  Class    : CShaderGenerator                                         //
//                                                                      //
//  Function : genreateShader()                                         //
//                                                                      //
//  Comments : This function generates a shader program that renders a  //
//             scene using raytracing.                                  //
//                                                                      //
//////////////////////////////////////////////////////////////////////////
void CShaderGenerator::generateShader(ostream& o,
                                      const CScene *pScene) const
{
  // All floats need to have a dot or a warning will be generated will
  // compiling the shader
  o.setf(ios::showpoint);

  // Generate the shader
  generateUniforms(o);
  generateStructures(o);

  // Surface specific code
  CPlane::generateShaderForClass(o);
  CQuadric::generateShaderForClass(o);
  CSphere::generateShaderForClass(o);
  CTriangle::generateShaderForClass(o);

  // Add a separator
  o <<
    "\n\n"
    "//////////////////////////////////////////////////////////////////////\n"
    "\n\n";

  // Scene specific code
  generateSceneSpecific(o, pScene);

  // Add a separator
  o <<
    "\n\n"
    "//////////////////////////////////////////////////////////////////////\n"
    "\n\n";

  // The rest of the shader code
  generateGenericCode(o);
}




//////////////////////////////////////////////////////////////////////////
//                                                                      //
//  Class    : CShaderGenerator                                         //
//                                                                      //
//  Function : generateUniforms()                                       //
//                                                                      //
//  Comments : This function generates the definition of the uniform    //
//             variables that can be set to customize the behavior of   //
//             the shader.                                              //
//                                                                      //
//////////////////////////////////////////////////////////////////////////
void CShaderGenerator::generateUniforms(ostream& o) const
{
  o <<
    "// Customable ways to see the scene\n"
    "uniform float u_aspectRatio;\n"
    "uniform float u_angle;\n"
    "uniform vec3  u_target;\n"
    "uniform vec3  u_up;\n"
    "uniform vec3  u_position;\n"
    "\n"
    "\n";
}




//////////////////////////////////////////////////////////////////////////
//                                                                      //
//  Class    : CShaderGenerator                                         //
//                                                                      //
//  Function : generateStructures()                                     //
//                                                                      //
//  Comments : This function generates the definition of the basic      //
//             structures that are needed by the raytracer.             //
//                                                                      //
//////////////////////////////////////////////////////////////////////////
void CShaderGenerator::generateStructures(ostream& o) const
{
  o <<
    "// Raytracing structures\n"
    "struct Ray\n"
    "{\n"
    "    vec3  origin;\n"
    "    vec3  direction;\n"
    "    float energy;\n"
    "    float refractionIndex;\n"
    "};\n"
    "\n"
    "struct Light {\n"
    "    vec3  color;\n"
    "    float intensity;\n"
    "    vec3  position;\n"
    "};\n"
    "\n"
    "struct Surface\n"
    "{\n"
    "    vec3  color;\n"
    "    float ambientCoef;\n"
    "    float diffuseCoef;\n"
    "    float specularCoef;\n"
    "    float shininessCoef;\n"
    "    float reflectionCoef;\n"
    "    float refractionCoef;\n"
    "    float refractionIndex;\n"
    "};\n"
    "\n"
    "struct Intersection {\n"
    "    Surface surface;\n"
    "    float distance;\n"
    "    vec3  normal;\n"
    "};\n"
    "\n"
    "\n";
}




//////////////////////////////////////////////////////////////////////////
//                                                                      //
//  Class    : CShaderGenerator                                         //
//                                                                      //
//  Function : generateStructures()                                     //
//                                                                      //
//  Comments : This function generates the definition of the basic      //
//             structures that are needed by the raytracer.             //
//                                                                      //
//////////////////////////////////////////////////////////////////////////
void CShaderGenerator::generateSceneSpecific(ostream& o,
                                             const CScene *pScene) const
{
  CScene::surface_container::const_iterator surfaceIt;
  CScene::light_container::const_iterator   lightIt;
  int i;

  // Scene specific constants
  o << "const vec3  g_backgroundColor = vec3("
    << pScene->m_PropBackColor.red() << " , "
    << pScene->m_PropBackColor.green() << " , "
    << pScene->m_PropBackColor.blue() << ");\n"

    << "const float g_refractionIndex = "
    << pScene->m_PropGenRefrIndex << ";\n"

    << "const float RAY_OFFSET = "
    << pScene->m_PropRayOffset << ";\n"

    << "const float MIN_ENERGY = "
    << pScene->m_PropMinEnergy << ";\n"
    << "\n\n";



  // Scene objects
  for (surfaceIt = pScene->m_vSurfaces.begin();
       surfaceIt != pScene->m_vSurfaces.end();
       ++surfaceIt)
    {
      (*surfaceIt)->GenerateShaderForObject(o);
    }

  // Scene lights
  i = 0;
  for (lightIt = pScene->m_vLights.begin();
       lightIt != pScene->m_vLights.end();
       ++lightIt)
    {
      o << "const Light light" << i++ << " = Light(\n    vec3("
        << (*lightIt)->m_PropLightColor.red() << " , "
        << (*lightIt)->m_PropLightColor.green() << " , "
        << (*lightIt)->m_PropLightColor.blue() << "),\n"
        << "    " << (*lightIt)->m_PropIntensity << ",\n"
        << "    vec3(" << (*lightIt)->m_PropPosition.x << " , "
        << (*lightIt)->m_PropPosition.y << " , "
        << (*lightIt)->m_PropPosition.z << ")\n);\n";
    }



  // Intersections with object
  generateGetIntersections(o, pScene);

  // Unrolling the getSurfaceColor recursive functions
  generateGetSurfaceColor(o, pScene);
}




//////////////////////////////////////////////////////////////////////////
//                                                                      //
//  Class    : CShaderGenerator                                         //
//                                                                      //
//  Function : generateGetIntersections()                               //
//                                                                      //
//  Comments : This function generates the functions that checks for    //
//             intersections with the objects in the scene.             //
//                                                                      //
//////////////////////////////////////////////////////////////////////////
void CShaderGenerator::generateGetIntersections(ostream& o,
                                                const CScene *pScene) const
{
  CScene::surface_container::const_iterator surfaceIt;
  CScene::light_container::const_iterator   lightIt;
  int i;

  o << "\n\n"
    "/////////////////////////////////////////////////////////////////////\n"
    "//                                                                 //\n"
    "//  Intersection checking functions (generated to fit this scene)  //\n"
    "//                                                                 //\n"
    "/////////////////////////////////////////////////////////////////////\n"
    "Intersection getIntersection(Ray ray)\n"
    "{\n"
    "    Intersection t;\n"
    "    t.distance = -1.0;\n"
    "\n";

  // Check with all scene surfaces
  i = 0;
  for (surfaceIt = pScene->m_vSurfaces.begin();
       surfaceIt != pScene->m_vSurfaces.end();
       ++surfaceIt, ++i)
    {
      o <<
      "    Intersection t" << i << " = intersect(" << (*surfaceIt)->getName() << ", ray);\n"
      "    if ( (t" << i << ".distance > 0.0) &&\n"
      "         ((t.distance < 0.0) || (t" << i << ".distance < t.distance)) ) {\n"
      "        t = t" << i << ";\n"
      "    }\n";
    }

  o <<
    "\n"
    "    return t;\n"
    "}\n\n";




  // Now the function for what surfaces are between a hit point and a
  // light
  o <<
    "vec3 getLightFilter(Ray lightRay)\n"
    "{\n";

  switch (pScene->m_PropShadowType)
    {
      // No shadows
    case (0):
      o << "    return vec3(1.0);\n";
      break;

      // Simple shadows: if an object is between the light and the hit
      // point, the light does not contribute to the lighting
    case (1):
      o <<
        "    Intersection tLight = getIntersection(lightRay);\n"
        "    if (tLight.distance > 0.0)\n"
        "        return vec3(0.0);\n"
        "    else\n"
        "        return vec3(1.0);\n";
      break;

      // First shadows: only the first object encountered by the ray
      // going from the ray to the light can filter the light
    case (2):
      o <<
        "    Intersection tLight = getIntersection(lightRay);\n"
        "    if (tLight.distance > 0.0)\n"
        "        return tLight.surface.refractionCoef * tLight.surface.color;\n"
        "    else\n"
        "        return vec3(1.0);\n";
      break;

      // "Complex" shadows: all the objets betweeen the light and the
      // hit point filter the light.  Using this imply generating
      // another function that holds the list of surfaces
    case (3):
      o <<
        "    vec3 lightFilter = vec3(1.0);\n"
        "    Intersection tLight;\n"
        "\n"
        "    // The \"intersect\" functions expect a unit direction vector for\n"
        "    // the ray\n"
        "    float distance = length(lightRay.direction);\n"
        "    lightRay.direction /= distance;\n"
        "\n";

      i = 0;
      for (surfaceIt = pScene->m_vSurfaces.begin();
           surfaceIt != pScene->m_vSurfaces.end();
           ++surfaceIt, ++i)
        {
          o <<
            "    tLight = intersect(" << (*surfaceIt)->getName() << ", lightRay);\n"
            "    if ( (tLight.distance > 0.0) && (tLight.distance < distance) ) {\n"
            "        lightFilter *= tLight.surface.color * tLight.surface.refractionCoef;\n"
            "    }\n";
        }

      o << "\n    return lightFilter;\n";
      break;
    };

  o << "}\n\n";




  // The contribution of a single light
  o <<
    "vec3 getLightContribution(Light light, Ray ray, Intersection t, vec3 hitPoint)\n"
    "{\n"
    "    vec3 color = vec3(0.0);\n"
    "\n"
    "    vec3 lightDirection = light.position - hitPoint;\n"
    "\n"
    "    // We offset the position of the light ray here to avoid artefacts\n"
    "    // on the surfaces.  The drivers seems to react strange if we do\n"
    "    // some tests that would otherwise be cleaner, like testing for\n"
    "    // the surface ID to avoid self-shadowing.\n"
    "    Ray lightRay = Ray(hitPoint+RAY_OFFSET*lightDirection,\n"
    "                       lightDirection,\n"
    "                       1.0,\n"
    "                       1.0);\n"
    "\n"
    "    if (dot(lightRay.direction, t.normal) > 0.0) {\n"
    "        // Get the color from the light\n"
    "        vec3 lightColor = light.color * getLightFilter(lightRay);\n"
    "	lightRay.direction = normalize(lightRay.direction);\n"
    "	lightDirection = normalize(lightDirection);\n"
    "\n"
    "        // Gouraud shading\n"
    "        float gouraudFactor = light.intensity * t.surface.diffuseCoef *\n"
    "            dot(t.normal, lightRay.direction);\n"
    "\n"
    "        color = t.surface.color * gouraudFactor * lightColor;\n"
    "\n"
    "        // Phong shading\n"
    "        // Reflected ray: since lightRay is oriented from the hit point\n"
    "        // towards the light, we get:\n"
    "        // Reflected = 2*(projection on normal) - LightRay\n"
    "        vec3 reflectedRay = reflect(lightRay.direction, t.normal);\n"
    "\n"
    "        float cosBetweenRay = dot(reflectedRay, ray.direction) /\n"
    "                              dot(ray.direction, ray.direction);\n"
    "\n"
    "        float phongFactor = 0.0;\n"
    "        if (cosBetweenRay >= 0.0) {\n"
    "            phongFactor = light.intensity * t.surface.specularCoef *\n"
    "                pow( cosBetweenRay, t.surface.shininessCoef );\n"
    "        }\n"
    "\n"
    "        color += phongFactor * lightColor;\n"
    "    }\n"
    "\n"
    "    return color;\n"
    "}\n"
    "\n";




  // The complete light contribution generation
  o <<
    "vec3 getLightsContribution(Ray ray, Intersection t, vec3 hitPoint)\n"
    "{\n"
    "    vec3 color = vec3(0.0);\n"
    "\n";
  // Scene lights
  i = 0;
  for (lightIt = pScene->m_vLights.begin();
       lightIt != pScene->m_vLights.end();
       ++lightIt)
    {
      o << "    color += getLightContribution(light"
        << i++ << ", ray, t, hitPoint);\n";
    }

  o <<
    "\n"
    "    return color;\n"
    "}\n"
    "\n"
    "\n";
}




//////////////////////////////////////////////////////////////////////////
//                                                                      //
//  Class    : CShaderGenerator                                         //
//                                                                      //
//  Function : generateGetSurfaceColor()                                //
//                                                                      //
//  Comments : This function generates the getSurfaceColor() function   //
//             of the shader.  It is needed because this function, in   //
//             a normal raytracer, can call itself recursively.  Since  //
//             GLSL specification does not allow recursive function     //
//             calls, we create an instance of the function for each    //
//             level of    recursion, which the last level not making   //
//             recursive calls.                                         //
//                                                                      //
//////////////////////////////////////////////////////////////////////////
void CShaderGenerator::generateGetSurfaceColor(ostream& o,
                                               const CScene *pScene) const
{
  o <<
    "/////////////////////////////////////////////////////////////////////\n"
    "//                                                                 //\n"
    "//  Unrolling of the getColor recursive function so that it has a  //\n"
    "//  maximum recursion depth.                                       //\n"
    "//                                                                 //\n"
    "/////////////////////////////////////////////////////////////////////\n";

  int nbLevel = 0;
  if ( (pScene->m_PropReflectionEnabled || pScene->m_PropRefractionEnabled) &&
       (pScene->m_PropMaxLevel > 0) ) {
    nbLevel = pScene->m_PropMaxLevel;
  }

  for (int level = nbLevel; level >= 0; --level) {
    ////////////////////////////////////
    //                                //
    //  getSurfaceColor() generation  //
    //                                //
    ////////////////////////////////////

    // Function prototype (and name)
    o << "vec3 getSurfaceColor";
    if (level != 0)
      o << level - 1;

    //Rest of the prototype, and light contribution
    o <<
      "(Ray ray, Intersection t)\n"
      "{\n"
      "    vec3 color = t.surface.color * t.surface.ambientCoef;\n"
      "\n"
      "    vec3 hitPoint = ray.origin + t.distance * ray.direction;\n"
      "    color += getLightsContribution(ray, t, hitPoint);\n"
      "\n";

    // On the last level, no "recursion"
    if (level != nbLevel) {
      if (pScene->m_PropReflectionEnabled) {
        o <<
          "    // Reflection\n"
          "    float reflectedRayEnergy = t.surface.reflectionCoef * ray.energy;\n"
          "    if (reflectedRayEnergy > MIN_ENERGY) {\n"
          "        // Since the ray is going towards the scene surface, the\n"
          "        // reflected ray is:\n"
          "        vec3 reflectedRayDirection = reflect(ray.direction, t.normal);\n"
          "\n"
          "        Ray reflectedRay = Ray(hitPoint + RAY_OFFSET*reflectedRayDirection,\n"
          "                               reflectedRayDirection,\n"
          "                               reflectedRayEnergy,\n"
          "                               ray.refractionIndex);\n"
          "\n"
          "        color += t.surface.reflectionCoef * getColor"
          << level <<
          "(reflectedRay);\n"
          "    }\n\n";
      }

      if (pScene->m_PropRefractionEnabled) {
        o <<
          "    // Refraction\n"
          "    float refractedRayEnergy = t.surface.refractionCoef * ray.energy;\n"
          "    if (refractedRayEnergy > MIN_ENERGY) {\n"
          "        float eta; // Ratio of indices of refraction\n"
          "        float refractedRayRefractionIndex;\n"
          "\n"
          "        if (ray.refractionIndex == t.surface.refractionIndex) {\n"
          "            // Inside going out\n"
          "            refractedRayRefractionIndex = g_refractionIndex;\n"
          "            eta = t.surface.refractionIndex / g_refractionIndex;\n"
          "        }\n"
          "        else {\n"
          "            // Outside going in\n"
          "            refractedRayRefractionIndex = t.surface.refractionIndex;\n"
          "            eta = g_refractionIndex / t.surface.refractionIndex;\n"
          "        }\n"
          "\n"
          "        vec3 refracted = refract(normalize(ray.direction), t.normal, eta);\n"
          "        if (all(equal(refracted, vec3(0.0)))) {\n"
          "            // Total internal refraction\n"
          "            refracted = reflect(ray.direction, t.normal);\n"
          "            refractedRayRefractionIndex = ray.refractionIndex;\n"
          "        }\n"
          "\n"
          "        Ray refractedRay = Ray(hitPoint + RAY_OFFSET*refracted,\n"
          "                               refracted,\n"
          "                               refractedRayEnergy,\n"
          "                               refractedRayRefractionIndex);\n"
          "\n"
          "        color += t.surface.refractionCoef * getColor"
          << level <<
          "(refractedRay);\n"
          "    }\n\n";
      }
    }

    // Return the computed color
    o <<
      "    return color;\n"
      "}\n";




    /////////////////////////////
    //                         //
    //  getColor() generation  //
    //                         //
    /////////////////////////////

    // Function prototype (and name)
    o << "vec3 getColor";
    if (level != 0)
      o << level - 1;

    // Rest of prototype and intersection
    o <<
      "(Ray ray)\n"
      "{\n"
      "    Intersection t = getIntersection(ray);\n"
      "\n"
      "    if (t.distance < 0.0)\n"
      "        return g_backgroundColor;\n"
      "    else {\n"
      "        t.normal = faceforward(t.normal, ray.direction, t.normal);\n"
      "        return getSurfaceColor";
    if (level != 0)
      o << level - 1;
    o <<
      "(ray, t);\n"
      "    }\n"
      "}\n\n";
  }
}




//////////////////////////////////////////////////////////////////////////
//                                                                      //
//  Class    : CShaderGenerator                                         //
//                                                                      //
//  Function : generateGenericCode()                                    //
//                                                                      //
//  Comments : This function generates the generic raytracing code      //
//             that calls the already generated code.                   //
//                                                                      //
//////////////////////////////////////////////////////////////////////////
void CShaderGenerator::generateGenericCode(ostream& o) const
{
  o <<
    "// Rest of the shader code\n"
    "mat4 getModelView(void)\n"
    "{\n"
    "    vec3 n = normalize(u_position - u_target);\n"
    "    vec3 v = normalize(u_up - n*dot(u_up, n)/dot(n,n));\n"
    "    vec3 u = cross(v,n);\n"
    "\n"
    "    mat4 rot = mat4(vec4(u, 0.0),\n"
    "                    vec4(v, 0.0),\n"
    "                    vec4(n, 0.0),\n"
    "                    vec4(0,0,0,1));\n"
    "\n"
    "    return rot;\n"
    "}\n"
    "\n"
    "vec3 getDirection()\n"
    "{\n"
    "    // Compute the ray at a 1 distance from the camera\n"
    "    float halfHeight = tan(radians(u_angle/2.0));\n"
    "    float halfWidth  = u_aspectRatio*halfHeight;\n"
    "\n"
    "    vec3 direction = vec3((2.0*gl_TexCoord[0].x - 1.0)*halfWidth,\n"
    "                          (2.0*gl_TexCoord[0].y - 1.0)*halfHeight,\n"
    "                          -1.0);\n"
    "\n"
    "    return normalize((getModelView()*vec4(direction,1)).xyz);\n"
    "}\n"
    "\n"
    "void main()\n"
    "{\n"
    "    Ray ray = Ray(u_position,\n"
    "                  getDirection(),\n"
    "                  1.0,\n"
    "                  g_refractionIndex);\n"
    "\n"
    "    gl_FragColor = vec4(getColor(ray), 1.0);\n"
    "}\n";
}
