#include <math.h>               /* exp(), pow(), M_PI */
#include <stdlib.h>

#include <string>
#include <vector>
#include <utility>              /* pair */
#include <map>
#include <algorithm>            /* sort() */
#include <functional>           /* binary_function */
#include <typeinfo>             /* typeid operator support */

#include "Scene.h"
#include "BMP.h"
#include "util/Color.h"
#include "shapes/Shape.h"
#include "Light.h"

using namespace std;

#define MAX_AMBIENT_OCCLUSION_DISTANCE 50.0

Scene::Scene(const map<string, const char *> & options,
             const char * filename)
{
    m_width = 800;
    m_height = 600;
    m_multisample_level = 1;
    m_vfov = 60.0;
    m_ambient_light = Color(0.2, 0.2, 0.2);
    m_max_depth = 10;
    m_exposure = 1.0f;
    m_ambient_occlusion_level = 0;
    m_transforms.push(Transform());

    load(filename);

    /* after loading the scene file, apply any command-line render options */
    for (map<const string, const char *>::const_iterator it = options.begin();
         it != options.end();
         it++)
    {
        if (it->first == "width")
        {
            m_width = atoi(it->second);
        }
        else if (it->first == "height")
        {
            m_height = atoi(it->second);
        }
        else if (it->first == "multisample")
        {
            m_multisample_level = atoi(it->second);
        }
        else if (it->first == "field-of-view")
        {
            m_vfov = atof(it->second);
        }
        else if (it->first == "max-depth")
        {
            m_max_depth = atoi(it->second);
        }
        else if (it->first == "ambient-occlusion")
        {
            m_ambient_occlusion_level = atoi(it->second);
        }
    }

    /* view plane distance is calculated based on the field of view */
    m_view_plane_dist = (m_height / 2.0) / tan(M_PI * m_vfov / 360.0);
    m_sample_span = 1.0 / m_multisample_level;
    m_half_sample_span = m_sample_span / 2.0;
    m_multisample_level_squared = m_multisample_level * m_multisample_level;
}

Scene::~Scene()
{
}

void Scene::renderPixel(int x, int y, unsigned char * pixel)
{
    /* calculate the ray going from the camera through this pixel */
    Color finalColor;
    for (int i = 0; i < m_multisample_level; i++)
    {
        for (int j = 0; j < m_multisample_level; j++)
        {
            double rx = (x + i * m_sample_span + m_half_sample_span)
                      - (m_width / 2.0);
            double rz = (m_height / 2.0)
                      - (y + j * m_sample_span + m_half_sample_span);

            Ray ray(Vector(0, 0, 0), Vector(rx, m_view_plane_dist, rz));

            finalColor += traceRay(ray);
        }
    }

    /* apply exposure formula so we aren't saturated */
    finalColor.r = 1.0f -
        exp(-m_exposure * finalColor.r / m_multisample_level_squared);
    finalColor.g = 1.0f -
        exp(-m_exposure * finalColor.g / m_multisample_level_squared);
    finalColor.b = 1.0f -
        exp(-m_exposure * finalColor.b / m_multisample_level_squared);

#if 0
    /* gamma correct */
    finalColor.r = pow(finalColor.r, 1 / 2.2);
    finalColor.g = pow(finalColor.g, 1 / 2.2);
    finalColor.b = pow(finalColor.b, 1 / 2.2);
#endif

    /* take the average of all the samples as the final pixel value */
    pixel[BMP_RED]   = (unsigned char) (0xFF * finalColor.r);
    pixel[BMP_GREEN] = (unsigned char) (0xFF * finalColor.g);
    pixel[BMP_BLUE]  = (unsigned char) (0xFF * finalColor.b);
}

Color Scene::traceRay(const Ray & ray)
{
    return traceRayRecurse(ray, m_max_depth, 1.0, NULL);
}

/**
 * factor: the proportion of the final color that this computation is worth
 */
