/************************************************************
Title    : A Photo-realistic Renderer Based on Photon Mapping
Students : Davood  Rasti
           Alireza Rasti

Version  : 0.01
Date     : 23-06-2012
Updated  : 01-10-2026
*************************************************************/

#include "./3rdparty/stb_collection/stb_image_write.h"
#include "./prrpm.h"

#include <charconv>
#include <chrono>
#include <cmath>
#include <filesystem>
#include <fstream>
#include <iterator>
#include <sstream>
#include <stdexcept>
#include <string_view>

using namespace PRRPM;

// Parse a complete numeric token so malformed scene values cannot silently become zero.
template <typename Value> Value parseNumber(std::string_view token)
{
    Value value{};
    const char* tokenEnd = token.data() + token.size();
    const auto result = std::from_chars(token.data(), tokenEnd, value);
    if (result.ec != std::errc{} || result.ptr != tokenEnd)
        throw std::invalid_argument("invalid numeric value '" + std::string(token) + "'");
    return value;
}

std::vector<std::string> tokenizer(const std::string& line)
{
    std::vector<std::string> stringList;
    std::string normalizedLine = line;
    // Scene files accept whitespace as well as '=', ',' and ';' as separators.
    for (char& character : normalizedLine)
    {
        if (character == '=' || character == ',' || character == ';')
            character = ' ';
    }

    std::istringstream lineStream(normalizedLine);
    std::string token;
    while (lineStream >> token)
        stringList.push_back(token);

    return stringList;
}

