/************************************************************
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
*************************************************************/

#pragma once

#include "./3rdparty/stb_collection/stb_image.h"
#include "vector3d.h"
#include <algorithm>
#include <iostream>
#include <sstream>
#include <stdio.h>
#include <vector>

namespace PRRPM
{

    // Structures
    struct Ray
    {
        Vector3d origin;
        Vector3d direction;
        int objectID;
        ui32 objectType;
        Vector3d intersectionPosition;
        f32 distance;
    };

    struct Triangle
    {
        Vector3d point1;
        Vector3d point2;
        Vector3d point3;
        Vector3d color;
        Vector3d normal;
        ui8 materialID;
    };

    struct Sphere
    {
        Vector3d position;
        float radius;
        Vector3d color;
        ui32 numberOfSpheres;
        ui8 materialID;
    };

    struct IntersectionPoint
    {
        Vector3d color;
        Vector3d positionAndOrigin;
        Vector3d normal;
        Vector3d fluxAndDirection;
        f32 radiusSquare;
        i32 nubmerOfPhotonAndMaterialID;
        int pixelIndex;
        int cellID;
    };

    struct LightSource
    {
        ui8 type;
        Vector3d position;
        Vector3d intensity;
        Vector3d size;
    };

    extern ui32 imageWidth;
    extern ui32 imageHeight;
    extern ui32 numberOfPhotons;
    extern ui32 numberOfPhotonPass;
    extern ui32 numberOfBounces;
    inline constexpr f32 maxSceneDepth = 1000.0f;
    inline constexpr f32 alpha = 0.7f;
    extern std::string outFileName;
    extern std::vector<Vector3d> imageBufferList;
    extern bool isLightOn;
    extern bool isPhotonMapping;
    extern bool isGUI;
    inline constexpr f32 refractiveIndexIn = 1.0f;
    inline constexpr f32 refractiveIndexOut = 1.5f;

    inline constexpr int gridSize = 16000000;
    inline constexpr f32 irridianceRadiusSquare = 0.000016f;
    inline constexpr int numberOfGridCell = 250;
    inline constexpr int numberOfGridCellXY = 252;
    inline constexpr int numberOfGridCellXMY = 63000;

    extern Vector3d neighborsDirection[27];
    extern std::vector<Triangle> triangleList;
    extern std::vector<Sphere> sphereList;
    extern std::vector<LightSource> lightSourcesList;
    extern std::vector<IntersectionPoint> intersectionPointsList;
    extern std::vector<IntersectionPoint> secondIntersectionPointsList;
    extern std::vector<int> hastValueList;

    inline bool compareByCellID(const IntersectionPoint& firstPoint,
                                const IntersectionPoint& secondPoint)
    {
        return (firstPoint.cellID < secondPoint.cellID);
    }

    inline bool compareByMetalID(const IntersectionPoint& firstPoint,
                                 const IntersectionPoint& secondPoint)
    {
        return (firstPoint.nubmerOfPhotonAndMaterialID < secondPoint.nubmerOfPhotonAndMaterialID);
    }

    inline bool isCellIDOut(const IntersectionPoint& point)
    {
        return (point.cellID == -1);
    }

    inline bool isMetalIDOut(const IntersectionPoint& point)
    {
        return (point.nubmerOfPhotonAndMaterialID == -2);
    }

    inline int computeCellID(const Vector3d& position)
    {
        int x = (position.x + 0.6) * numberOfGridCell;
        int y = (position.y + 0.6) * numberOfGridCell;
        int z = (position.z - 1.0) * numberOfGridCell;
        return (z * numberOfGridCellXMY + y * numberOfGridCellXY + x);
    }

    inline void generatePhotonRayFromLighSource(Ray* ray, const LightSource& lightSource,
                                                std::default_random_engine& randomNumberGenerator)
    {
        std::uniform_real_distribution<f32> uniform(0.0, 1.0);

        if (lightSource.type == 0)
        {
            f32 u = MathUtility::TwoPI * uniform(randomNumberGenerator);
            f32 v = acosf(sqrtf(uniform(randomNumberGenerator))) * 2.0;
            f32 w = sinf(v);
            ray->direction = Vector3d(cosf(u) * w, cosf(v), sinf(u) * w);
            ray->origin = lightSource.position;
        }
        else if (lightSource.type == 1)
        {
            std::uniform_real_distribution<f32> uniformX(-lightSource.size.x, lightSource.size.x);
            std::uniform_real_distribution<f32> uniformZ(-lightSource.size.z, lightSource.size.z);
            std::uniform_real_distribution<f32> uniformNP(-1.0, 1.0);

            f32 u = uniform(randomNumberGenerator);
            f32 v = MathUtility::TwoPI * uniform(randomNumberGenerator);
            f32 uSquare = sqrtf(u);
            f32 x = cosf(v) * uSquare;
            f32 y = sin(v) * uSquare;
            f32 z = sqrtf(1.0 - u);

            ray->direction = Vector3d(x, -y, z * uniformNP(randomNumberGenerator));

            ray->origin.x = uniformX(randomNumberGenerator);
            ray->origin.y = 0;
            ray->origin.z = uniformZ(randomNumberGenerator);
            ray->origin = ray->origin + lightSource.position;
        }

        ray->direction.normalize();
    }

