#include <cmath>
#include <iostream>
#include "cloudshader.h"
#include "common/noise/cloudnoise.h"


Color CloudShader::shade(const Scene &scene, const Ray &ray) const
{
    Vector3d hitPoint = ray.origin + ray.direction * ray.length; // Potentially add epsilon

    // Collect getNoise through the cloud
    float cloudLength = 0.0f;   // Length of cloud in ray direction

    // Get background color behind cloud and information about the clouds length
    Ray cloudRay = ray;
    cloudRay.origin = ray.origin + (ray.length + REFR_EPS) * ray.direction;
    cloudRay.length = INFINITY;
    cloudRay.primitive = nullptr;

    // Get out of cloud primitive first
    if (ray.primitive->intersect(cloudRay))
    {
        // Get length
        cloudLength = cloudRay.length;

        // Prepare ray for background color
        cloudRay.setRemainingBounces(cloudRay.getRemainingBounces() + 1);
        cloudRay.origin = cloudRay.origin + (cloudRay.length + REFR_EPS) * cloudRay.direction;
        cloudRay.length = INFINITY;
        cloudRay.primitive = nullptr;
    }
    Color background = scene.traceRay(cloudRay);

    if (cloudLength == 0.0f) return background; // No cloud or at edge

    // Calculate step length
    float stepLength =  settings.densitySteps;
    int noiseSamples = std::floor(cloudLength / stepLength);

    // Step through cloud
    float transmittance = 1.0f;
    Color cloudColor = Color(0, 0, 0);
    for (int i = 0; i < noiseSamples; ++i)
    {
        // Get sample point
        Vector3d lengthDirection = (float) i * stepLength * ray.direction;
        Vector3d samplePoint = hitPoint + lengthDirection;

        // Get data at point
        float sampleDensity = getCloudDensity(samplePoint, ray.primitive) * stepLength;

        if (sampleDensity <= 0)
        {
            continue;
        }

        cloudColor += lightMarch(scene, samplePoint, lengthDirection, ray.primitive) * sampleDensity * transmittance;
        transmittance *= calcBeer(sampleDensity * settings.lightAbsorptionThroughCloud);

        if (transmittance < TRANSMITTANCE_BREAK) break;
    }

    return background * transmittance + cloudColor;
}

bool CloudShader::isTransparent() const
{
    return true;
}

CloudShader::CloudShader(const CloudSettings &settings) : settings(settings),
                                                          cloudNoise(CloudNoise(NOISE_SIZE, settings.seed))
{
    cloudNoise.invert = true;
}

float CloudShader::getCloudDensity(Vector3d point, Primitive const *primitive) const
{
    //! Requires the unscaled point
    float edgeDensity = getEdgeDensity(point, primitive);

    point /= settings.scale;
    float density = cloudNoise.getNoise(point);

    // Threshold
    // TODO: Smooth out!
    density = std::max(0.0f, density + settings.densityOffset) * settings.densityIntensity;

    return density * edgeDensity;
}

float CloudShader::getEdgeDensity(const Vector3d &point, const Primitive *primitive) const
{
    if (primitive == nullptr)
    {
        return 1;
    }

    Vector3d size = Vector3d(0, 0, 0);
    size.x = primitive->maximumBounds(0) - primitive->minimumBounds(0);
    size.y = primitive->maximumBounds(1) - primitive->minimumBounds(1);
    size.z = primitive->maximumBounds(2) - primitive->minimumBounds(2);

    Vector3d center = Vector3d(0, 0, 0);
    center.x = primitive->maximumBounds(0) + primitive->minimumBounds(0);
    center.y = primitive->maximumBounds(1) + primitive->minimumBounds(1);
    center.z = primitive->maximumBounds(2) + primitive->minimumBounds(2);
    center /= 2;

    Vector3d distance = point - center;
    distance.x = std::abs(distance.x);
    distance.y = std::abs(distance.y);
    distance.z = std::abs(distance.z);

    Vector3d distanceFromEdge = size / 2 - distance;
    distanceFromEdge.x = std::max(0.0f, distanceFromEdge.x);
    distanceFromEdge.y = std::max(0.0f, distanceFromEdge.y);
    distanceFromEdge.z = std::max(0.0f, distanceFromEdge.z);

    float fallOff = std::min(distanceFromEdge.x, std::min(distanceFromEdge.y, distanceFromEdge.z)) /
                    settings.edgeFadeOffDistance;

    return std::min(1.0f, fallOff);
}