Color Scene::traceRayRecurse(const Ray & ray, int depth, double factor,
        refptr<Material> last_material)
{
    static refptr<Material> air = new Material();
    Color color(0, 0, 0);

    if (last_material.isNull())
        last_material = air;

    Shape::Intersection hit = getRayClosestHit(ray);

    if ( ! hit.shape.isNull() )
    {
        /* compute the Phong lighting for each hit */
        refptr<Material> material = hit.shape->getMaterial();

        /* check for backfaces */
        bool frontface = ray.getDirection() % hit.normal < 0.0;

        if (frontface)
        {
            color = computePhong(material,
                                 ray,
                                 hit.position,
                                 hit.normal);

            if (depth > 0 && factor > SCENE_FACTOR_THRESHOLD)
            {
                double reflectance = material->getReflectance();
                if (factor * reflectance > SCENE_FACTOR_THRESHOLD)
                {
                    color *= (1.0 - reflectance);
                    Vector reflected_direction =
                        ray.getDirection().reflect(hit.normal);
                    Ray newRay(hit.position, reflected_direction);
                    Vector jitter_surface_point = newRay[0.0001];
                    Ray jitterNewRay(jitter_surface_point, reflected_direction);
                    Color c = traceRayRecurse(jitterNewRay,
                                              depth - 1,
                                              factor * reflectance,
                                              material);
                    color += c * reflectance;
                }

                double transparency = material->getTransparency();
                if (factor * transparency > SCENE_FACTOR_THRESHOLD)
                {
                    color *= (1.0 - transparency);
                    Vector jitter_surface_point = hit.position
                                                  + ray.getDirection() * 0.0001;
                    Vector refracted_direction =
                        ray.getDirection().refract(hit.normal,
                                last_material->getRefraction(),
                                material->getRefraction());
                    Ray newRay(jitter_surface_point, refracted_direction);
                    Color c = traceRayRecurse(newRay,
                                              depth - 1,
                                              factor * transparency,
                                              material);
                    color += c * transparency;
                }
            }
        }
        else
        {
            material = air;
            Vector jitter_surface_point = hit.position
                                          + ray.getDirection() * 0.0001;
            Vector refracted_direction =
                ray.getDirection().refract(-hit.normal,
                        last_material->getRefraction(),
                        material->getRefraction());
            Ray newRay(jitter_surface_point, refracted_direction);
            color = traceRayRecurse(newRay, depth, factor, material);
        }
    }

    return color;
}

Shape::Intersection Scene::getRayClosestHit(const Ray & ray)
{
    Shape::Intersection hit;
    double min_dist = 0.0;
    bool foundOne = false;

    /* loop through all shapes in the scene */
    for (vector< refptr<Shape> >::iterator it = m_shapes.begin();
         it != m_shapes.end();
         it++)
    {
        Shape::IntersectionList intersections = (*it)->intersect(*it, ray);

        for (int i = 0, num_results = intersections.size();
             i < num_results;
             i++)
        {
            refptr<Shape> shape = intersections[i].shape;
            const Vector & isect_point = intersections[i].position;
            double intersect_dist = ray.getOrigin().dist_to(isect_point);
            if (foundOne == false || intersect_dist < min_dist)
            {
                hit = intersections[i];
                min_dist = intersect_dist;
                foundOne = true;
            }
        }
    }

    return hit;
}