bool loadSceneFile(const std::string& fileName)
{
    std::ifstream modelFile(fileName);
    if (!modelFile)
    {
        std::cerr << "Unable to open scene file: " << fileName << '\n';
        return false;
    }

    std::string code;
    PRRPM::Triangle triangle;
    PRRPM::Sphere sphere;

    std::string lineToRead;
    std::size_t lineNumber = 0;
    while (std::getline(modelFile, lineToRead))
    {
        ++lineNumber;
        try
        {
            std::vector<std::string> stringList;
            stringList = tokenizer(lineToRead);

            if (stringList.size() < 1)
                continue;

            code = stringList[0];

            if (code[0] == '*')
                break;
            else if (code[0] == '/')
                continue;
            // Each non-comment line begins with a directive such as imagesize, light, t, or s.
            else if (code == "imagesize" && stringList.size() >= 3)
            {
                imageWidth = parseNumber<ui32>(stringList[1]);
                imageHeight = parseNumber<ui32>(stringList[2]);
            }
            else if (code == "light" && stringList.size() >= 2)
            {
                if (stringList[1] == "on" || stringList[1] == "off")
                    isLightOn = stringList[1] == "on";
                // Accept either the full numeric light record or the compact "light pos" form.
                else if (stringList.size() >= 11)
                {
                    PRRPM::LightSource lightSource{};
                    lightSource.type = parseNumber<ui8>(stringList[1]);
                    lightSource.position =
                        Vector3d(parseNumber<f32>(stringList[2]), parseNumber<f32>(stringList[3]),
                                 parseNumber<f32>(stringList[4]));
                    lightSource.intensity =
                        Vector3d(parseNumber<f32>(stringList[5]), parseNumber<f32>(stringList[6]),
                                 parseNumber<f32>(stringList[7]));
                    lightSource.size =
                        Vector3d(parseNumber<f32>(stringList[8]), parseNumber<f32>(stringList[9]),
                                 parseNumber<f32>(stringList[10]));
                    lightSourcesList.push_back(lightSource);
                }
                else if (stringList[1] == "pos" && stringList.size() >= 5)
                {
                    PRRPM::LightSource lightSource{};
                    lightSource.type = 0;
                    lightSource.position =
                        Vector3d(parseNumber<f32>(stringList[2]), parseNumber<f32>(stringList[3]),
                                 parseNumber<f32>(stringList[4]));
                    lightSource.intensity = Vector3d(1.0f, 1.0f, 1.0f);
                    lightSourcesList.push_back(lightSource);
                }
            }
            else if ((code == "numberOfPhotons" || code == "numberOfPhotonPass" ||
                      code == "numberOfBounces" || code == "numberOfPhotonBounces") &&
                     stringList.size() >= 2)
            {
                const ui32 value = parseNumber<ui32>(stringList[1]);
                if (code == "numberOfPhotons")
                    numberOfPhotons = value;
                else if (code == "numberOfPhotonPass")
                    numberOfPhotonPass = value;
                else
                    numberOfBounces = value;
            }
            else if (code == "photonCasting" && stringList.size() >= 2)
                isPhotonMapping = stringList[1] == "on";
            else if (code == "gui" && stringList.size() >= 2)
            {
                if (stringList[1] == "off")
                    isGUI = false;
                else if (stringList[1] == "on")
                    isGUI = true;
            }
            else if (code == "t" && stringList.size() >= 13)
            {
                triangle.point1.x = parseNumber<f32>(stringList[1]);
                triangle.point1.y = parseNumber<f32>(stringList[2]);
                triangle.point1.z = parseNumber<f32>(stringList[3]);

                triangle.point2.x = parseNumber<f32>(stringList[4]);
                triangle.point2.y = parseNumber<f32>(stringList[5]);
                triangle.point2.z = parseNumber<f32>(stringList[6]);

                triangle.point3.x = parseNumber<f32>(stringList[7]);
                triangle.point3.y = parseNumber<f32>(stringList[8]);
                triangle.point3.z = parseNumber<f32>(stringList[9]);

                triangle.color.x = parseNumber<f32>(stringList[10]);
                triangle.color.y = parseNumber<f32>(stringList[11]);
                triangle.color.z = parseNumber<f32>(stringList[12]);

                triangle.materialID =
                    stringList.size() >= 14 ? parseNumber<ui8>(stringList[13]) : 0;

                // The ray-triangle test uses the precomputed face normal.
                PRRPM::Vector3d edge1 = triangle.point2 - triangle.point1;
                PRRPM::Vector3d edge2 = triangle.point3 - triangle.point1;
                triangle.normal = edge1.cross(edge2);
                triangle.normal.normalize();

                triangleList.push_back(triangle);
            }
            else if (code == "s" && stringList.size() >= 8)
            {
                sphere.position.x = parseNumber<f32>(stringList[1]);
                sphere.position.y = parseNumber<f32>(stringList[2]);
                sphere.position.z = parseNumber<f32>(stringList[3]);

                sphere.radius = parseNumber<f32>(stringList[4]);

                sphere.color.x = parseNumber<f32>(stringList[5]);
                sphere.color.y = parseNumber<f32>(stringList[6]);
                sphere.color.z = parseNumber<f32>(stringList[7]);

                sphere.materialID = stringList.size() >= 9 ? parseNumber<ui8>(stringList[8]) : 0;

                sphereList.push_back(sphere);
            }
        }
        catch (const std::invalid_argument& error)
        {
            // Include the source location so scene-file typos are easy to find.
            std::cerr << fileName << ':' << lineNumber << ": " << error.what() << '\n';
            return false;
        }
    }

    return true;
}

