#ifndef __PARTICIPATING_MEDIA_HLSL__
#define __PARTICIPATING_MEDIA_HLSL__

/*
    Interpretation of "Physically Based Rendering: From Theory to Implementation"

    Emission: radiance that is added to the environment from luminous particles

    Absorption: reduction in radiance due to the conversion of light to another form of energy, such as heat

    Scattering: radiance heading in one direction that is scattered to other directions due to collisions with particles
          ~In-scattering: increases radiance due to rays scattered into the path of the current ray
          ~Out-scattering: reduces the radiance exiting a region of the beam because some of it is deflected to different directions
          ~Real-scattering: in-scattering and out-scattering
          ~Null-scattering: scattering that does not correspond to any type of physical scattering process, specified to have no effect on the distribution of light
                            + allows us to treat inhomogeneous media as if they were homogeneous by handling null scattering interactions

    Extinction: the combination of absorption and scattering coefficients
                + may also be called "attenuation"
                + describes the rate of interactions (absorption/scattering) per unit distance
                + higher extinction (like thick fog) increases the rate of interactions 
                ++ SampleExponential will return lower values, making interactions more likely

    Null-scattering coefficient: probability of a null-scattering event per unit distance traveled in the medium
                                 + defined by the constant "majorant", such that absorption + scattering + null scattering = majorant
                                 + the majorant represents the maximum extinction along the ray

    Single-Scattering Albedo: 0 to 1, probability of scattering vs absorption at a scattering event
                              + high albedo means most interactions are scattering, more light bounce
                              + low albedo means most interactions are absorption, darkening the result 
                              + defined as the scattering coefficient divided by extinction

    Mean Free Path: average distance a ray travels in a medium (with attenuation coefficient) before interacting with a particle

    Homogeneous: properties are constant throughout some region of space 

    Inhomogeneous: properties vary throughout space
                   + for example, sampling 3D density textures, or a low frequency noise texture to modify the properties
                   + also known as "heterogeneous"

    The differential reduction in radiance along a beam is a linear function of its initial radiance
                   + the fraction of light absorbed does not vary based on the ray's radiance, but is always a fixed fraction

    Phase function: describes the angular distribution of scattered radiance at a point
                    + phase functions are probability distributions for scattering in a particular direction
                    + for most naturally occurring media (isotropic), the phase is a 1D function of the angle between incoming/outgoing directions

    Isotropic/Symmetric Media: response to illumination is invariant to rotation

    Anisotropic Media: particles arranged in a coherent structure 
                       + phase function can be a 4D function of the two directions 
                       + true for crystals, or media made of coherently oriented fibers

    Isotropic Phase: scatters light equally in all directions
                     + in Henyey-Greenstein terms: asymmetry = 0

    Anisotropic Phase: scatters light biased in a particular direction
                       + in Henyey-Greenstein terms: asymmetry > 0 or < 0
                       ++ forward scattering if > 0
                       ++ backward scattering if < 0

    The asymmetry parameter (g) in the HG model is the integral of the product of the given phase function and the cosine of the angle
                    + referred to as the "mean cosine" of the scattering angle
                    + the greater the magnitude of "g", the more concentrated the scattering distribution will be in the respective direction (forward or backward)

    Transmittance gives the fraction of radiance that is transmitted between two points
                    + or put another way, the fraction of radiance that passes through unimpeded by extinction
                    + transmittance is multiplicative along points on a ray, making it possible to incrementally compute transmittance at multiple points along a ray

    Delta Tracking: recursive (or in this case iterative) estimation of transmittance that continues until termination due to extinction, Russian Roulette (not shown here) or until the sampled point is past the endpoint of the ray
                    + also known as "Woodcock tracking"
                    + randomly decides whether the ray interacts with a true particle or a null particle at each scattering event
                    ++ interactions with null particles (null scattering) are ignored and the algorithm continues, restarting from the sampled point 
                    ++ interactions with true particles cause extinction
*/

static const uint MIN_DELTA_ITERATIONS = 16;
static const uint MAX_DELTA_ITERATIONS = 1024;

struct Medium
{
    float3x4 worldToVolume;
    float3 absorption;
    float3 scattering;
    float asymmetry;
    float density;
    float emissiveTemperature;
    uint volumeTextureIndex;
    bool isHomogeneous;
};