Color Scene::computePhong(const refptr<Material> material,
                          const Ray & viewRay,
                          const Vector & surfacePoint,
                          const Vector & surfaceNormal)
{
    Color result = m_ambient_light * material->getAmbientColor();
    if (m_ambient_occlusion_level > 0)
    {
        result *= calculateAmbientOcclusion(
                Ray(surfacePoint, surfaceNormal).shift(1e-7));
    }

    Vector viewDirection = -viewRay.getDirection();
    double shininess = material->getShininess();
    const Color & diffuseColor = material->getDiffuseColor();
    const Color & specularColor = material->getSpecularColor();

    for (std::vector< refptr<Light> >::const_iterator it = m_lights.begin();
         it != m_lights.end();
         it++)
    {
        Vector lightC = (*it)->getPosition();
        double lightRadius = (*it)->getRadius();
        Vector directionToLightC = lightC - surfacePoint;
        Vector lightPlaneX = directionToLightC.getPerpendicular().normalize();
        Vector lightPlaneY = (directionToLightC * lightPlaneX).normalize();
        int jitter_samples = 0, jitter_level = (*it)->getJitter();;
        Color jitterResult;
        for (int jitter_index = 0; jitter_index < jitter_level; jitter_index++)
        {
            double jitterRadius = jitter_index * lightRadius
                / (jitter_level - 0.5);
            for (int i = 0, num = (int) (M_PI * jitter_index) + 1; i < num; i++)
            {
                jitter_samples++;
                double jitter_angle = i * 2.0 * M_PI / num;
                Vector jitterPosition = lightC
                    + lightPlaneX * jitterRadius * cos(jitter_angle)
                    + lightPlaneY * jitterRadius * sin(jitter_angle);
                Vector directionToLight = jitterPosition - surfacePoint;
                directionToLight.normalize();
                Vector reflectedLightDirection =
                    (-directionToLight).reflect(surfaceNormal);

                Ray surfaceToLight(surfacePoint, directionToLight);
                Color light_contribution =
                    calculateLightContribution(surfaceToLight.shift(0.0001),
                            jitterPosition);

                if (   light_contribution.r > 0.0
                    || light_contribution.g > 0.0
                    || light_contribution.b > 0.0 )
                {
                    /* calculate the diffuse term */
                    double diffuse_coef = directionToLight % surfaceNormal;
                    if (diffuse_coef > 0.0)
                    {
                        jitterResult += diffuseColor
                            * (*it)->getDiffuseColor()
                            * diffuse_coef
                            * light_contribution;
                    }

                    /* calculate the specular term */
                    double specular_coef = reflectedLightDirection % viewDirection;
                    if (specular_coef > 0.0)
                    {
                        jitterResult += specularColor
                            * (*it)->getSpecularColor()
                            * pow(specular_coef, shininess)
                            * light_contribution;
                    }
                }
            }
        }
        jitterResult /= jitter_samples;
        result += jitterResult;
    }

    return result;
}

Color Scene::calculateLightContribution(const Ray & toLight,
                                         const Vector & lightPosition)
{
    Color contrib(1.0, 1.0, 1.0);
    double dist_to_light = (lightPosition - toLight.getOrigin()).mag();
    double dist_so_far = 0.0;

    Ray currentRay = toLight;

    for (;;)
    {
        Shape::Intersection hit = getRayClosestHit(currentRay);

        if ( hit.shape.isNull() )
            break;

        double offset = currentRay.getOrigin().dist_to(hit.position) + 0.0001;

        if ( dist_so_far + offset > dist_to_light )
            break;

        contrib *= hit.shape->getMaterial()->getTransparency();
        contrib *= hit.shape->getMaterial()->getDiffuseColor();

        if (    contrib.r < SCENE_FACTOR_THRESHOLD
             && contrib.g < SCENE_FACTOR_THRESHOLD
             && contrib.b < SCENE_FACTOR_THRESHOLD )
            break;

        dist_so_far += offset;
        currentRay = currentRay.shift(offset);
    }

    return contrib;
}

Color Scene::calculateAmbientOcclusion(const Ray & surfaceNormal)
{
    Color result(1, 1, 1);
    const int nISteps = m_ambient_occlusion_level * 6;
    const int nJSteps = nISteps / 2;
    int nRays = 0;
    Vector perpX = surfaceNormal.getDirection().getPerpendicular().normalize();
    Vector perpY = (surfaceNormal.getDirection() * perpX).normalize();
    double istep = 2.0 * M_PI / nISteps;
    double jstep = M_PI_2 / nJSteps;
    for (int i = 0; i < nISteps; i++)
    {
        int lim = i > 0 ? nJSteps : 1;
        for (int j = 0; j < lim; j++)
        {
            Vector direction = cos(i * istep) * sin(j * jstep) * perpX
                + sin(i * istep) * sin(j * jstep) * perpY
                + cos(j * jstep) * surfaceNormal.getDirection();
            Ray thisRay(surfaceNormal.getOrigin(), direction);
            double dist = 0.0;
            Color contrib(1, 1, 1);
            while (contrib.r > 0.2 && contrib.g > 0.2 && contrib.b > 0.2)
            {
                Shape::Intersection hit = getRayClosestHit(thisRay);
                if (hit.shape.isNull())
                    break;
                double hitDist = (hit.position - thisRay.getOrigin()).mag();
                dist += hitDist;
                if (dist > MAX_AMBIENT_OCCLUSION_DISTANCE)
                    break;
                contrib *= hit.shape->getMaterial()->getTransparency()
                    * hit.shape->getMaterial()->getDiffuseColor();
                thisRay = thisRay.shift(hitDist + 1E-7);
            }
            result += contrib;
            nRays++;
        }
    }
    result /= nRays;
    return result;
}