// Advance one queued secondary ray. Status values are -2 (finished), 0 (diffuse),
// 1 (mirror), and 2 (dielectric); storing these states makes recursion unnecessary.
void nextRayTracingPass(ui32 index)
{

    if (secondIntersectionPointsList[index].nubmerOfPhotonAndMaterialID == -2)
        return;

    PRRPM::Ray ray;
    ray.direction = secondIntersectionPointsList[index].fluxAndDirection;
    ray.direction.normalize();
    ray.origin = secondIntersectionPointsList[index].positionAndOrigin;
    ray.objectID = -1;
    ray.distance = maxSceneDepth;

    PRRPM::Vector3d color;
    intersectObjects(&ray);
    if (ray.objectID != -1)
    {
        if (getObjectMaterial(ray.objectID, ray.objectType) == 0)
        {
            if (secondIntersectionPointsList[index].radiusSquare != -1)
                color = getObjectColor(ray.objectID, ray.objectType);
            else
                color = secondIntersectionPointsList[index].color *
                        getObjectColor(ray.objectID, ray.objectType);

            int cellID = computeCellID(ray.intersectionPosition);
            if (cellID >= 0 && cellID < gridSize)
            {
                secondIntersectionPointsList[index].color = color;
                secondIntersectionPointsList[index].positionAndOrigin = ray.intersectionPosition;
                secondIntersectionPointsList[index].normal = getSurfaceNormal(ray);
                secondIntersectionPointsList[index].cellID = cellID;
                // secondIntersectionPointsList[index].materialID = 0;

                secondIntersectionPointsList[index].radiusSquare = irridianceRadiusSquare;
                secondIntersectionPointsList[index].nubmerOfPhotonAndMaterialID = 0;
                secondIntersectionPointsList[index].fluxAndDirection = PRRPM::Vector3d();
            }
        }
        else if (getObjectMaterial(ray.objectID, ray.objectType) == 1)
        {
            PRRPM::Vector3d normal = getSurfaceNormal(ray);
            secondIntersectionPointsList[index].positionAndOrigin = ray.intersectionPosition;
            secondIntersectionPointsList[index].fluxAndDirection =
                ray.direction - normal * normal.dot(ray.direction) * 2.0f;
            secondIntersectionPointsList[index].nubmerOfPhotonAndMaterialID = 1;
        }
        else if (getObjectMaterial(ray.objectID, ray.objectType) == 2)
        {
            secondIntersectionPointsList[index].positionAndOrigin = ray.intersectionPosition;
            secondIntersectionPointsList[index].normal = getSurfaceNormal(ray);
            secondIntersectionPointsList[index].fluxAndDirection = ray.direction;

            if (secondIntersectionPointsList[index].radiusSquare == -1)
                secondIntersectionPointsList[index].color =
                    secondIntersectionPointsList[index].color *
                    getObjectColor(ray.objectID, ray.objectType);
            else
                secondIntersectionPointsList[index].color =
                    getObjectColor(ray.objectID, ray.objectType);
            secondIntersectionPointsList[index].nubmerOfPhotonAndMaterialID = 2;
        }
    }
    else
        secondIntersectionPointsList[index].nubmerOfPhotonAndMaterialID = -2;
}

// Replace each dielectric hit with its reflected and refracted paths, then discard finished paths.
void calculateReflectionFractionForAllPoints()
{
    PRRPM::ui32 step = 0;
    for (PRRPM::ui32 iPoint = 0; iPoint < secondIntersectionPointsList.size(); ++iPoint)
    {
        IntersectionPoint intersectionPoint = secondIntersectionPointsList[iPoint];
        if (intersectionPoint.nubmerOfPhotonAndMaterialID == 2)
        {
            PRRPM::Vector3d reflectionRayDirection;
            PRRPM::Vector3d refractionRayDirection;
            f32 cosOutAngle;
            f32 reflectionIntensity;

            ++step;

            calculateReflectionFractionDirection(intersectionPoint.fluxAndDirection,
                                                 intersectionPoint.normal, &reflectionRayDirection,
                                                 &refractionRayDirection, &cosOutAngle,
                                                 &reflectionIntensity);
            if (cosOutAngle < 0)
            {
                secondIntersectionPointsList[iPoint] = intersectionPoint;
                secondIntersectionPointsList[iPoint].fluxAndDirection = reflectionRayDirection;
                secondIntersectionPointsList[iPoint].radiusSquare = -1;
                ++iPoint;
                continue;
            }

            secondIntersectionPointsList[iPoint - 1] = intersectionPoint;
            secondIntersectionPointsList[iPoint - 1].fluxAndDirection = refractionRayDirection;
            secondIntersectionPointsList[iPoint - 1].color =
                secondIntersectionPointsList[iPoint].color * (1.0 - reflectionIntensity);
            secondIntersectionPointsList[iPoint - 1].radiusSquare = -1;

            secondIntersectionPointsList[iPoint] = intersectionPoint;
            secondIntersectionPointsList[iPoint].fluxAndDirection = reflectionRayDirection;
            // secondIntersectionPointsList[iPoint].nubmerOfPhotonAndMaterialID = -2;
            secondIntersectionPointsList[iPoint].color =
                secondIntersectionPointsList[iPoint].color * reflectionIntensity;
            secondIntersectionPointsList[iPoint].radiusSquare = -1;
        }
    }

    std::sort(secondIntersectionPointsList.begin(), secondIntersectionPointsList.end(),
              compareByMetalID);
    PRRPM::ui32 numberOfNonDiffusedPoint = int(std::count_if(
        secondIntersectionPointsList.begin(), secondIntersectionPointsList.end(), isMetalIDOut));
    secondIntersectionPointsList.erase(secondIntersectionPointsList.begin(),
                                       secondIntersectionPointsList.begin() +
                                           numberOfNonDiffusedPoint);
}

