#include "RaymarchRenderer.h"
#include <common_device.hpp>
#include <importance.hpp>
#include <phase.hpp>

namespace RAYMARCH {
    __device__ void firstPass(CUDA::Ray ray, CUDA::Volume* volume, float start, float entry, float exit, float stepSize,
                              float& gr, float& gb, float& gf, float thresholdMin, float thresholdMax, bool sdf,
                              bool position) {
        // first pass
        gr = 0;
        gb = 0;
        gf = 0;
        float t = entry;
        while (t < exit) {
            float3 pos = ray.eye + t * ray.dir;
            float gi = sampleImpSDF(pos, volume->volumeData, volume->bbMin, volume->bbMax, thresholdMin, thresholdMax,
                                    sdf, volume->importanceSDF, position);
            gr += gi * gi;

            if (t < start) {
                gb += gi * gi;
            }
            t += stepSize;
        }
    }

    __device__ float optimizeLight(CUDA::Ray ray, CUDA::Volume* volume, uint2& random, float d, float stepSize,
                                   float thresholdMin, float thresholdMax, float p, float q, float r, float lambda,
                                   bool jittered, bool sdf, bool position) {
        float transmittance = 1;
        auto lightHit = intersectAABB(ray, volume->bbMin, volume->bbMax);
        if (lightHit.hit) {
            // volume hit
            float gr = 0;
            float gb = 0;
            float gf = 0;
            float dt = 0;
            if (jittered) {
                dt = (1 - uniform_float(random)) * stepSize;
            }
            firstPass(ray, volume, 0, lightHit.tMin + dt, lightHit.tMax, stepSize, gr, gb, gf, thresholdMin,
                      thresholdMax, sdf, position);

            float t = lightHit.tMin;
            while (t < d) {
                float3 pos = ray.eye + (t + dt) * ray.dir;
                float extinction = getEx(pos, volume->volumeData, volume->bbMin, volume->bbMax);
                float gi = sampleImpSDF(pos, volume->volumeData, volume->bbMin, volume->bbMax, thresholdMin,
                                        thresholdMax, sdf, volume->importanceSDF, position);
                float4 sigmaVec = getCol(extinction, &volume->transferFunction);
                float sigma = sigmaVec.w;
                gf = gr - gi * gi - gb;
                float sigmaOpti =
                        (p * sigma) / (p + powf(1 - gi, 2 * lambda) * (r * gb * stepSize + q * gf * stepSize));
                float pointTransmittance =
                        expf(-sigmaOpti * volume->density * (volume->absorption + volume->scattering) * stepSize);
                transmittance *= pointTransmittance;
                gb += gi * gi;

                t += stepSize;
            }
        }
        return transmittance;
    }

