#include "Material/SimpleLit.h"
#include "Algorithm/BSDF.h"
#include <float.h>

static const float DIELECTRIC_REFLECTIVE_F0 = 0.04f; // Standard dielectric reflectivity coef at incident angle (= 4%)
static const vec3s DIELECTRIC_REFLECTIVE = {0.04f, 0.04f, 0.04f}; // Standard dielectric reflectivity coef at incident angle (= 4%)

// Simple lit, but keep it unbiased as much as possible
vec3s sample_bsdf_simple_lit(const void* data, const vec3s normal, const vec3s wo, sobol_state_t* sobol_state, float* pdf_out)
{
    const simple_lit_data_t shading_data = *(const simple_lit_data_t*)data;

    //TODO: having a bsdf data struct to avoid recomputing the same thing in both sample and evaluate
    vec3s f0 = glms_vec3_lerp(DIELECTRIC_REFLECTIVE, shading_data.albedo, shading_data.metallic);
    float cos_theta_0 = fmaxf(glms_vec3_dot(normal, wo), 0.0f);
    float F = glms_vec3_max(fresnel_schlick_vec3(f0, cos_theta_0)); // We use the max component of the Fresnel term for simplicity

    float prob_specular = glm_lerp(F, 1.0f, shading_data.metallic);
    float prob_diffuse = (1.0f - shading_data.metallic) * (1.0f - F); // Diffuse only for non-metals, reduced by reflection
    float total_prob = prob_diffuse + prob_specular;
    if (total_prob < FLT_EPSILON)
    {
        *pdf_out = 0.0f;
        return glms_vec3_zero();
    }

    // Normalize probabilities
    // total_prob should be 1.0f, worth it? Maybe still need to avoid floating point errors
    prob_diffuse /= total_prob;
    prob_specular /= total_prob;

    vec3s wi;

    if (random_float() < prob_diffuse) // Diffuse Lobe
    {
        wi = random_cosine_direction(normal, sobol_state);
    }
    else // Specular Lobe
    {
        // For simplification we use blinn-phong lobe distribution, we will implement GGX for standard lit later
        // When talking about simplification, wen even can use a simple interpolation bwtween roughness and wi, but it's too biased.
        // A common simplification involves sampling spherical coordinates(theta and phi angles) related to normal such that cose(theta) is distributed according to the Blinn-Phong distribution
        // We can use a inversion sampling where cos(theta) = powf(random_float(), 1.0f / (specular_exponent + 1.0f)) and phi = 2 * PI * random_float()
        float specular_exponent = roughness_to_blinn_phong_specular_exponent(shading_data.roughness);
        float theta = acosf(powf(sobol_next(sobol_state), 1.0f / (specular_exponent + 1.0f)));
        float phi = 2.0f * (float)M_PI * sobol_next(sobol_state);
        vec3s h_ts = (vec3s){
            sinf(theta) * cosf(phi),
            sinf(theta) * sinf(phi),
            cosf(theta)
        };

        vec3s tangent_u; // World-space tangent (U)
        vec3s bitangent_v; // World-space bitangent (V)
        create_orthonormal_basis(normal, &tangent_u, &bitangent_v);

        vec3s scaled_u = glms_vec3_scale(tangent_u, h_ts.x);
        vec3s scaled_v = glms_vec3_scale(bitangent_v, h_ts.y);
        vec3s scaled_n = glms_vec3_scale(normal, h_ts.z);

        // Transform h from tangent space to world space
        vec3s h_ws;
        h_ws = glms_vec3_add(scaled_u, scaled_v);
        h_ws = glms_vec3_add(h_ws, scaled_n);
        h_ws = glms_vec3_normalize(h_ws); // Normalize the half-vector

        // wi is simple now, just reflect wo around normal
        wi = glms_vec3_reflect(glms_vec3_negate(wo), h_ws);
    }

    // Final check to ensure wi is in the correct hemisphere
    if (glms_vec3_dot(wi, normal) < 0.0f)
    {
        *pdf_out = 0.0f;
        return glms_vec3_zero();
    }

    float pdf_diffuse = pdf_cosine_weighted_hemisphere(normal, wi);
    float pdf_specular = pdf_blinn_phong_lobe(normal, wi, wo, shading_data.roughness);
    *pdf_out = prob_diffuse * pdf_diffuse + prob_specular * pdf_specular;

    return wi;
}

