#include "ArrayRenderer.h"
#include <cuda_helper.hpp>
#include <utility>
#include <vector>
#include <Loader.hpp>
#include <array.hpp>
#include <common.hpp>
#include <common_device.hpp>
#include <importance.hpp>
#include <phase.hpp>

constexpr int NEW_SAMPLES = 30;

namespace ARRAYRENDERER {
    __device__ float ratioTrackingOptimized(CUDA::OptimizationParameters* optParams, CUDA::Volume* volume, float start,
                                            float end, float d, float sigma, CUDA::Ray ray, int numberOfSamples,
                                            uint2& random) {
        float4 newVals[150];
        int size = 1;
        float len = 0.f;
        float fullIntegral = 0.f;
        float dtImp = (end - start) / (float) numberOfSamples;

        // first value
        float ran = uniform_float(random);
        float startStrata = start;
        float endStrata = start + dtImp;
        float rand = startStrata + ran * (endStrata - startStrata);
        float value = sampleImpSDF(ray.eye + rand * ray.dir, volume->volumeData, volume->bbMin, volume->bbMax,
                                   optParams->thresholdMin, optParams->thresholdMax, optParams->sdf,
                                   volume->importanceSDF, optParams->position);
        fullIntegral += (rand - start) * value * value * 0.5f;
        float4 curr = make_float4(rand, value * value, 0, fullIntegral);

        ran = uniform_float(random);
        startStrata = start + 1 * dtImp;
        endStrata = start + 2 * dtImp;
        rand = startStrata + ran * (endStrata - startStrata);
        value = sampleImpSDF(ray.eye + rand * ray.dir, volume->volumeData, volume->bbMin, volume->bbMax,
                             optParams->thresholdMin, optParams->thresholdMax, optParams->sdf, volume->importanceSDF,
                             optParams->position);

        for (int a = 2; a < numberOfSamples; ++a) {
            ran = uniform_float(random);
            startStrata = start + a * dtImp;
            endStrata = start + (a + 1) * dtImp;
            float2 candidate = make_float2(rand, value * value);
            rand = startStrata + ran * (endStrata - startStrata);
            value = sampleImpSDF(ray.eye + rand * ray.dir, volume->volumeData, volume->bbMin, volume->bbMax,
                                 optParams->thresholdMin, optParams->thresholdMax, optParams->sdf,
                                 volume->importanceSDF, optParams->position);
            insert_light(newVals, 0, candidate, size, len, fullIntegral, curr);
            if (size > 150) {
                printf("array too small\n");
            }
        }

        newVals[size - 1] = curr;
        fullIntegral = curr.w + curr.y * curr.z + (end - curr.x) * curr.y * 0.5f;

        float t = start;
        float T = 1;
        do {
            float dt = -logf(uniform_float(random)) / sigma;
            t += dt;
            if (t >= d) {
                break;
            }
            float3 pos = ray.eye + t * ray.dir;
            float extinction = getEx(pos, volume->volumeData, volume->bbMin, volume->bbMax);
            float smallIntegral = 0;
            if (fabsf(optParams->r - optParams->q) > 0) {
                smallIntegral = getIntegralArray(t, newVals, size, 0, start);
            }
            float sumExtinction = optParams->q * fullIntegral + (optParams->r - optParams->q) * smallIntegral;
            float gi =
                    sampleImpSDF(pos, volume->volumeData, volume->bbMin, volume->bbMax, optParams->thresholdMin,
                                 optParams->thresholdMax, optParams->sdf, volume->importanceSDF, optParams->position);
            float exVec = getCol(extinction, &volume->transferFunction).w;
            float ex = (optParams->p * exVec) / (optParams->p + powf(1 - gi, 2 * optParams->lambda) * sumExtinction);
            T *= 1 - ex;
        } while (true);
        return T;
    }

