#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
    int noiseSamples = settings.densitySamples;
    float stepLength = cloudLength / (float) noiseSamples;

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

        // Get data at point
        float sampleDensity = getCloudDensity(samplePoint) * stepLength;

        if (sampleDensity > 0)
        {
            cloudColor += lightMarch(scene, samplePoint, lengthDirection, ray.primitive) * sampleDensity;
        }

        accumulatedDensity += sampleDensity;
    }

    if (accumulatedDensity > 1)
        cloudColor /= accumulatedDensity;

    float transmittance = exp(-accumulatedDensity * settings.lightAbsorptionThroughCloud);

    return background * transmittance + cloudColor;
}

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

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

float CloudShader::getCloudDensity(Vector3d point) const
{
    point /= settings.scale;

    float density = cloudNoise.getNoise(point);

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

    return density;
}

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
        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);
        density *= settings.lightAbsorptionTowardsLight;

        // Proper light calculation
        float transmittance = getDensityTransmittance(density);
        float scatter = scatterFactor(normalized(lengthDistance), illumination.direction);

        float factor = transmittance;
        if (density > 0)
        {
            factor = settings.darknessThreshold +
                     (1.0f - settings.darknessThreshold) * factor * scatter;
        }

        cloudColor += factor * illumination.color;
    }

    return cloudColor;
}

float CloudShader::getDensityTransmittance(float density) const
{
    return exp(-density) * (1 - exp(-density * 2)) / 0.4f;
}

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

    float transmittance = exp(-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 + 0.0001f);

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

    // Get out of cloud primitive first
    if (ray.primitive != nullptr && !ray.primitive->intersect(cloudRay) || cloudRay.length == INFINITY ||
        cloudRay.length <= 0)
    {
        // Something went wrong => No density
        return 0;
    }
    float cloudLength = std::min(cloudRay.length, maxLength - ray.length);

    // Calculate step length
    int noiseSamples = settings.lightSamples;
    float stepLength = cloudLength / (float) noiseSamples;

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

        // Get data at point
        density += getCloudDensity(samplePoint) * 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::scatterFactor(Vector3d visualRay, Vector3d illuminationRay) const
{
    // The asymmetry parameter
    float g = settings.scatterWeight;

    // The angle between the visual and illumination rays
    float cosTheta = dotProduct(visualRay, illuminationRay);

    // The Dual-Lob Henyey-Greenstein function
    float blend = .5;
    float scatter = HenyeyGreenstein(cosTheta, g) * (1 - blend) + HenyeyGreenstein(cosTheta, -g) * blend;

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

    return scatter;
}

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