// Cast one primary camera ray per pixel. Diffuse hits are photon-map gather points;
// mirror and dielectric hits are queued for the later non-recursive passes.
void firstRayTracingPass()
{
    PRRPM::f32 screenRatio = (PRRPM::f32)imageWidth / (PRRPM::f32)imageHeight;
    PRRPM::f32 inverseWidth = 1.0f / (PRRPM::f32)imageWidth;
    PRRPM::f32 inverseHeight = 1.0f / (PRRPM::f32)imageHeight;
    PRRPM::f32 inverseWidthOverTwo = inverseWidth / 4.0f;
    PRRPM::f32 inverseHeightOverTwo = inverseHeight / 2.0f;

    for (ui32 y = 0; y < imageHeight; ++y)
    {
        Ray ray;
        unsigned int yPixel = y * imageWidth;

        for (ui32 x = 0; x < imageWidth; ++x)
        {

            ray.direction.x = ((x + inverseWidthOverTwo) * inverseWidth - 0.5) * screenRatio;
            ray.direction.y = ((y + inverseHeightOverTwo) * inverseHeight - 0.5);
            ray.direction.z = screenRatio;
            ray.direction.normalize();
            ray.origin = Vector3d(0.0, 0.0, 0.0);
            ray.objectID = -1;
            ray.distance = maxSceneDepth;

            intersectObjects(&ray);
            if (ray.objectID != -1)
            {
                PRRPM::ui32 index = x + (y * imageWidth);
                if (getObjectMaterial(ray.objectID, ray.objectType) == 0)
                {
                    PRRPM::Vector3d color = getObjectColor(ray.objectID, ray.objectType);
                    int cellID = computeCellID(ray.intersectionPosition);
                    if (cellID >= 0 && cellID < gridSize)
                    {
                        intersectionPointsList[index].color = color;
                        intersectionPointsList[index].positionAndOrigin = ray.intersectionPosition;
                        intersectionPointsList[index].normal = getSurfaceNormal(ray);
                        intersectionPointsList[index].pixelIndex = x + yPixel;
                        intersectionPointsList[index].cellID = cellID;
                        // intersectionPointsList[index].materialID = 0;

                        intersectionPointsList[index].radiusSquare = irridianceRadiusSquare;
                        intersectionPointsList[index].nubmerOfPhotonAndMaterialID = 0;
                        intersectionPointsList[index].fluxAndDirection = Vector3d();
                    }
                }
                else if (getObjectMaterial(ray.objectID, ray.objectType) == 1)
                {
                    PRRPM::Vector3d normal = getSurfaceNormal(ray);
                    intersectionPointsList[index].positionAndOrigin = ray.intersectionPosition;
                    intersectionPointsList[index].fluxAndDirection =
                        ray.direction - normal * normal.dot(ray.direction) * 2.0f;
                    intersectionPointsList[index].pixelIndex = x + yPixel;
                    intersectionPointsList[index].nubmerOfPhotonAndMaterialID = 1;
                }
                else if (getObjectMaterial(ray.objectID, ray.objectType) == 2)
                {
                    intersectionPointsList[index].positionAndOrigin = ray.intersectionPosition;
                    intersectionPointsList[index].normal = getSurfaceNormal(ray);
                    intersectionPointsList[index].pixelIndex = x + yPixel;
                    intersectionPointsList[index].fluxAndDirection = ray.direction;
                    intersectionPointsList[index].color =
                        getObjectColor(ray.objectID, ray.objectType);
                    intersectionPointsList[index].nubmerOfPhotonAndMaterialID = 2;
                }
            }
        }
    }
}