//TODO: Most of the calculation here is same as in sample_bsdf_simple_lit, we can optimize this by using a bsdf data struct to avoid recomputing the same thing in both sample and evaluate
float sample_bsdf_pdf_simple_lit(const void* data, const vec3s normal, const vec3s wo, const vec3s wi)
{
    // If wi is below the horizon relative to the normal, PDF must be 0
    if (glms_vec3_dot(normal, wi) <= 0.0f) // Use <= to be safe
    {
        return 0.0f;
    }

    const simple_lit_data_t shading_data = *(const simple_lit_data_t*)data;

    // Again, we need bsdf data;
    vec3s f0 = glms_vec3_lerp(DIELECTRIC_REFLECTIVE, shading_data.albedo, shading_data.metallic);
    float cos_theta_o = fmaxf(glms_vec3_dot(normal, wo), 0.0f); // Use 'o' for outgoing (wo)
    float F = glms_vec3_max(fresnel_schlick_vec3(f0, cos_theta_o));

    float prob_specular = glm_lerp(F, 1.0f, shading_data.metallic);
    float prob_diffuse = (1.0f - shading_data.metallic) * (1.0f - F);
    float total_prob = prob_diffuse + prob_specular;

    if (total_prob < FLT_EPSILON)
    {
        return 0.0f; // No probability of scattering
    }
    prob_diffuse /= total_prob;
    prob_specular /= total_prob;

    float pdf_diff = pdf_cosine_weighted_hemisphere(normal, wi);
    float diffuse_pdf_component = prob_diffuse * pdf_diff;

    float pdf_spec = pdf_blinn_phong_lobe(normal, wo, wi, shading_data.roughness);
    float specular_pdf_component = prob_specular * pdf_spec;

    return diffuse_pdf_component + specular_pdf_component;
}

vec3s evaluate_bsdf_simple_lit(const shading_context_t* context, const void* data)
{
    const simple_lit_data_t shading_data = *(const simple_lit_data_t*)data;
    const shading_context_t shading_context = *context;

    vec3s h = glms_vec3_normalize(glms_vec3_add(shading_context.wi, shading_context.wo));
    float n_dot_l = fmaxf(FLT_EPSILON, glms_vec3_dot(shading_context.normal, shading_context.wi));
    float n_dot_v = fmaxf(FLT_EPSILON, glms_vec3_dot(shading_context.normal, shading_context.wo));
    float n_dot_h = glms_vec3_dot(shading_context.normal, h);
    float v_dot_h = glms_vec3_dot(shading_context.wo, h);

    vec3s f0 = glms_vec3_lerp(DIELECTRIC_REFLECTIVE, shading_data.albedo, shading_data.metallic);

    vec3s diffuse_color = glms_vec3_scale(shading_data.albedo, 1.0f - shading_data.metallic);
    vec3s diffuse_term = glms_vec3_scale(diffuse_color, (float)M_1_PI);

    float specular_exponent = roughness_to_blinn_phong_specular_exponent(shading_data.roughness);
    // Normalization factor D (Blinn-Phong distribution)
    float D_norm = (specular_exponent + 2.0f) / (2.0f * (float)M_PI); // Common normalization
    float D = D_norm * powf(n_dot_h, specular_exponent);
    vec3s F = fresnel_schlick_vec3(f0, v_dot_h);

    float denominator = 4.0f * n_dot_l * n_dot_v;
    if (denominator < FLT_EPSILON)
    {
        return diffuse_term;
    }

    vec3s specular_term = glms_vec3_scale(F, D / denominator); // Specular term (Blinn-Phong), we assume that G = 1.0f for simplicity

    return glms_vec3_add(diffuse_term, specular_term);
}