#pragma once
#include "cuda_helper.hpp"
#include "common.hpp"
#include "importance.hpp"

__device__ inline float getEx(float3 p, cudaTextureObject_t volumeData, float3 bbMin, float3 bbMax) {
    float x = shiftInterval01(p.x, bbMin.x, bbMax.x);
    float y = shiftInterval01(p.y, bbMin.y, bbMax.y);
    float z = shiftInterval01(p.z, bbMin.z, bbMax.z);

    return tex3D<float>(volumeData, x, y, z);
}

__device__ inline float4 getCol(const float t, CUDA::TF* transferFunction) {
    if (transferFunction->size == 0) {
        return make_float4(0);
    }
    if (t < transferFunction->minValue) {
        return __ldg(&transferFunction->data[0]);
    }
    if (t > transferFunction->maxValue) {
        return __ldg(&transferFunction->data[transferFunction->size - 1]);
    }
    float x = (t - transferFunction->minValue) / (transferFunction->maxValue - transferFunction->minValue);
    x *= (float) transferFunction->size - 1.f;
    int idx = clamp(int(x), 0, transferFunction->size - 2);
    x -= (float) idx;
    float4 v0 = __ldg(&transferFunction->data[idx]);
    float4 v1 = __ldg(&transferFunction->data[idx + 1]);
    float4 r = (1.f - x) * v0 + x * v1;
    return r;
}

__device__ inline float ratioTracking(const CUDA::Ray& ray, cudaTextureObject_t volumeData, float3 bbMin, float3 bbMax,
                                      CUDA::TF* transferFunction, float tMin, float d, float sigma, uint2& rand) {
    float random = uniform_float(rand);
    float t = tMin - logf(random) / sigma;
    random = uniform_float(rand);
    float3 pos = ray.eye + t * ray.dir;
    float extinction = getEx(pos, volumeData, bbMin, bbMax);
    float T = 1;

    while (t < d) {
        t -= logf(random) / sigma;
        float4 vec = getCol(extinction, transferFunction);
        random = uniform_float(rand);
        pos = ray.eye + t * ray.dir;
        extinction = getEx(pos, volumeData, bbMin, bbMax);
        T *= 1 - vec.w;
    }
    return T;
}

__device__ inline float ratioTrackingDebug(const CUDA::Ray& ray, cudaTextureObject_t volumeData, float3 bbMin,
                                           float3 bbMax, float tMin, float d, float sigma, uint2& random,
                                           float thresholdMin, float thresholdMax, bool sdf,
                                           cudaTextureObject_t sdfTexture, bool position) {
    float rand = uniform_float(random);
    float t = tMin - logf(rand) / sigma;
    rand = uniform_float(random);
    float3 pos = ray.eye + t * ray.dir;
    float T = 1;

    while (t < d) {
        t -= logf(rand) / sigma;
        float ex = sampleImpSDF(pos, volumeData, bbMin, bbMax, thresholdMin, thresholdMax, sdf, sdfTexture, position);
        rand = uniform_float(random);
        pos = ray.eye + t * ray.dir;
        T *= 1 - ex;
    }
    return T;
}

__device__ inline float read_from_texture(float3 pos, cudaTextureObject_t volumeData, float3 bbMin, float3 bbMax,
                                          float3 dt, int x, int y, int z) {
    float3 sample = make_float3(pos.x + x * dt.x, pos.y + y * dt.y, pos.z + z * dt.z);
    float ex = getEx(sample, volumeData, bbMin, bbMax);
    return ex;
}

__device__ inline float3 computeGradient(float3 pos, cudaTextureObject_t volumeData, float3 bbMin, float3 bbMax,
                                         float3 dt) {
    const int w[9] = {1, 3, 1, 3, 6, 3, 1, 3, 1};
    float3 d = make_float3(0.0f);
    int wc = 0;
    for (int i = -1; i < 2; i++) {
        for (int j = -1; j < 2; j++) {
            d.x += read_from_texture(pos, volumeData, bbMin, bbMax, dt, +1, +i, +j) * w[wc];
            d.x -= read_from_texture(pos, volumeData, bbMin, bbMax, dt, -1, +i, +j) * w[wc];
            d.y += read_from_texture(pos, volumeData, bbMin, bbMax, dt, +i, +1, +j) * w[wc];
            d.y -= read_from_texture(pos, volumeData, bbMin, bbMax, dt, +i, -1, +j) * w[wc];
            d.z += read_from_texture(pos, volumeData, bbMin, bbMax, dt, +i, +j, +1) * w[wc];
            d.z -= read_from_texture(pos, volumeData, bbMin, bbMax, dt, +i, +j, -1) * w[wc];
            wc++;
        }
    }
    d /= 22.0f;
    return d;
}

__device__ inline float computeBlinnPhong(float3 pos, cudaTextureObject_t volumeData, float3 bbMin, float3 bbMax,
                                          float3 dt, float3 V, float3 L, float shiny, float scale) {
    float3 gradient = computeGradient(pos, volumeData, bbMin, bbMax, dt);
    float ks = length(gradient);
    if (ks > 0) {
        gradient /= -ks;
        float3 H = normalize(L + V);
        return (shiny + 1.f) / (2.f * Pi) * ks * fmaxf(dot(gradient, L), 0) * powf(dot(H, gradient), shiny) * scale;
    }
    return 0;
}