void prepareIntersectionPointsList()
{
    // -1 cell/pixel indices mark pixels that did not produce a valid gather point.
    intersectionPointsList.resize((imageWidth * imageHeight));
    PRRPM::IntersectionPoint intersectionPoint;
    intersectionPoint.color = PRRPM::Vector3d();
    intersectionPoint.positionAndOrigin = PRRPM::Vector3d();
    intersectionPoint.normal = PRRPM::Vector3d();
    intersectionPoint.pixelIndex = -1;
    intersectionPoint.cellID = -1;
    intersectionPoint.radiusSquare = irridianceRadiusSquare;
    intersectionPoint.nubmerOfPhotonAndMaterialID = 0;
    intersectionPoint.fluxAndDirection = PRRPM::Vector3d();
    std::fill(intersectionPointsList.begin(),
              intersectionPointsList.begin() + (imageWidth * imageHeight), intersectionPoint);
}

void rayTracingPass()
{
    firstRayTracingPass();

    for (ui32 iPoint = 0; iPoint < intersectionPointsList.size(); ++iPoint)
    {
        if (intersectionPointsList[iPoint].nubmerOfPhotonAndMaterialID != 0)
        {
            PRRPM::IntersectionPoint intersectionPoint = intersectionPointsList[iPoint];
            intersectionPoint.nubmerOfPhotonAndMaterialID = -2;
            secondIntersectionPointsList.push_back(intersectionPoint);
            secondIntersectionPointsList.push_back(intersectionPointsList[iPoint]);
        }
    }

    // Seed the work list with secondary paths, then advance it a fixed number of bounces.
    calculateReflectionFractionForAllPoints();

    for (ui32 iBounce = 1; iBounce < 5; ++iBounce)
    {
        for (ui32 iPoints = 0; iPoints < secondIntersectionPointsList.size(); ++iPoints)
            nextRayTracingPass(iPoints);

        for (ui32 iPoint = 0; iPoint < secondIntersectionPointsList.size(); ++iPoint)
        {
            if (secondIntersectionPointsList[iPoint].nubmerOfPhotonAndMaterialID == 0)
                intersectionPointsList.push_back(secondIntersectionPointsList[iPoint]);
        }

        std::sort(secondIntersectionPointsList.begin(), secondIntersectionPointsList.end(),
                  compareByCellID);
        ui32 numberOfNonDiffusedPoint = int(std::count_if(
            secondIntersectionPointsList.begin(), secondIntersectionPointsList.end(), isCellIDOut));
        secondIntersectionPointsList.erase(
            secondIntersectionPointsList.end() -
                (secondIntersectionPointsList.size() - numberOfNonDiffusedPoint),
            secondIntersectionPointsList.end());

        std::vector<PRRPM::IntersectionPoint> tempsecondIntersectionPointsList;
        for (ui32 iPoint = 0; iPoint < secondIntersectionPointsList.size(); ++iPoint)
        {
            PRRPM::IntersectionPoint intersectionPoint;
            intersectionPoint.nubmerOfPhotonAndMaterialID = -2;
            tempsecondIntersectionPointsList.push_back(intersectionPoint);
            tempsecondIntersectionPointsList.push_back(secondIntersectionPointsList[iPoint]);
        }

        secondIntersectionPointsList.clear();
        secondIntersectionPointsList = tempsecondIntersectionPointsList;
        calculateReflectionFractionForAllPoints();
    }
}