    __device__ CUDA::PosSample samplePos(CUDA::Ray ray, CUDA::Volume* volume, CUDA::OptimizationParameters* optParams,
                                         uint2& random, float4* array, int size, int offset, float fullIntegral,
                                         bool debug, bool first, float tMax) {
        float sigma = volume->density * (volume->absorption + volume->scattering);

        float t = 0;
        while (true) {
            float dt = -logf(uniform_float(random)) / sigma;
            t += dt;
            if (t > tMax) {
                break;
            }
            float3 pos = ray.eye + t * ray.dir;
            float extinction = getEx(pos, volume->volumeData, volume->bbMin, volume->bbMax);
            float ex;
            float4 exVec;

            if (debug) {
                exVec = make_float4(0.93, 0.44, 0.08, 1);
                ex = sampleImpSDF(pos, volume->volumeData, volume->bbMin, volume->bbMax, optParams->thresholdMin,
                                  optParams->thresholdMax, optParams->sdf, volume->importanceSDF, optParams->position);
            } else {
                exVec = getCol(extinction, &volume->transferFunction);
                ex = exVec.w;
                if (optParams->view && first) {
                    float smallIntegral = 0;
                    if (fabsf(optParams->r - optParams->q) > 0) {
                        smallIntegral = getIntegralArray(t, array, size, offset, 0);
                    }
                    float sumExtinction = optParams->q * fullIntegral + (optParams->r - optParams->q) * smallIntegral;
                    float gi = sampleImpSDF(pos, volume->volumeData, volume->bbMin, volume->bbMax,
                                            optParams->thresholdMin, optParams->thresholdMax, optParams->sdf,
                                            volume->importanceSDF, optParams->position);
                    ex = (optParams->p * ex) / (optParams->p + powf(1 - gi, 2 * optParams->lambda) * sumExtinction);
                }
            }

            float pNull = fmaxf(1 - ex, 0);
            float mode = uniform_float(random);
            if (mode <= 1 - pNull) {
                return {make_float3(exVec), pos,
                        mode <= (ex * volume->absorption) / (volume->absorption + volume->scattering)};
            }
        }
        return {make_float3(-100), make_float3(-0), false};
    }

    __device__ float3 traceVolume(CUDA::Ray ray, CUDA::Volume* volume, CUDA::OptimizationParameters* optParams,
                                  CUDA::PointLight* lights, float numLights, uint2& random, float tMax, float4* array,
                                  int size, int offset, float fullIntegral, bool debug, int numSamples, bool sdf,
                                  bool position, bool multiscatter, bool full) {
        float3 color = make_float3(0);
        float sigma = volume->density * (volume->absorption + volume->scattering);
        bool first = true;

        while (true) {
            auto val = samplePos(ray, volume, optParams, random, array, size, offset, fullIntegral, debug, first, tMax);
            if (val.color.x < -99) {
                // escaped medium
                break;
            }
            if (val.absorb) {
                // absorbtion
                if (debug) {
                    return val.color;
                }
                return color;
            }
            // scattering
            float3 pos = val.pos;
            float lightPdf = 0;
            CUDA::PointLight light = SampleLights(lights, numLights, uniform_float(random), lightPdf);

            auto sample = SamplePointLight(pos, light);

            CUDA::Ray lightRay = {sample.pos, -sample.wi};
            float d = length(sample.pos - pos);
            auto inter = intersectAABB(lightRay, volume->bbMin, volume->bbMax);

            float T;
            if (debug) {
                val.color = make_float3(1);
                T = ratioTrackingDebug(lightRay, volume->volumeData, volume->bbMin, volume->bbMax, inter.tMin, d, sigma,
                                       random, optParams->thresholdMin, optParams->thresholdMax, sdf,
                                       volume->importanceSDF, position);
            } else {
                if (optParams->light && (first || full)) {
                    T = ratioTrackingOptimized(optParams, volume, inter.tMin, inter.tMax, d, sigma, lightRay,
                                               numSamples, random);
                } else {
                    T = ratioTracking(lightRay, volume->volumeData, volume->bbMin, volume->bbMax,
                                      &volume->transferFunction, inter.tMin, d, sigma, random);
                }
            }

            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);
            }

            float3 inscattering = sample.Li / (sample.pdf * lightPdf) * T * (f_v + f_s) * val.color;
            color += inscattering;

            if (!multiscatter) {
                return color;
            }
            float3 u = make_float3(uniform_float2(random), uniform_float(random));
            auto phaseSample = samplePhaseFunction(ray.dir, u, volume->g, 0, HenyeyGreenstein);

            ray.dir = phaseSample.dir;
            ray.eye = pos;

            auto interNewRay = intersectAABB(ray, volume->bbMin, volume->bbMax);
            if (!interNewRay.hit || interNewRay.tMin > 0.f) {
                // escaped medium
                break;
            }

            tMax = interNewRay.tMax;
            first = false;
        }

        // Background
        color += make_float3(1);

        return color;
    }
} // namespace ARRAYRENDERER

