/***********************************************/
/*           Copyright (c) 2025 Belmu          */
/*             All Rights Reserved             */
/***********************************************/

/*
    [Credits]:
        Jessie - providing CIE CMF curve, spectral to XYZ function, CIE binary search function and CMF weights method (https://github.com/Jessie-LC)
*/

#include "/include/fragment/bsdf/heitz.glsl"
#include "/include/fragment/bsdf/mbnm.glsl"

uniform sampler2D depthtex0;

vec3 evaluateHeitzBSDF(float throughput, vec3 origin, vec3 direction, float wavelength, vec2 fragCoords, vec3 cmfWeights) {
    /* Atmospherics */
    vec3 sunPosition  = calculateSunPosition (LATITUDE, LONGITUDE);
    vec3 moonPosition = calculateMoonPosition(LATITUDE, LONGITUDE);

    moonAngularRadius = CELESTIAL_SIZE_MULTIPLIER * moonRadius / (moonPosition.z * 1e3);

    vec3 sunDirection  = normalize(sphericalToCartesian(sunPosition ));
    vec3 moonDirection = normalize(sphericalToCartesian(moonPosition));

    vec3 sunConeDirection  = generateConeVector(sunDirection , rand2F(), sunAngularRadius );
    vec3 moonConeDirection = generateConeVector(moonDirection, rand2F(), moonAngularRadius);

    float rayleighCoefficient = rayleighCrossSection(wavelength, airIOR);
    float mieCoefficient      = mieCrossSection(wavelength);
    float ozoneCoefficient    = ozoneCrossSection(wavelength);

    vec3 attenuationCoefficients = vec3(rayleighCoefficient, mieCoefficient / mieScatteringAlbedo, ozoneCoefficient);

    float asymmetry = aerosolAsymmetry(wavelength);

    float sunIrradianceSpectrum = solarIrradianceAM0(wavelength);
    float sunRadianceSpectrum   = sunIrradianceSpectrum / coneAngleToSolidAngle(asin(sunAngularRadius));

    float moonRadianceSpectrum   = physicalMoon(moonConeDirection, sunConeDirection, moonDirection, sunIrradianceSpectrum, wavelength);
    float moonIrradianceSpectrum = moonRadianceSpectrum * coneAngleToSolidAngle(asin(moonAngularRadius));

    float irradiance;
    vec3  celestialBodyDirection;

    bool celestialPathProbability = 0.5 > randF();

    if (celestialPathProbability) {
        irradiance             = sunIrradianceSpectrum / 0.5;
        celestialBodyDirection = sunConeDirection;
    } else {
        irradiance             = moonIrradianceSpectrum / 0.5;
        celestialBodyDirection = moonConeDirection;
    }

    float estimate = 0.0;

    if (texture(depthtex0, fragCoords).r == 1.0) {
        estimate += throughput * physicalSun (direction, sunDirection, wavelength) * sunRadianceSpectrum;
        estimate += throughput * physicalMoon(direction, sunDirection, moonDirection, sunIrradianceSpectrum, wavelength);
    }

    /* Ray Data */
    vec3 rayPosition  = origin;
    vec3 rayDirection = direction;
    
    VoxelIntersection hit;
    Material material;

    int raySteps = MAX_RAYTRACE_STEPS;

    for (int scatteringOrder = 0; scatteringOrder < MAX_PT_BOUNCES; scatteringOrder++) {
        
        #if RUSSIAN_ROULETTE == 1
            if (throughput < randF()) {
                throughput = 0.0;
                break;
            }
            throughput /= saturate(throughput);
        #endif

        /* Voxel Tracing */
        hit = raytraceVoxel(rayPosition, rayDirection, raySteps);

        /* Atmospheric Contribution */
        if (!hit.intersect) {
            #if defined WORLD_OVERWORLD
                vec3 atmosphereRayPosition = rayPosition + vec3(0.0, planetRadius + cameraPosition.y, 0.0);
                
                float atmosphere = estimateAtmosphericScattering(
                    atmosphereRayPosition, rayDirection, sunConeDirection, moonConeDirection, attenuationCoefficients, sunRadianceSpectrum, sunIrradianceSpectrum, moonIrradianceSpectrum, asymmetry, wavelength
                );
                
                estimate += throughput * atmosphere;
            #endif
            break;
        }

        /* Voxel Data */
        material    = getVoxelMaterial(hit.voxel.packedData, hit.textureCoords, hit.position, hit.normal, wavelength);
        rayPosition = hit.position + hit.normal * 1e-5;

        int blockId = unpackId(hit.voxel.packedData);

        bool canRefract = isTranslucent(blockId) || blockId == 974;
        bool sphere     = isSphere(blockId);

        mat3 tbnMatrix  = sphere ? getGeometryTBNMatrix(hit.normal) : getShadingTBNMatrix(hit.normal);
        vec3 wi         = rayDirection * tbnMatrix;
             wi.z       = sign(dot(hit.normal, rayDirection)) * abs(wi.z);
        bool wi_outside = wi.z < 0.0;

        /* Shadow Ray */
        #if defined WORLD_OVERWORLD
            bool shadowRayHit = raytraceVoxel(rayPosition, celestialBodyDirection, MAX_PT_SHADOW_STEPS).intersect;

            if (!shadowRayHit) {
                #if RAIN_PUDDLES == 1
                    rainPuddles(material, getWaterMaterial(0.0, vec3(0.0), vec3(0.0), vec3(0.0)).F0, rayPosition, hit.normal);
                #endif

                vec3 wo = celestialBodyDirection * tbnMatrix;

                float phase = 0.0;
                if (canRefract) {
                    bool wo_outside = wo.z > 0.0;

                    if (sphere) {
                        phase = evaluateMicrosurfaceTranslucent(wi, wo, wi_outside, wo_outside, material);
                    } else {
                        phase = evaluateMicrosurfaceTranslucentMBNM(wi, wo, wi_outside, wo_outside, material);
                    }
                } else {
                    if (sphere) {
                        phase = evaluateMicrosurfaceOpaque(wi, wo, material, wavelength);
                    } else {
                        phase = evaluateMicrosurfaceOpaqueMBNM(wi, wo, material, wavelength);
                    }
                }
                
                vec3 atmosphereRayPosition = rayPosition + vec3(0.0, planetRadius + cameraPosition.y, 0.0);

                float transmittance = ratioTrackingEstimator(atmosphereRayPosition, celestialBodyDirection, attenuationCoefficients);

                estimate += throughput * phase * irradiance * transmittance;
            }
        #endif

        /* Phase Sampling */
        if (canRefract) {
            if (sphere) {
                sampleMicrosurfaceTranslucent(throughput, wi, wi_outside, material);
            } else {
                sampleMicrosurfaceTranslucentMBNM(throughput, wi, wi_outside, material);
            }
            rayDirection = tbnMatrix * wi;
        } else {
            if (sphere) {
                sampleMicrosurfaceOpaque(estimate, throughput, wi, material, wavelength);
            } else {
                sampleMicrosurfaceOpaqueMBNM(estimate, throughput, wi, material, wavelength);
            }
            rayDirection = tbnMatrix * wi;
            
            if (dot(hit.normal, rayDirection) < 0.0) break;
        }

        if (isnan(throughput)) throughput = 0.0;

        raySteps = MAX_PT_BOUNCE_STEPS;
    }

    if (isnan(estimate)) estimate = 0.0;

    return max0(spectrumToLinear(estimate, wavelength, cmfWeights));
}