void findPointAndComputeFlux(const Ray& photonRay, const PRRPM::Vector3d& color)
{
    // Gather this photon into nearby camera gather points and shrink their radii as they
    // accumulate.
    PRRPM::Vector3d surfaceNormal = getSurfaceNormal(photonRay);
    for (std::size_t iNeighbor = 0; iNeighbor < std::size(neighborsDirection); ++iNeighbor)
    {
        int cellID = computeCellID(photonRay.intersectionPosition + neighborsDirection[iNeighbor]);
        if (cellID < 0 || cellID >= gridSize)
            continue;

        const int firstCell = hastValueList[cellID];
        if (firstCell >= 0)
        {
            std::size_t pointIndex = static_cast<std::size_t>(firstCell);
            while (true)
            {

                Vector3d distanceVector = intersectionPointsList[pointIndex].positionAndOrigin -
                                          photonRay.intersectionPosition;
                f32 distance = distanceVector.dot(distanceVector);
                IntersectionPoint* intersectionPoint = &intersectionPointsList[pointIndex];
                f32 dotNormals = intersectionPoint->normal.dot(surfaceNormal);
                if (dotNormals > -0.001 && distance <= intersectionPoint->radiusSquare)
                {
                    f32 radiusReductionCoefficient =
                        (intersectionPointsList[pointIndex].nubmerOfPhotonAndMaterialID * alpha +
                         alpha) /
                        (intersectionPoint->nubmerOfPhotonAndMaterialID * alpha + 1.0);
                    intersectionPoint->radiusSquare *= radiusReductionCoefficient;
                    ++intersectionPoint->nubmerOfPhotonAndMaterialID;
                    intersectionPoint->fluxAndDirection =
                        (intersectionPoint->fluxAndDirection +
                         intersectionPoint->color * color * MathUtility::InverstPI) *
                        radiusReductionCoefficient;
                }
                ++pointIndex;
                if (pointIndex >= intersectionPointsList.size() ||
                    intersectionPointsList[pointIndex].cellID != cellID)
                    break;
            }
        }
    }
}

void photonCasting(std::default_random_engine& randomNumberGenerator,
                   const LightSource& lightSource)
{
    std::uniform_real_distribution<f32> uniform(0.0, 1.0);

    PRRPM::Ray photonRay;
    generatePhotonRayFromLighSource(&photonRay, lightSource, randomNumberGenerator);
    photonRay.objectID = -1;
    photonRay.distance = maxSceneDepth;

    ui8 bounces = 0;
    Vector3d color = lightSource.intensity;
    bool isRefraction = false;
    while (true)
    {
        intersectObjects(&photonRay);
        if (photonRay.objectID != -1)
        {
            Vector3d objectColor = getObjectColor(photonRay.objectID, photonRay.objectType);
            ui8 objectMaterialID = getObjectMaterial(photonRay.objectID, photonRay.objectType);
            if (objectMaterialID == 0)
            {
                // Diffuse hits contribute flux; the next direction is sampled from the hemisphere.
                findPointAndComputeFlux(photonRay, color);
                ++bounces;
                f32 reflectionProbability =
                    std::max(objectColor.x, std::max(objectColor.y, objectColor.z));
                if (bounces >= numberOfBounces ||
                    (bounces > 1 && uniform(randomNumberGenerator) <= reflectionProbability))
                    break;

                if (isRefraction)
                {
                    color = lightSource.intensity;
                    isRefraction = false;
                }
                else
                    color = color * getObjectColor(photonRay.objectID, photonRay.objectType) /
                            reflectionProbability; // * 0.7f;

                f32 r1 = MathUtility::TwoPI * uniform(randomNumberGenerator);
                f32 r2 = uniform(randomNumberGenerator);
                f32 r2s = sqrtf(r2);
                PRRPM::Vector3d tempNormal = getSurfaceNormal(photonRay);
                if (tempNormal.dot(photonRay.direction) >= 0)
                    tempNormal = tempNormal * -1.0f;
                PRRPM::Vector3d w = tempNormal;
                PRRPM::Vector3d u = PRRPM::Vector3d(1, 0, 0);

                if (fabs(w.x) > 0.1)
                    u = PRRPM::Vector3d(0, 1, 0);

                u = u.cross(w);
                u.normalize();
                PRRPM::Vector3d v = w.cross(u);
                photonRay.direction =
                    u * cosf(r1) * r2s + v * sinf(r1) * r2s + w * sqrtf(1.0f - r2);
                photonRay.direction.normalize();

                photonRay.origin = photonRay.intersectionPosition;
                photonRay.objectID = -1;
                photonRay.distance = maxSceneDepth;
                continue;
            }
            else if (objectMaterialID == 1)
            {
                // Mirror material: reflect the photon and continue its path.
                PRRPM::Vector3d normal = getSurfaceNormal(photonRay);
                photonRay.direction =
                    photonRay.direction - normal * normal.dot(photonRay.direction) * 2.0;
                photonRay.origin = photonRay.intersectionPosition;
                photonRay.objectID = -1;
                photonRay.distance = maxSceneDepth;
                ++bounces;
                if (bounces >= numberOfBounces)
                    break;

                continue;
            }
            else if (objectMaterialID == 2)
            {
                // Dielectric material: choose reflection or refraction using Fresnel probability.

                PRRPM::Vector3d reflectionRayDirection;
                PRRPM::Vector3d refractionRayDirection;
                f32 cosOutAngle;
                f32 reflectionIntensity;
                calculateReflectionFractionDirection(
                    photonRay.direction, getSurfaceNormal(photonRay), &reflectionRayDirection,
                    &refractionRayDirection, &cosOutAngle, &reflectionIntensity);
                if (cosOutAngle < 0)
                {
                    photonRay.origin = photonRay.intersectionPosition;
                    photonRay.direction = reflectionRayDirection;
                    photonRay.objectID = -1;
                    photonRay.distance = maxSceneDepth;
                    ++bounces;
                    if (bounces >= numberOfBounces)
                        break;

                    continue;
                }

                photonRay.origin = photonRay.intersectionPosition;
                if (uniform(randomNumberGenerator) >= reflectionIntensity)
                {
                    photonRay.direction = refractionRayDirection;
                    color = getObjectColor(photonRay.objectID, photonRay.objectType) *
                            lightSource.intensity;
                }
                else
                {
                    photonRay.direction = reflectionRayDirection;
                    color = getObjectColor(photonRay.objectID, photonRay.objectType);
                }
                photonRay.objectID = -1;
                photonRay.distance = maxSceneDepth;
                isRefraction = true;
                ++bounces;
                if (bounces >= numberOfBounces)
                    break;

                continue;
            }
            else
                break;
        }
        else
            break;
    }
}