__global__ void buildArray(uint2 screen_size, CUDA::Camera camera, CUDA::Volume volume,
                           CUDA::OptimizationParameters optParams, unsigned long long int seed, int currSamples,
                           CUDA::ArrayValuesSmall* arraySizes, int3* offsets, float4* oldValues, float4* newValues,
                           float2* tempValues, bool bootstrapped) {
    CUDA::ThreadInfo ti;
    if (ti.threadId >= screen_size.x * screen_size.y) {
        return;
    }

    CUDA::ArrayValuesSmall sizeInt = arraySizes[ti.threadId];
    int size = sizeInt.size;
    int3 offset = offsets[ti.threadId];

    __shared__ float2 temp[NEW_SAMPLES * CUDA::THREADS_PER_BLOCK];
    unsigned int sharedIdCoalesced = ti.localThreadId;

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

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

    CUDA::ArrayValues arrayVals{};

    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) {
        return;
    }
    ray.eye += tMin * ray.dir;
    tMax -= tMin;
    tMin = 0;

    if (bootstrapped) {
        int s = 0;
        float start = tMin;
        for (int a = 0; a < size; ++a) {
            auto val = oldValues[offset.x + a];
            float startStrata = start;
            float endStrata = val.x - val.z;
            float ran = uniform_float(seedRandom);
            float t = startStrata + ran * (endStrata - startStrata);
            float value = sampleImpSDF(ray.eye + t * ray.dir, volume.volumeData, volume.bbMin, volume.bbMax,
                                       optParams.thresholdMin, optParams.thresholdMax, optParams.sdf,
                                       volume.importanceSDF, optParams.position);

            tempValues[offset.z + s++] = make_float2(t, value * value);
            start = val.x;

            if (val.z > 0) {
                startStrata = val.x - val.z;
                endStrata = val.x;
                ran = uniform_float(seedRandom);
                t = startStrata + ran * (endStrata - startStrata);
                value = sampleImpSDF(ray.eye + t * ray.dir, volume.volumeData, volume.bbMin, volume.bbMax,
                                     optParams.thresholdMin, optParams.thresholdMax, optParams.sdf,
                                     volume.importanceSDF, optParams.position);

                tempValues[offset.z + s++] = make_float2(t, value * value);
            }
        }
        if (size > 0) {
            float startStrata = start;
            float endStrata = tMax;
            float ran = uniform_float(seedRandom);
            float t = startStrata + ran * (endStrata - startStrata);
            float value = sampleImpSDF(ray.eye + t * ray.dir, volume.volumeData, volume.bbMin, volume.bbMax,
                                       optParams.thresholdMin, optParams.thresholdMax, optParams.sdf,
                                       volume.importanceSDF, optParams.position);

            tempValues[offset.z + s++] = make_float2(t, value * value);
            arrayVals = insertArray(size, offset.x, offset.y, offset.z, oldValues, newValues, tempValues, s, tMin, tMax,
                                    -1);
        }
    } else {
        float dtImp = (tMax - tMin) / (float) (NEW_SAMPLES);
        for (int a = 0; a < NEW_SAMPLES; ++a) {
            float startStrata = tMin + a * dtImp;
            float endStrata = tMin + (a + 1) * dtImp;
            float ran = uniform_float(seedRandom);
            float rand = startStrata + ran * (endStrata - startStrata);
            float value = sampleImpSDF(ray.eye + rand * ray.dir, volume.volumeData, volume.bbMin, volume.bbMax,
                                       optParams.thresholdMin, optParams.thresholdMax, optParams.sdf,
                                       volume.importanceSDF, optParams.position);
            temp[sharedIdCoalesced + a * CUDA::THREADS_PER_BLOCK] = make_float2(rand, value * value);
        }

        arrayVals = insertArray(size, offset.x, offset.y, sharedIdCoalesced, oldValues, newValues, temp, NEW_SAMPLES,
                                tMin, tMax, CUDA::THREADS_PER_BLOCK);
    }
    arraySizes[ti.threadId] = {arrayVals.size, arrayVals.fullIntegral};
}