struct MediumStack
{
    Medium layers; //That's a lot of registers, especially to make this an array (not a good idea, this is just a stub), can be optimized
    uint layerDepth;
};

struct MediumSample
{
    float3 throughput;
    float3 scatterPosition;
    float3 scatterDirection;
    float3 emission; //add to radiance, regardless of scattering/termination
    bool isScattered;
    bool isTerminated;
};

struct MediumEventSample
{
    float3 absorption;
    float3 scattering;
    float asymmetry;
    float density;
    float emissiveTemperature;
};

/* 
    Blackbody: A perfect emitter. Not perfectly realizable, but has useful properties for modeling emissive behaviors anyway
               + Absorbs all incident power and reflects none

    Planck's Law (this function): The radiance emitted by a blackbody as a function of wavelength and temperature (measured in kelvins)

    Kirchhoff's Law (non-blackbodies): The emitted radiance distribution at any frequency is equal to the emission of a perfect blackbody at that frequency 
                                       times the fraction of incident radiance at that frequency that is absorbed by the object

    Stefan-Boltzmann Law (note): Doubling the temperature of a blackbody emitter increases the total energy emitted by a factor of 16
                                 + Stefan-Boltzmann Constant x temperature^4
*/
float CalculateBlackbodySpectral(float wavelength, float temperatureInKelvin) 
{
    float c = 299792458.f;     //meters/second, speed of light in the medium (this value is a vacuum)
    float h = 6.62606957e-34f; //joule-seconds (energy * time), Planck's constant
    float kb = 1.3806488e-23f; //joules/kelvin, Boltzmann Constant (note: different from Stefan-Boltzmann Constant!)

    float l = wavelength * 1e-9f;
    float Le = (2 * h * c * c) / (pow(l, 5) * (exp((h * c) / (l * kb * temperatureInKelvin)) - 1));

    return Le;
}

/*
    The above is handy to have as reference, but this shader works with linearized sRGB, so just using an approximation for now

    Specifically using "Color Temperature Conversion of Homestar.io" by Neil Bartlett
    https://www.zombieprototypes.com/p-210/

    Which is an improvement of "How to Convert Temperature (K) to RGB" by Tanner Helland
    https://tannerhelland.com/2012/09/18/convert-temperature-rgb-algorithm-code.html
*/

float3 CalculateBlackbodyLinearizedSRGB(float temperatureInKelvin)
{
    float scaledTemperature = temperatureInKelvin / 100.0;
    float3 blackbodyAsLinearizedSRGB = 0;

    if (scaledTemperature < 66.0)
    {
        blackbodyAsLinearizedSRGB.r = 255.0;
    }
    else
    {
        float x = scaledTemperature - 55.0;
        blackbodyAsLinearizedSRGB.r = 351.97690566805693 + 0.114206453784165 * x - 40.25366309332127 * log(x);
    }

    if (scaledTemperature < 10.0)
    {
        blackbodyAsLinearizedSRGB.g = 0.0;
    }
    else if (scaledTemperature < 66.0)
    {
        float x = scaledTemperature - 2.0;
        blackbodyAsLinearizedSRGB.g = -155.25485562709179 - 0.44596950469579133 * x + 104.49216199393888 * log(x);
    }
    else
    {
        float x = scaledTemperature - 50.0;
        blackbodyAsLinearizedSRGB.g = 325.4494125711974 + 0.07943456536662342 * x - 28.0852963507957 * log(x);
    }

    if (scaledTemperature >= 66.0)
    {
        blackbodyAsLinearizedSRGB.b = 255.0;
    }
    else if (scaledTemperature < 20.0)
    {
        blackbodyAsLinearizedSRGB.b = 0.0;
    }
    else
    {
        float x = scaledTemperature - 10.0;
        blackbodyAsLinearizedSRGB.b = -254.76935184120902 + 0.8274096064007395 * x + 115.67994401066147 * log(x);
    }

    //The approximation is sRGB encoded
    blackbodyAsLinearizedSRGB = DecodeFromSRGB(saturate(blackbodyAsLinearizedSRGB / 255.0));

    //As defined in Helland's guide
    static const float WHITE_POINT_TEMPERATURE = 6500;

    //Stefan-Boltzmann
    float energyScaling = pow(temperatureInKelvin / WHITE_POINT_TEMPERATURE, 4.0);
    blackbodyAsLinearizedSRGB *= energyScaling;

    return blackbodyAsLinearizedSRGB;
}