void photonCastingPass()
{
    if (lightSourcesList.empty())
    {
        std::cerr << "No light sources defined in the scene\n";
        return;
    }

    // Each pass launches a fresh set of photons; progress is updated in place on the console.
    for (ui32 iPhotonPass = 0; iPhotonPass < numberOfPhotonPass; ++iPhotonPass)
    {
        std::default_random_engine randomNumberGenerator;
        for (ui32 iPhoton = 0; iPhoton < numberOfPhotons; ++iPhoton)
        {
            LightSource lightSource = lightSourcesList[iPhoton % lightSourcesList.size()];
            photonCasting(randomNumberGenerator, lightSource);
        }

        std::cout << "\rPass: " << iPhotonPass + 1 << '/' << numberOfPhotonPass << std::flush;
    }
    std::cout << '\n';
}

// Build the spatial hash table for efficient photon gathering.
void buildHashTable()
{
    // Sort gather points by cell so each hash entry can point to the start of its cell's range.
    std::sort(intersectionPointsList.begin(), intersectionPointsList.end(), compareByCellID);
    PRRPM::ui32 numberOfNonDiffusedPoint = int(
        std::count_if(intersectionPointsList.begin(), intersectionPointsList.end(), isCellIDOut));
    intersectionPointsList.erase(intersectionPointsList.begin(),
                                 intersectionPointsList.begin() + numberOfNonDiffusedPoint);

    hastValueList.resize(gridSize);
    std::fill(hastValueList.begin(), hastValueList.begin() + gridSize, -1);

    for (std::size_t i = 0; i < intersectionPointsList.size(); ++i)
    {
        const int cellID = intersectionPointsList[i].cellID;
        if (hastValueList[cellID] == -1)
            hastValueList[cellID] = i;
    }
}