    inline void intersectTriangle(Ray* ray, unsigned int triangleIndex)
    {

        Vector3d edge1 = triangleList[triangleIndex].point2 - triangleList[triangleIndex].point1;
        Vector3d edge2 = triangleList[triangleIndex].point3 - triangleList[triangleIndex].point1;

        Vector3d pVector = ray->direction.cross(edge2);
        f32 determinat = edge1.dot(pVector);

        if (determinat < 0.000001)
            return;

        Vector3d tVector = ray->origin - triangleList[triangleIndex].point1;

        f32 u = tVector.dot(pVector);

        if (u < 0.0 || u > determinat)
            return;

        Vector3d qVector = tVector.cross(edge1);

        f32 v = ray->direction.dot(qVector);

        if (v < 0.0 || (u + v) > determinat)
            return;

        float inverseDeterminat = 1.0f / determinat;

        float t = edge2.dot(qVector) * inverseDeterminat;
        u *= inverseDeterminat;
        v *= inverseDeterminat;

        ray->distance = t;
        ray->intersectionPosition = triangleList[triangleIndex].point1 + (edge1 * u) + (edge2 * v);
    }

    inline void intersectSphere(Ray* ray, unsigned int sphereIndex)
    {
        Vector3d dist = sphereList[sphereIndex].position - ray->origin;
        f32 b = ray->direction.dot(dist);
        f32 delta = b * b - dist.dot(dist) +
                    (sphereList[sphereIndex].radius * sphereList[sphereIndex].radius);
        if (delta < 0.0f)
            return;
        f32 t0 = b - sqrt(delta);
        f32 t1 = b + sqrt(delta);
        f32 distance = maxSceneDepth;
        if ((t0 > 0.001f) && (t0 < distance))
            distance = t0;
        else if ((t1 > 0.001f) && (t1 < distance))
            distance = t1;

        ray->distance = distance;
    }

    inline void intersectObjects(Ray* ray)
    {
        Ray tempRay;
        tempRay.direction = ray->direction;
        tempRay.origin = ray->origin;
        tempRay.distance = maxSceneDepth;
        ray->distance = maxSceneDepth;

        for (unsigned int iTriangle = 0; iTriangle < triangleList.size(); ++iTriangle)
        {
            intersectTriangle(&tempRay, iTriangle);
            if (tempRay.distance < ray->distance)
            {
                ray->intersectionPosition = tempRay.intersectionPosition;
                ray->objectID = iTriangle;
                ray->objectType = 0;
                ray->distance = tempRay.distance;
            }
        }

        tempRay.objectID = -1;
        tempRay.distance = maxSceneDepth;

        for (unsigned int iSphere = 0; iSphere < sphereList.size(); ++iSphere)
        {
            intersectSphere(&tempRay, iSphere);
            if (tempRay.distance < ray->distance)
            {
                ray->objectID = iSphere;
                ray->objectType = 1;
                ray->distance = tempRay.distance;
                ray->intersectionPosition = (ray->direction * ray->distance) + ray->origin;
            }
        }
    }

    inline Vector3d getSurfaceNormal(const Ray& ray)
    {
        Vector3d surfaceNormal;
        if (ray.objectType == 0)
            surfaceNormal = triangleList[ray.objectID].normal;
        else
        {
            surfaceNormal = ray.intersectionPosition - sphereList[ray.objectID].position;
            surfaceNormal.normalize();
        }

        return surfaceNormal;
    }

    inline Vector3d getObjectColor(ui32 objectID, ui32 objectType)
    {
        if (objectType == 0)
            return triangleList[objectID].color;
        else
            return sphereList[objectID].color;
    }

    inline ui8 getObjectMaterial(ui32 objectID, ui32 objectType)
    {
        if (objectType == 0)
            return triangleList[objectID].materialID;
        else
            return sphereList[objectID].materialID;
    }

    inline void calculateReflectionFractionDirection(Vector3d rayDirection, Vector3d normal,
                                                     Vector3d* reflectionDirection,
                                                     Vector3d* refractionDirection,
                                                     f32* cosOutAngle, f32* reflectionIntensity)
    {
        Vector3d secondNormal;
        if (normal.dot(rayDirection) < 0)
            secondNormal = normal;
        else
            secondNormal = normal * -1;

        *reflectionDirection = rayDirection - normal * normal.dot(rayDirection) * 2.0f;

        bool isRayGoingToObject = (normal.dot(secondNormal) > 0.0f);
        f32 refractiveIndex = isRayGoingToObject ? refractiveIndexIn / refractiveIndexOut
                                                 : refractiveIndexOut / refractiveIndexIn;
        f32 cosInAngle = rayDirection.dot(secondNormal);
        *cosOutAngle = 1.0f - refractiveIndex * refractiveIndex * (1.0f - cosInAngle * cosInAngle);

        if (*cosOutAngle >= 0)
        {
            *refractionDirection = rayDirection * refractiveIndex -
                                   normal * ((isRayGoingToObject ? 1 : -1) *
                                             (cosInAngle * refractiveIndex + sqrt(*cosOutAngle)));
            refractionDirection->normalize();
            f32 reflectivity = pow((refractiveIndex - 1), 2) / pow((refractiveIndex + 1), 2);
            f32 angle = (isRayGoingToObject ? -cosInAngle : refractionDirection->dot(normal));
            *reflectionIntensity = reflectivity + (1.0f - reflectivity) * pow((1.0f - angle), 5);
        }
    }
} // namespace PRRPM