__global__ void __maxnreg__(96)
        d_render(uint2 screen_size, CUDA::Camera camera, CUDA::Volume volume, CUDA::OptimizationParameters optParams,
                 cudaSurfaceObject_t surface, float4* dataArray, unsigned long long int seed, int currSamples,
                 int samplesPerFrame, bool prog, CUDA::ArrayValuesSmall* arraySizes, int3* offsets, float4* newValues,
                 bool debug, int numSamples, CUDA::PointLight* lights, int numLights, bool multiscatter, bool full,
                 float3 pos, float intensity) {
    CUDA::ThreadInfo ti;
    if (ti.threadId >= screen_size.x * screen_size.y) {
        return;
    }

    CUDA::ArrayValuesSmall arrayVals = arraySizes[ti.threadId];
    int3 offset = offsets[ti.threadId];

    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);

    if (numLights > 0) {
        lights[0] = {intensity, pos};
    }

    for (int sample = 1; sample <= samplesPerFrame; ++sample) {
        float3 color = make_float3(1);
        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) {
            ray.eye += tMin * ray.dir;
            color = ARRAYRENDERER::traceVolume(ray, &volume, &optParams, lights, numLights, seedRandom, tMax - tMin,
                                               newValues, arrayVals.size, offset.y, arrayVals.fullIntegral, debug,
                                               numSamples, optParams.sdf, optParams.position, multiscatter, full);
            if (isnan(color.x)) {
                printf("NaN\n");
                color = make_float3(0);
            }
        }

        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);
}

ArrayRenderer::ArrayRenderer(int width, int height, int maxSamples, int samplesPerFrame,
                             std::vector<CUDA::PointLight> lightsVector) {
    initRenderer(width, height, maxSamples, samplesPerFrame);

    cudaStreamCreate(&streamKernel);

    // make first light moveable
    if (lightsVector.size() > 0) {
        pos = lightsVector[0].pos;
        temp_pos[0] = pos.x;
        temp_pos[1] = pos.y;
        temp_pos[2] = pos.z;
        intensity = lightsVector[0].intensity;
        temp_intensity = intensity;
    }

    numLights = lightsVector.size();
    CUDA::malloc(lights, numLights);
    CUDA::memcopyToDevice(lights, lightsVector.data(), numLights);
}

void ArrayRenderer::render(cudaSurfaceObject_t surface, CUDA::Camera h_camera, CUDA::Volume h_volume,
                           CUDA::OptimizationParameters h_optParams) {
    if (currSamples < maxSamples) {
        bool bootstrappe = false;
        if (bootstrapped) {
            bootstrappe = (currSamples / samplesPerFrame) % 4 != 0;
            if (maxSize > 200) {
                bootstrappe = false;
            }
        }

        if (h_optParams.view) {
            int newAllocationSize = 0;
            int newAllocationSizeTemp = 0;

            maxSize = 0;
            for (int j = 0; j < width * height; ++j) {
                int3 oldOffset = offsets[j];
                offsets[j] = make_int3(oldOffset.y, newAllocationSize, newAllocationSizeTemp);
                int size = arraySizes[j].size;
                if (bootstrappe) {
                    newAllocationSize += 4 * size + 1;
                    newAllocationSizeTemp += 2 * size + 1;
                } else {
                    newAllocationSize += 2 * size + NEW_SAMPLES;
                }

                maxSize = max(maxSize, size);
            }
            assert(newAllocationSize >= 0);

            CUDA::mallocAsync(newArray, newAllocationSize, streamKernel);
            CUDA::mallocAsync(tempArray, newAllocationSizeTemp, streamKernel);
            CUDA::memcopyToDeviceAsync(offsetsDev, offsets, width * height, streamKernel);

            buildArray<<<CUDA::getBlockCount(width * height), CUDA::THREADS_PER_BLOCK, 0, streamKernel>>>(
                    make_uint2(width, height), h_camera, h_volume, h_optParams, dev(), currSamples, arraySizesDev,
                    offsetsDev, oldArray, newArray, tempArray, bootstrappe);
            CUDA::freeAsync(oldArray, streamKernel);
            CUDA::freeAsync(tempArray, streamKernel);
        }

        d_render<<<CUDA::getBlockCount(width * height), CUDA::THREADS_PER_BLOCK, 0, streamKernel>>>(
                make_uint2(width, height), h_camera, h_volume, h_optParams, surface, deviceData, dev(), currSamples,
                samplesPerFrame, progressive, arraySizesDev, offsetsDev, newArray, debug, numSamples, lights, numLights,
                multiscatter, full, pos, intensity);

        CUDA::memcopyToHostAsync(arraySizes, arraySizesDev, width * height, streamKernel);

        CUDA::checkCudaError(cudaStreamSynchronize(streamKernel));

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

        oldArray = newArray;
        newArray = nullptr;

        progressive = true;
        currSamples += samplesPerFrame;
    }
}