//Transmittance: The attenuation of light in participating media between points
float3 CalculateTransmittanceHomogeneous(float3 extinction, float distance)
{
    //Beer's Law: light intensity decreases exponentially as it travels through a medium, proportional to extinction and distance
    return exp(-extinction * distance);
}

//It's expensive to run the sampling per-channel, so one channel is picked, but the bias needs to be accounted for
float3 CorrectTransmittanceChannelBias(float3 transmittanceEstimate, float3 channelSelectionProbabilities)
{
    float weightedAverageTransmittance = dot(transmittanceEstimate, channelSelectionProbabilities);
    return weightedAverageTransmittance > 0 ? transmittanceEstimate / weightedAverageTransmittance : 0;
}

//Phase: The distribution of light scattered at a point in space
//Needed for next event estimation and multiple importance sampling
//"Diffuse Radiation in the Galaxy", Henyey & Greenstein
float CalculatePhaseHenyeyGreenstein(float cosineTheta, float g) 
{
    float g2 = g * g;
    float denom = 1.0 + g2 + 2.0 * g * cosineTheta;
    return INV_FOUR_PI * (1.0 - g2) / (denom * SafeSqrt(denom));
}

//Samples in-scattering directions proportional to the phase function. The phase function is normalized (integrates to 1 over the sphere)
float3 SamplePhaseHenyeyGreenstein(float3 outgoingDirection, float g, float2 randomSample) 
{
    float cosTheta;
    if (abs(g) < BIG_EPSILON)
    {
        //Isotropic (or close enough): sample uniformly on sphere
        cosTheta = 1.0 - 2.0 * randomSample.x;
    }
    else
    {
        //Anisotropic: sample according to the directional bias
        float g2 = g * g;
        float sqrTerm = (1.0 - g2) / (1.0 + g - 2.0 * g * randomSample.x);
        cosTheta = -1.0 / (2.0 * g) * (1.0 + g2 - sqrTerm * sqrTerm);
    }

    float sinTheta = SafeSqrt(1.0 - cosTheta * cosTheta);
    float phi = TWO_PI * randomSample.y;

    float3 localDirection = SphericalDirection(sinTheta, cosTheta, phi);

    float3 u;
    float3 v;
    CreateOrthonormalBasis(outgoingDirection, u, v);
    
    return localDirection.x * u + localDirection.y * v + localDirection.z * outgoingDirection;
}

MediumEventSample SampleMediumAtPoint(Medium medium, float3 worldPosition)
{
    //Full medium sample stubbed, but only actually sampling density
    MediumEventSample eventSample;
    eventSample.scattering = medium.scattering;
    eventSample.absorption = medium.absorption;
    eventSample.asymmetry = medium.asymmetry;
    eventSample.density = medium.density;
    eventSample.emissiveTemperature = medium.emissiveTemperature;

    if (medium.volumeTextureIndex == INVALID_RESOURCE_INDEX)
    {
        return eventSample;
    }
 
    float3 uvw = mul(medium.worldToVolume, float4(worldPosition, 1.0));
    if (any(saturate(uvw) != uvw))
    {
        //Just assuming zero outside of bounds for now
        eventSample.density = 0.0;
        return eventSample;
    }

    float sampleDensity = saturate(GetTexture3DNonUniform(medium.volumeTextureIndex).SampleLevel(GetSampler(linearClampSampler), uvw, 0).x);
    if (sampleDensity < BIG_EPSILON)
    {
        sampleDensity = 0.0;
    }
    
    eventSample.density *= sampleDensity;

    return eventSample;
}