Color CloudShader::lightMarch(const Scene &scene, Vector3d currentInCloudPosition, Vector3d lengthDistance,
                              const Primitive *cloudObject) const
{
    Color cloudColor = Color(0, 0, 0);

    // For alle lights
    for (const auto &light: scene.lights())
    {
        Ray ray = Ray(currentInCloudPosition - lengthDistance, normalized(lengthDistance));
        ray.length = length(lengthDistance);
        ray.primitive = cloudObject;
        auto illumination = light->illuminate(scene, ray);

        // Handle ambient lights
        if (illumination.distance == 0.0f)
        {
            cloudColor += illumination.color;
            continue;
        }

        // Light ray from object to light
        Ray lightRay;
        lightRay.origin = currentInCloudPosition;
        lightRay.direction = -illumination.direction;
        lightRay.primitive = cloudObject;
        lightRay.length = 0;    // Starting in cloud itself

        float density = this->rayDensity(lightRay, illumination.distance);

        // Proper light calculation
        float transmittance = calcBeer(density * settings.lightAbsorptionTowardsLight);

        transmittance *= phase(normalized(lengthDistance), lightRay.direction);

        cloudColor += transmittance * illumination.color;
    }

    float darknessFactor = settings.darknessThreshold +
                           (1.0f - settings.darknessThreshold);
    return cloudColor * darknessFactor;
}

Color CloudShader::transparency(const Scene &scene, const Ray &ray, float maxLength) const
{
    float density = rayDensity(ray, maxLength);

    float transmittance = calcBeer(density * settings.shadowLightAbsorption);
    transmittance = 1 - (1 - transmittance) * settings.shadowIntensity;

    return Color(1, 1, 1) * transmittance;
}

float CloudShader::rayDensity(const Ray &ray, float maxLength) const
{
    Vector3d startPoint = ray.origin + ray.direction * (ray.length + REFR_EPS);

    // Determine length of cloud
    Ray cloudRay;
    cloudRay.origin = startPoint;
    cloudRay.length = INFINITY;
    cloudRay.direction = ray.direction;
    cloudRay.primitive = nullptr;

    // Get out of cloud primitive first
    if (ray.primitive != nullptr && !ray.primitive->intersect(cloudRay) || cloudRay.length == INFINITY ||
        cloudRay.length <= 0)
    {
        // Assume ray started in cloud
        cloudRay.length = ray.length;
        cloudRay.primitive = ray.primitive;
    }
    float cloudLength = std::min(cloudRay.length, maxLength - ray.length);

    // Calculate step length
    float stepLength = settings.densitySteps;
    int noiseSamples = std::floor(cloudLength / stepLength);

    // Step through cloud
    float density = 0.0f;
    for (int i = 0; i < noiseSamples; ++i)
    {
        // Get sample point
        Vector3d samplePoint = startPoint + (float) i * stepLength * ray.direction;

        // Get data at point
        density += getCloudDensity(samplePoint, ray.primitive) * stepLength;
    }

    // If there is length left, check if it is in the cloud recursively
    if (cloudLength < maxLength - ray.length)
    {
        Vector3d endPoint = startPoint + (cloudLength + REFR_EPS) * ray.direction;
        Ray recursiveRay = cloudRay;
        recursiveRay.origin = endPoint;
        recursiveRay.length = 0;
        recursiveRay.primitive = nullptr;

        if (ray.primitive != nullptr && ray.primitive->intersect(recursiveRay) && recursiveRay.length > 0)
        {
            density += rayDensity(recursiveRay, maxLength - (cloudLength + REFR_EPS));
        }
    }

    return density;
}

float CloudShader::phase(Vector3d visualRay, Vector3d illuminationRay) const
{
    // The angle between the visual and illumination rays
    float cosTheta = dotProduct(visualRay, illuminationRay);

    // The Dual-Lob Henyey-Greenstein function
    float blend = .5;
    float phaseBlend = calcDualHenyeyGreenstein(cosTheta, settings.phaseA) * (1 - blend) +
            calcDualHenyeyGreenstein(cosTheta, -settings.phaseB) * blend;

    // Clamp the result to the range [0, 1]
    phaseBlend = std::max(std::min(phaseBlend, 1.0f), 0.0f);

    return settings.phaseOffset + phaseBlend * settings.phaseIntensity;
}

float CloudShader::calcBeer(float d)
{
    float beer = std::exp(-d);
    return beer;
}


float CloudShader::calcDualHenyeyGreenstein(float cosTheta, float g)
{
    float g2 = g * g;
    return (1 - g2) / (4 * 3.1415f * std::pow(1 + g2 - 2 * g * (cosTheta), 1.5f));
}