void ArrayRenderer::renderSimple(cudaSurfaceObject_t surface, CUDA::Camera h_camera, CUDA::Volume h_volume,
                                 CUDA::OptimizationParameters h_optParams) {
    CUDA::memset(arraySizesDev, width * height);
    for (int i = 0; i < width * height; ++i) {
        offsets[i] = make_int3(0, i * NEW_SAMPLES, i * NEW_SAMPLES);
    }
    CUDA::memcopyToDevice(offsetsDev, offsets, width * height);

    CUDA::free(newArray);

    if (h_optParams.view) {
        CUDA::malloc(newArray, NEW_SAMPLES * width * height);
        CUDA::malloc(tempArray, NEW_SAMPLES * width * height);
        buildArray<<<CUDA::getBlockCount(width * height), CUDA::THREADS_PER_BLOCK, 0, streamKernel>>>(
                make_uint2(width, height), h_camera, h_volume, h_optParams, 0, 0, arraySizesDev, offsetsDev, nullptr,
                newArray, tempArray, false);
    }

    d_render<<<CUDA::getBlockCount(width * height), CUDA::THREADS_PER_BLOCK, 0, streamKernel>>>(
            make_uint2(width, height), h_camera, h_volume, h_optParams, surface, deviceData, 0, 0, samplesPerFrame,
            false, arraySizesDev, offsetsDev, newArray, debug, numSamples, lights, numLights, multiscatter, full, pos,
            intensity);
    CUDA::checkCudaError(cudaDeviceSynchronize());
    CUDA::free(newArray);
    CUDA::free(tempArray);

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

    currSamples = samplesPerFrame;
}

void ArrayRenderer::drawGui(float deltaTime, float frametime) {
    Renderer::drawGui(deltaTime, frametime);
    ImGui::Begin("Parameters");

    ImGui::InputInt("Num Samples light", &temp_numSamples, 1.f);
    ImGui::Checkbox("Debug", &temp_debug);
    if (ImGui::CollapsingHeader("Renderer options", ImGuiTreeNodeFlags_DefaultOpen)) {
        ImGui::Checkbox("Multi-Scattering", &temp_multiscatter);
        if (temp_multiscatter) {
            ImGui::Checkbox("Full Transport", &temp_full);
        }
        ImGui::Checkbox("Bootstrapped", &temp_bootstrapped);
    }
    ImGui::SliderFloat3("Light Pos", temp_pos, -3, 3);
    ImGui::SliderFloat("Light intensity", &temp_intensity, 0, 100);

    if (temp_numSamples != numSamples || temp_multiscatter != multiscatter || temp_full != full ||
        temp_debug != debug || temp_bootstrapped != bootstrapped || temp_pos[0] != pos.x || temp_pos[1] != pos.y ||
        temp_pos[2] != pos.z || temp_intensity != intensity) {
        submit();
    }

    if (ImGui::Button("Save PNG", {100, 20})) {
        saveImage();
    }
    ImGui::End();
}

void ArrayRenderer::submit() {
    full = temp_full;
    debug = temp_debug;
    bootstrapped = temp_bootstrapped;
    multiscatter = temp_multiscatter;

    numSamples = temp_numSamples;

    pos = make_float3(temp_pos[0], temp_pos[1], temp_pos[2]);
    intensity = temp_intensity;
    resetRendering(width, height);
}

void ArrayRenderer::resetRendering(int width, int height) {
    Renderer::resetRendering(width, height);

    CUDA::free(arraySizesDev);
    CUDA::malloc(arraySizesDev, width * height);
    CUDA::memset(arraySizesDev, width * height);

    CUDA::free(offsetsDev);
    CUDA::malloc(offsetsDev, width * height);

    delete[] offsets;
    offsets = new int3[width * height];
    for (int i = 0; i < width * height; ++i) {
        offsets[i] = make_int3(0, i * NEW_SAMPLES, i * NEW_SAMPLES);
    }
    CUDA::memcopyToDevice(offsetsDev, offsets, width * height);

    delete[] arraySizes;
    arraySizes = new CUDA::ArrayValuesSmall[width * height];
    memset(arraySizes, 0, width * height * sizeof(CUDA::ArrayValuesSmall));

    CUDA::free(tempArray);
    CUDA::free(oldArray);
    CUDA::free(newArray);
}