MediumSample SampleMediumHomogeneous(Medium medium, float3 rayOrigin, float3 rayDirection, float distanceToSurface, float3 throughput, inout uint randomSeed)
{
    MediumSample result;
    result.throughput = float3(1, 1, 1);
    result.scatterPosition = rayOrigin;
    result.scatterDirection = rayDirection;
    result.emission = float3(0, 0, 0);
    result.isScattered = false;
    result.isTerminated = false;

    float3 extinction = (medium.absorption + medium.scattering) * medium.density;

    //There are 3 different channels available to sample distance using the exponential
    //Only one can be used for that purpose though, so need some way to decide that. PBRT picks the first wavelength through the whole path
    //For now just picking channel weighted by throughput * extinction, better options must exist 
    float3 channelImportance = throughput * extinction;
    float totalChannelImportance = channelImportance.r + channelImportance.g + channelImportance.b;
    if (totalChannelImportance < EPSILON)
    {
        return result;
    }

    float channelRandom = GeneratePCGRandom1D(randomSeed) * totalChannelImportance;
    uint selectedChannel = 2;
    if (channelRandom < channelImportance.r)
    {
        selectedChannel = 0;
    }
    else if (channelRandom < channelImportance.r + channelImportance.g)
    {
        selectedChannel = 1;
    }

    float3 emissionFactor = 0;
    if (medium.emissiveTemperature > 0)
    {
        float3 absorptionRatio = select(extinction > 0, medium.absorption * medium.density / extinction, 0);
        float3 emission = CalculateBlackbodyLinearizedSRGB(medium.emissiveTemperature);

        emissionFactor = absorptionRatio * emission;
    }

    float sampleDistance = SampleExponential(GeneratePCGRandom1D(randomSeed), extinction[selectedChannel]);
    if (sampleDistance >= distanceToSurface)
    {
        //Handle surface hit
        float3 transmittance = CalculateTransmittanceHomogeneous(extinction, distanceToSurface);
        float3 selectionProbabilities = channelImportance / totalChannelImportance;
        
        result.throughput = CorrectTransmittanceChannelBias(transmittance, selectionProbabilities);
        result.scatterPosition = rayOrigin + rayDirection * distanceToSurface;
        result.emission = emissionFactor * (1.0 - transmittance);
        return result;
    }
    
    float singleScatteringAlbedo = medium.scattering[selectedChannel] / (medium.absorption[selectedChannel] + medium.scattering[selectedChannel]);
    if (GeneratePCGRandom1D(randomSeed) < singleScatteringAlbedo)
    {
        //Scattering
        float3 scatteringCoefficient = medium.scattering * medium.density;
        float3 transmittance = CalculateTransmittanceHomogeneous(extinction, sampleDistance);
        float3 scatteredTransmittance = scatteringCoefficient * transmittance;
        float3 selectionProbabilities = channelImportance / totalChannelImportance;

        result.throughput = CorrectTransmittanceChannelBias(scatteredTransmittance, selectionProbabilities);
        result.scatterPosition = rayOrigin + rayDirection * sampleDistance;
        result.scatterDirection = SamplePhaseHenyeyGreenstein(-rayDirection, medium.asymmetry, GeneratePCGRandom2D(randomSeed));
        result.emission = emissionFactor * (1.0 - transmittance);
        result.isScattered = true;
    }
    else
    {
        //Absorption
        result.emission = emissionFactor * (1.0 - CalculateTransmittanceHomogeneous(extinction, sampleDistance));
        result.isTerminated = true;
    }
    
    return result;
}

