#pragma once
#include "cuda_helper.hpp"

enum PhaseFunction { HenyeyGreenstein};

__host__ __device__ inline float evalHenyeyGreenstein(float cosTheta, float g) {
    float denom = 1 + g * g - 2 * g * cosTheta;
    return Inv4Pi * (1 - g * g) / (denom * sqrtf(denom));
}

__host__ __device__ inline void sampleHenyeyGreenstein(float rand, float& cosTheta, float& value, float& pdf, float g) {
    if (abs(g) < 1e-3f) {
        cosTheta = 2 * rand - 1;
    } else {
        float t = (1.f - g * g) / (1.f - g + 2.f * g * rand);
        cosTheta = (1.f + g * g - t * t) / (2 * g);
    }
    value = evalHenyeyGreenstein(cosTheta, g);
    pdf = value;
}

__host__ __device__ inline float evalPhaseFunction(float cosTheta, float g, float d, PhaseFunction phase_function) {
    switch (phase_function) {
        case HenyeyGreenstein:
            return evalHenyeyGreenstein(cosTheta, g);
        default:
            return Inv4Pi;
    }
}

__host__ __device__ inline float3 to_world(float3 vec, float3 normal) {
    float sign = copysignf(1.f, normal.z);
    float a = -1.f / (sign + normal.z);
    float b = normal.x * normal.y * a;
    float3 x = make_float3(1 + sign * normal.x * normal.x * a, sign * b, -sign * normal.x);
    float3 y = make_float3(b, sign + normal.y * normal.y * a, -normal.y);

    return vec.x * x + vec.y * y + vec.z * normal;
}

// samples specified phase function with random number rand in [0,1)
__host__ __device__ inline CUDA::PhaseSample samplePhaseFunction(float3 curr_dir, float3 rand, float g, float d,
                                                                 PhaseFunction phase_function) {
    float cosTheta = 0;
    float pdf = 0;
    float value = 0;
    switch (phase_function) {
        case HenyeyGreenstein:
            sampleHenyeyGreenstein(rand.x, cosTheta, value, pdf, g);
            break;
        default:
            cosTheta = 2 * rand.x - 1;
            value = Inv4Pi;
            pdf = value;
    }

    float sinTheta = clamp(sqrtf(fmaxf(1 - cosTheta * cosTheta, 0.f)), -1.f, 1.f);
    cosTheta = clamp(cosTheta, -1.f, 1.f);
    float phi = 2.f * Pi * rand.y;

    // direction in spherical coordinates
    float3 dir = make_float3(sinTheta * cosf(phi), sinTheta * sinf(phi), cosTheta);
    float3 newDir = to_world(dir, curr_dir);

    return {newDir, value, pdf};
}