    __device__ float3 traceVolume(CUDA::Ray ray, CUDA::Volume* volume, CUDA::PointLight* lights, int numLights,
                                  uint2& random, float tMin, float tMax, CUDA::OptimizationParameters optParams,
                                  float stepSizeView, float stepSizeLight, bool jittered) {
        float3 color = make_float3(0);
        float transmittance = 1;

        float d = 2;

        float gr = 0;
        float gb = 0;
        float gf = 0;
        float dt = 0;
        if (jittered) {
            dt = (1 - uniform_float(random)) * stepSizeView;
        }
        firstPass(ray, volume, 0, tMin + dt, tMax, stepSizeView, gr, gb, gf, optParams.thresholdMin,
                  optParams.thresholdMax, optParams.sdf, optParams.position);

        float t = tMin;
        while (t < tMax) {
            float3 pos = ray.eye + (t + dt) * ray.dir;
            float extinction = getEx(pos, volume->volumeData, volume->bbMin, volume->bbMax);
            float4 sigmaVec = getCol(extinction, &volume->transferFunction);
            float ex;
            if (optParams.view) {
                float gi =
                        sampleImpSDF(pos, volume->volumeData, volume->bbMin, volume->bbMax, optParams.thresholdMin,
                                     optParams.thresholdMax, optParams.sdf, volume->importanceSDF, optParams.position);

                float sigma = sigmaVec.w;
                gf = gr - gi * gi - gb;

                ex = (optParams.p * sigma) /
                     (optParams.p + powf(1 - gi, 2 * optParams.lambda) *
                                            (optParams.r * gb * stepSizeView + optParams.q * gf * stepSizeView));

                gb += gi * gi;
            } else {
                ex = sigmaVec.w;
            }

            float pointTransmittance =
                    expf(-ex * volume->density * (volume->absorption + volume->scattering) * stepSizeView);
            transmittance *= pointTransmittance;

            float sigma_a = ex * volume->density * volume->absorption;
            float sigma_s = ex * volume->density * volume->scattering;

            // emmision
            float3 emmision = make_float3(0) * sigma_a;

            // inscattering
            float3 inscattering = make_float3(0);
            for (int i = 0; i < numLights; ++i) {
                float lightTransmittance = 1;
                auto sample = SamplePointLight(pos, lights[i]);
                auto lightHit = intersectAABB({pos, sample.wi}, volume->bbMin, volume->bbMax);
                if (optParams.light) {
                    lightTransmittance = optimizeLight({sample.pos, -sample.wi}, volume, random,
                                                       length(pos - sample.pos), stepSizeLight, optParams.thresholdMin,
                                                       optParams.thresholdMax, optParams.p, optParams.q, optParams.r,
                                                       optParams.lambda, jittered, optParams.sdf, optParams.position);
                } else {
                    if (lightHit.hit) {
                        float tLight = lightHit.tMin;
                        float dtLight = 0;
                        if (jittered) {
                            dtLight = (1 - uniform_float(random)) * stepSizeLight;
                        }
                        while (tLight < lightHit.tMax) {

                            float3 posLight = pos + (tLight + dtLight) * sample.wi;
                            float extinctionLight = getEx(posLight, volume->volumeData, volume->bbMin, volume->bbMax);
                            float4 sigmaVecLight = getCol(extinctionLight, &volume->transferFunction);
                            float sigma = sigmaVecLight.w;
                            float pointTransmittanceLight =
                                    expf(-sigma * volume->density * (volume->absorption + volume->scattering) *
                                         stepSizeLight);
                            lightTransmittance *= pointTransmittanceLight;

                            tLight += stepSizeLight;
                        }
                    }
                }

                float f_v = evalPhaseFunction(dot(ray.dir, sample.wi), volume->g, 0, HenyeyGreenstein);
                float f_s = 0;
                if (optParams.phong) {
                    f_s = computeBlinnPhong(pos, volume->volumeData, volume->bbMin, volume->bbMax, volume->gridSize,
                                            -ray.dir, sample.wi, optParams.shiny, optParams.scale);
                }

                inscattering +=
                        sample.Li / sample.pdf * lightTransmittance * (f_v + f_s) * sigma_s * make_float3(sigmaVec);
            }
            color += (emmision + inscattering) * transmittance * stepSizeView;

            if (transmittance < 1e-3) {
                if (uniform_float(random) > 1.f / d) {
                    break;
                }
                transmittance *= d;
            }

            t += stepSizeView;
        }

        color += make_float3(1) * transmittance;

        return color;
    }

} // namespace RAYMARCH