// Correct the flux at each gather point to compute the final irradiance.
void fluxCorrection()
{
    // Convert accumulated photon flux to irradiance using each gather point's final disk area.
    imageBufferList.resize((imageWidth * imageHeight));
    std::fill(imageBufferList.begin(), imageBufferList.begin() + (imageWidth * imageHeight),
              Vector3d(0, 0, 0));

    for (std::size_t i = 0; i < intersectionPointsList.size(); ++i)
    {
        const int pixelIndex = intersectionPointsList[i].pixelIndex;
        imageBufferList[pixelIndex] =
            imageBufferList[pixelIndex] +
            intersectionPointsList[i].fluxAndDirection *
                (1.0 / (MathUtility::PI * intersectionPointsList[i].radiusSquare * numberOfPhotons *
                        numberOfPhotonPass));
    }
}

void saveImage(void)
{
    const std::size_t expectedPixelCount = static_cast<std::size_t>(imageWidth) * imageHeight;
    if (imageBufferList.size() != expectedPixelCount)
    {
        std::cerr << "Image buffer size does not match image dimensions\n";
        return;
    }

    // Convert linear floating-point RGB into display-encoded 8-bit RGB for the BMP writer.
    std::vector<unsigned char> pixels(imageBufferList.size() * 3);
    const auto toByte = [](f32 channel)
    {
        channel = std::clamp(channel, 0.0f, 1.0f);
        const f32 srgb = channel <= 0.0031308f ? 12.92f * channel
                                               : 1.055f * std::pow(channel, 1.0f / 2.4f) - 0.055f;
        return static_cast<unsigned char>(std::lround(srgb * 255.0f));
    };

    for (std::size_t y = 0; y < imageHeight; ++y)
    {
        // Camera rows increase upward; image files store the top row first.
        const std::size_t sourceRow = imageHeight - 1 - y;
        for (std::size_t x = 0; x < imageWidth; ++x)
        {
            const std::size_t sourceIndex = sourceRow * imageWidth + x;
            const std::size_t targetIndex = (y * imageWidth + x) * 3;
            pixels[targetIndex] = toByte(imageBufferList[sourceIndex].x);
            pixels[targetIndex + 1] = toByte(imageBufferList[sourceIndex].y);
            pixels[targetIndex + 2] = toByte(imageBufferList[sourceIndex].z);
        }
    }

    if (!stbi_write_bmp(outFileName.c_str(), static_cast<int>(imageWidth),
                        static_cast<int>(imageHeight), 3, pixels.data()))
        std::cerr << "Failed to write image: " << outFileName << '\n';
}

void cleanUp(void)
{
    // delete[] imageBuffer;
}

void displayItroAndInfo(void)
{
    std::cout << " Scene Information";
    std::cout << "\n   Number of triangles : " << triangleList.size();
    std::cout << "\n   Light               : ";
    if (isLightOn)
        std::cout << "on";
    else
        std::cout << "off";
    std::cout << "\n   Image size          : " << imageWidth << " X " << imageHeight;
    std::cout << "\n   Output file         : " << outFileName << "\n\n";
}

int main(int argc, char* argv[])
{
    if (argc != 3)
    {
        std::cerr << "Usage: " << argv[0] << " <scene-file> <output-file>\n";
        return 1;
    }

    const std::filesystem::path sceneFilePath = argv[1];
    outFileName = argv[2];

    if (!loadSceneFile(sceneFilePath.string()))
        return 1;

    displayItroAndInfo();

    // Pipeline: trace camera paths, build the photon lookup table, cast photons, then write output.
    const auto renderingStart = std::chrono::steady_clock::now();

    prepareIntersectionPointsList();

    std::cout << "raytracing:1:start\n";
    rayTracingPass();
    std::cout << "raytracing:1:done\nphotontracing:start\n";

    buildHashTable();

    photonCastingPass();

    std::cout << "photontracing:done\n";

    // return 0;

    fluxCorrection();

    const auto renderingEnd = std::chrono::steady_clock::now();
    const double totalRenderingTime =
        std::chrono::duration<double, std::milli>(renderingEnd - renderingStart).count();
    std::cout << "renderingtime:" << totalRenderingTime << std::endl;

    // Save result image buffer in file
    saveImage();
    cleanUp();
    std::cout << "done\n";

    return 0;
}