MediumSample SampleMediumDeltaTracking(Medium medium, float3 rayOrigin, float3 rayDirection, float distanceToSurface, float3 throughput, inout uint randomSeed)
{
    MediumSample result;
    result.throughput = float3(1, 1, 1);
    result.scatterPosition = rayOrigin;
    result.scatterDirection = rayDirection;
    result.emission = float3(0, 0, 0);
    result.isScattered = false;
    result.isTerminated = false;

    float3 perChannelMajorant = (medium.absorption + medium.scattering) * medium.density;
    float majorant = max(perChannelMajorant.r, max(perChannelMajorant.g, perChannelMajorant.b));

    //There are 3 different channels available to sample distance using the exponential
    //Only one can be used for that purpose though, so need some way to decide that. PBRT picks the first wavelength through the whole path
    //For now just picking channel weighted by throughput * extinction, better options must exist
    float3 channelImportance = throughput * perChannelMajorant;
    float totalChannelImportance = channelImportance.r + channelImportance.g + channelImportance.b;
    if (totalChannelImportance < EPSILON)
    {
        return result;
    }

    float channelRandom = GeneratePCGRandom1D(randomSeed) * totalChannelImportance;
    uint selectedChannel = 2;
    if (channelRandom < channelImportance.r)
    {
        selectedChannel = 0;
    }
    else if (channelRandom < (channelImportance.r + channelImportance.g))
    {
        selectedChannel = 1;
    }
    
    float3 transmittanceEstimate = float3(1, 1, 1);
    float relativePositionOnRay = 0.0;

    //Estimate iterations required based on distance & majorant, needs to be increased if there is visible/unexpected darkening
    uint maxIterations = clamp((uint)(distanceToSurface * majorant) * 3, MIN_DELTA_ITERATIONS, MAX_DELTA_ITERATIONS);

    for (uint iteration = 0; iteration < maxIterations; iteration++)
    {
        relativePositionOnRay += SampleExponential(GeneratePCGRandom1D(randomSeed), majorant);
        if (relativePositionOnRay >= distanceToSurface)
        {
            //Handle surface hit
            float3 selectionProbabilities = channelImportance / totalChannelImportance;
            
            result.throughput = CorrectTransmittanceChannelBias(transmittanceEstimate, selectionProbabilities);
            if(all(result.throughput <= 0))
            {
                result.isTerminated = true;
                return result;
            }
            
            result.scatterPosition = rayOrigin + rayDirection * distanceToSurface;
            return result;
        }
        
        float3 samplePosition = rayOrigin + rayDirection * relativePositionOnRay;
        MediumEventSample eventSample = SampleMediumAtPoint(medium, samplePosition);

        float3 absorptionCoefficient = eventSample.absorption * eventSample.density;
        float3 scatteringCoefficient = eventSample.scattering * eventSample.density;
        float3 extinctionCoefficient = (absorptionCoefficient + scatteringCoefficient);
        float3 nullScatteringCoefficient = max(0, majorant - extinctionCoefficient);

        if (eventSample.emissiveTemperature > 0)
        {
            float3 emissionCoefficient = absorptionCoefficient * CalculateBlackbodyLinearizedSRGB(eventSample.emissiveTemperature);
            result.emission += (transmittanceEstimate * emissionCoefficient / majorant);
        }

        float absorptionProbability = absorptionCoefficient[selectedChannel] / majorant;
        float scatterProbability = scatteringCoefficient[selectedChannel] / majorant;
        
        float interaction = GeneratePCGRandom1D(randomSeed);
        if (interaction < absorptionProbability)
        {
            //Absorption
            result.isTerminated = true;
            return result;
        }
        else if (interaction < absorptionProbability + scatterProbability)
        {
            //Scattering
            transmittanceEstimate *= scatteringCoefficient / majorant;
            float3 selectionProbabilities = channelImportance / totalChannelImportance;
            
            result.throughput = CorrectTransmittanceChannelBias(transmittanceEstimate, selectionProbabilities);
            if(all(result.throughput <= 0))
            {
                result.isTerminated = true;
                return result;
            }
            
            result.scatterPosition = samplePosition;
            result.scatterDirection = SamplePhaseHenyeyGreenstein(-rayDirection, eventSample.asymmetry, GeneratePCGRandom2D(randomSeed));
            result.isScattered = true;
            return result;
        }
        else
        {
            //Null scattering
            //This is where PBRT recommends using things like coarse grids to be able to identify accurate local majorants and/or local density,
            //so that large swaths of space can be skipped through, else you can get stuck here looping through empty space, wasting performance 
            transmittanceEstimate *= nullScatteringCoefficient / majorant;
            continue;
        }
    }
    
    result.throughput = float3(1, 1, 1);
    result.scatterPosition = rayOrigin;
    result.scatterDirection = rayDirection;
    result.isScattered = false;
    result.isTerminated = false;
    return result;
}

MediumSample SampleMedium(Medium medium, float3 rayOrigin, float3 rayDirection, float distanceToSurface, float3 throughput, inout uint randomSeed)
{
    if (medium.isHomogeneous)
    {
        return SampleMediumHomogeneous(medium, rayOrigin, rayDirection, distanceToSurface, throughput, randomSeed);
    }
    else
    {
        return SampleMediumDeltaTracking(medium, rayOrigin, rayDirection, distanceToSurface, throughput, randomSeed);
    }
}
#endif