__global__ void d_render(uint2 screen_size, cudaSurfaceObject_t surface, float4* dataArray, unsigned long long int seed,
                         CUDA::Camera camera, int currSamples, int samplesPerFrame, CUDA::Volume volume, bool prog,
                         CUDA::OptimizationParameters optParams, float stepSizeView, float stepSizeLight,
                         CUDA::PointLight* lights, int numLights, bool jittered) {
    CUDA::ThreadInfo ti;
    if (ti.threadId >= screen_size.x * screen_size.y)
        return;

    unsigned int i = ti.threadId % screen_size.x;
    unsigned int j = ti.threadId / screen_size.x;

    // read prev color
    float4 outColor = surf2Dread<float4>(surface, i * sizeof(float4), j);
    if (!prog) {
        outColor = make_float4(0, 0, 0, 1);
    }

    uint2 seedRandom = make_uint2(i, j) ^ make_uint2(currSamples << 16, (currSamples + seed) << 16);

    for (int sample = 1; sample <= samplesPerFrame; ++sample) {
        float3 color;
        CUDA::Ray ray = generateRayPinhole((float) i, (float) j, (float) screen_size.x, (float) screen_size.y,
                                           camera.eye, camera.U, camera.V, camera.W);

        auto [hit, tMin, tMax] = intersectAABB(ray, volume.bbMin, volume.bbMax);

        if (hit) {
            color = RAYMARCH::traceVolume(ray, &volume, lights, numLights, seedRandom, tMin, tMax, optParams,
                                          stepSizeView, stepSizeLight, jittered);
            if (isnan(color.x)) {
                printf("NaN\n");
                color = make_float3(0);
            }
        } else {
            color = make_float3(1);
        }

        outColor += (make_float4(color, 1) - outColor) / (float) (currSamples + sample);
    }

    outColor.w = 1;

    dataArray[j * screen_size.x + i] = outColor;

    // write back
    surf2Dwrite(outColor, surface, i * sizeof(float4), j);
}

void RaymarchRenderer::render(cudaSurfaceObject_t surface, CUDA::Camera camera, CUDA::Volume volume,
                              CUDA::OptimizationParameters optParams) {
    if (currSamples < maxSamples) {
        d_render<<<CUDA::getBlockCount(width * height), CUDA::THREADS_PER_BLOCK>>>(
                make_uint2(width, height), surface, deviceData, dev(), camera, currSamples, samplesPerFrame, volume,
                progressive, optParams, stepSizeView, stepSizeLight, lights, numLights, jittered);

        CUDA::checkCudaError(cudaDeviceSynchronize());

        // Checking failures on kernel
        CUDA::checkCudaKernelError(cudaGetLastError());

        progressive = true;
        currSamples += samplesPerFrame;
    }
}

void RaymarchRenderer::renderSimple(cudaSurfaceObject_t surface, CUDA::Camera camera, CUDA::Volume volume,
                                    CUDA::OptimizationParameters optParams) {
    d_render<<<CUDA::getBlockCount(width * height), CUDA::THREADS_PER_BLOCK>>>(
            make_uint2(width, height), surface, deviceData, 0, camera, 0, 1, volume, false, optParams, stepSizeView,
            stepSizeLight, lights, numLights, jittered);

    // Checking failures on kernel
    CUDA::checkCudaKernelError(cudaGetLastError());

    CUDA::checkCudaError(cudaDeviceSynchronize());
    currSamples = 1;
}

void RaymarchRenderer::drawGui(float deltaTime, float frametime) {
    ImGui::Begin("Parameters");
    Renderer::drawGui(deltaTime, frametime);
    if (ImGui::CollapsingHeader("Optimization options", ImGuiTreeNodeFlags_DefaultOpen)) {
        ImGui::SliderFloat("Step Size view", &temp_stepSizeView, 0.001, 0.1f);
        ImGui::SliderFloat("Step Size light", &temp_stepSizeLight, 0.001, 0.1f);
        ImGui::Checkbox("Jittered", &temp_jittered);
    }

    if (stepSizeView != temp_stepSizeView || stepSizeLight != temp_stepSizeLight || jittered != temp_jittered) {
        submit();
    }
    if (ImGui::Button("Save PNG", {100, 20})) {
        saveImage();
    }
    ImGui::End();
}

void RaymarchRenderer::submit() {
    stepSizeView = temp_stepSizeView;
    stepSizeLight = temp_stepSizeLight;
    jittered = temp_jittered;

    resetRendering(width, height);
}
