#include "ALgorithm/BSDF.h"
#include "cglm/struct/mat3.h"
#include "cglm/util.h"

float power_heuristic(float pdf_a, float pdf_b)
{
    float a2 = pdf_a * pdf_a;
    float b2 = pdf_b * pdf_b;
    return a2 / (a2 + b2);
}

float roughness_to_blinn_phong_specular_exponent(float roughness)
{
    return glm_clamp(2.0f * 1.0f / (fmaxf(roughness * roughness, FLT_EPSILON)) - 2.0f, FLT_EPSILON, 1.0f / FLT_EPSILON);
}

float blinn_phong_specular_exponent_to_roughness(float specular_exponent)
{
    return sqrtf(2.0f / (specular_exponent + 2.0f));
}

vec3s fresnel_schlick_vec3(vec3s f0, float cos_theta)
{
    float x = 1.0f - cos_theta;
    float x5 = x * x * x * x * x;
    return glms_vec3_adds(glms_vec3_scale(f0, (1.0f - x5)), x5);
}

vec3s normal_unpack(vec3s normal)
{
    vec3s unpacked_normal = glms_vec3_scale(normal, 2.0f);
    unpacked_normal = glms_vec3_sub(unpacked_normal, glms_vec3_one());

    float dot_xy = glm_clamp_zo(unpacked_normal.x * unpacked_normal.x + unpacked_normal.y * unpacked_normal.y);
    unpacked_normal.z = fmaxf(FLT_MIN, sqrtf(1.0f - dot_xy));

    return glms_vec3_normalize(unpacked_normal);
}

vec3s normal_ts_to_ws(vec3s normal, vec3s geo_normal, vec3s tangent)
{
    // 1. Sanitize inputs
    tangent = glms_vec3_normalize(tangent);
    geo_normal = glms_vec3_normalize(geo_normal);

    // 2. Gram-Schmidt with safety check
    float proj = glms_vec3_dot(geo_normal, tangent);
    vec3s t_prime_unorm = glms_vec3_sub(tangent, glms_vec3_scale(geo_normal, proj));

    // SAFETY: If tangent is parallel to normal, t_prime is zero.
    // Fallback to original tangent or arbitrary axis to avoid NaN.
    if (glms_vec3_norm2(t_prime_unorm) < FLT_EPSILON)
    {
        t_prime_unorm = tangent;
        // If tangent was also bad, pick an arbitrary axis
        if (glms_vec3_norm2(t_prime_unorm) < FLT_EPSILON)
        {
            create_orthonormal_basis(geo_normal, &t_prime_unorm, &tangent); // Recycle variable
        }
    }
    vec3s t_prime = glms_vec3_normalize(t_prime_unorm);

    // 3. Calculate Bitangent
    vec3s b_prime = glms_vec3_cross(geo_normal, t_prime);

    // 4. Apply Tangent Handedness (w)
    // NOTE: Check if tangent W component is stored/passed correctly.
    // If not sure, assuming 1.0 is safer than calculating it from cross products of unnormalized vectors.
    float w = (glms_vec3_dot(glms_vec3_cross(geo_normal, tangent), b_prime) < 0.0f) ? -1.0f : 1.0f;
    b_prime = glms_vec3_scale(b_prime, w);

    mat3s tbn = {t_prime.x, t_prime.y, t_prime.z, b_prime.x, b_prime.y, b_prime.z, geo_normal.x, geo_normal.y, geo_normal.z};

    // 5. Transform and Re-normalize
    return glms_vec3_normalize(glms_mat3_mulv(tbn, normal));
}

// BSDF sampling functions

float pdf_cosine_weighted_hemisphere(vec3s normal, vec3s wi)
{
    return fmaxf(glms_vec3_dot(wi, normal), 0.0f) / PI;
}

float pdf_blinn_phong_lobe(vec3s normal, vec3s wi, vec3s wo, float roughness)
{
    // Check if wo and wi are on the same side of the surface normal geometry
    if (glms_vec3_dot(wo, normal) <= 0.0f || glms_vec3_dot(wi, normal) <= 0.0f)
    {
        return 0.0f; // Cannot scatter from below horizon to above, or vice versa
    }

    // Calculate the half-vector h based on input wo and wi
    vec3s wo_n = glms_vec3_normalize(wo); // Ensure normalized inputs if not guaranteed
    vec3s wi_n = glms_vec3_normalize(wi);
    vec3s h = glms_vec3_add(wo_n, wi_n);
    float h_len_sq = glms_vec3_norm2(h);
    if (h_len_sq < FLT_EPSILON)
    {
        return 0.0f; // wo and wi are opposite, highly unlikely for reflection
    }
    h = glms_vec3_scale(h, 1.0f / sqrtf(h_len_sq)); // Normalize h

    // Calculate Blinn-Phong specular exponent
    float specular_exponent = roughness_to_blinn_phong_specular_exponent(roughness);

    // PDF of sampling h (Blinn-Phong distribution)
    // D(h) = (specular_exponent + 1) / (2 * PI) * pow(max(0, dot(n, h)), specular_exponent)
    float n_dot_h = fmaxf(0.0f, glms_vec3_dot(normal, h));
    float pdf_h = (specular_exponent + 1.0f) / (2.0f * PI) * powf(n_dot_h, specular_exponent);

    // Jacobian of the transformation from h to wi
    // jacobian = 1 / (4 * dot(wo, h))
    float wo_dot_h = fmaxf(FLT_EPSILON, glms_vec3_dot(wo_n, h)); // Use normalized wo, ensure > 0
    float jacobian = 1.0f / (4.0f * wo_dot_h);

    // PDF of sampling wi is pdf(h) * jacobian
    float pdf_spec = pdf_h * jacobian;

    return pdf_spec;
}

vec3s sample_cosine_weighted_hemisphere_z_angular(float angular, uint32_t index, uint16_t d1, uint16_t d2, uint32_t scramble)
{
    float r1 = sobol_sample_scrambled(index, d1, scramble);
    float r2 = sobol_sample_scrambled(index, d2, scramble);

    float phi = 2.0f * PI * r1;

    float cos_angular = cosf(angular);
    // Correctly sample cos(theta) for cosine weighting within the cone [cos_angular, 1]
    // cos_theta = sqrt(cos(angular)^2 + r2 * (1 - cos(angular)^2))
    float cos_theta = sqrtf(cos_angular * cos_angular + r2 * (1.0f - cos_angular * cos_angular));
    float sin_theta = sqrtf(1.0f - cos_theta * cos_theta);

    // Convert spherical coordinates (1, theta, phi) to Cartesian (Z-up)
    float x = sin_theta * cosf(phi);
    float y = sin_theta * sinf(phi);
    float z = cos_theta;

    vec3s local_dir = {{x, y, z}};
    return local_dir;
}

// Function to generate a direction with cosine weighting around (0, 0, 1)
// This is the local coordinate sample.
vec3s sample_cosine_weighted_hemisphere_z(uint32_t index, uint16_t d1, uint16_t d2, uint32_t scramble)
{
    float r1 = sobol_sample_scrambled(index, d1, scramble);
    float r2 = sobol_sample_scrambled(index, d2, scramble);

    float r = sqrtf(r1);
    float phi = 2.0f * PI * r2;

    float disk_x = r * cosf(phi);
    float disk_y = r * sinf(phi);

    // Map point (disk_x, disk_y) on disk to hemisphere (Z-up)
    float x = disk_x;
    float y = disk_y;
    float z = sqrtf(1.0f - disk_x * disk_x - disk_y * disk_y); // z = sqrt(1 - r*r) = sqrt(1 - r1)

    vec3s local_dir = {{x, y, z}};
    return local_dir;
}

// Function to create an orthonormal basis (coordinate system) from a single vector (normal)
// w will be aligned with normal, u and v will be perpendicular.
void create_orthonormal_basis(vec3s direction, vec3s* u, vec3s* v)
{
    vec3s a;
    if (fabsf(direction.x) > 0.9f)
    {
        a = (vec3s){0.0f, 1.0f, 0.0f}; // Use y-axis
    }
    else
    {
        a = (vec3s){1.0f, 0.0f, 0.0f}; // Use x-axis
    }

    *u = glms_vec3_normalize(glms_vec3_cross(a, direction));
    *v = glms_vec3_normalize(glms_vec3_cross(direction, *u));
}

vec3s random_cosine_direction_angular(vec3s direction, float angular, uint32_t index, uint32_t d1, uint32_t d2, uint32_t scramble)
{
    vec3s local_dir = sample_cosine_weighted_hemisphere_z_angular(angular, index, d1, d2, scramble);

    vec3s u, v;
    create_orthonormal_basis(direction, &u, &v);

    vec3s term_u = glms_vec3_scale(u, local_dir.x);
    vec3s term_v = glms_vec3_scale(v, local_dir.y);
    vec3s term_w = glms_vec3_scale(direction, local_dir.z);

    vec3s world_dir = glms_vec3_add(glms_vec3_add(term_u, term_v), term_w);

    return world_dir;
}

// Samples a direction from the hemisphere oriented along 'normal'
// with a cosine-weighted distribution.
vec3s random_cosine_direction(vec3s direction, uint32_t index, uint32_t d1, uint32_t d2, uint32_t scramble)
{
    vec3s local_dir = sample_cosine_weighted_hemisphere_z(index, d1, d2, scramble);

    vec3s u, v;
    create_orthonormal_basis(direction, &u, &v);

    vec3s term_u = glms_vec3_scale(u, local_dir.x);
    vec3s term_v = glms_vec3_scale(v, local_dir.y);
    vec3s term_w = glms_vec3_scale(direction, local_dir.z);

    vec3s world_dir = glms_vec3_add(glms_vec3_add(term_u, term_v), term_w);
    return world_dir;
}

vec3s random_uniform_cdf_direction(vec3s direction, uint32_t index, uint16_t d1, uint16_t d2, uint32_t scramble)
{
    float r1 = sobol_sample_scrambled(index, d1, scramble);
    float r2 = sobol_sample_scrambled(index, d2, scramble);

    float phi = 2.0f * PI * r1;
    float cos_theta = 1.0f - r2 * 2.0f;
    float sin_theta = sqrtf(fmaxf(0.0f, 1.0f - cos_theta * cos_theta));

    float x = sin_theta * cosf(phi);
    float y = sin_theta * sinf(phi);
    float z = cos_theta;

    vec3s u, v;
    create_orthonormal_basis(direction, &u, &v);

    vec3s term_u = glms_vec3_scale(u, x);
    vec3s term_v = glms_vec3_scale(v, y);
    vec3s term_w = glms_vec3_scale(direction, z);

    vec3s world_dir = glms_vec3_add(glms_vec3_add(term_u, term_v), term_w);
    return world_dir;
}

vec3s random_uniform_cdf_direction_angular(vec3s direction, uint32_t index, float angular, uint16_t d1, uint16_t d2, uint32_t scramble)
{
    float r1 = sobol_sample_scrambled(index, d1, scramble);
    float r2 = sobol_sample_scrambled(index, d2, scramble);

    float cos_alpha = cosf(angular);
    float cos_theta = 1.0f - r1 * (1.0f - cos_alpha);
    float sin_theta = sqrtf(fmaxf(0.0f, 1.0f - cos_theta * cos_theta));

    float phi = 2.0f * PI * r2;

    float x = sin_theta * cosf(phi);
    float y = sin_theta * sinf(phi);
    float z = cos_theta;

    vec3s u, v;
    create_orthonormal_basis(direction, &u, &v);

    vec3s term_u = glms_vec3_scale(u, x);
    vec3s term_v = glms_vec3_scale(v, y);
    vec3s term_w = glms_vec3_scale(direction, z);

    vec3s world_dir = glms_vec3_add(glms_vec3_add(term_u, term_v), term_w);
    return world_dir;
}

// Must use this function to weight any nee light contribution before accumulate.
vec3s weight_nee_light(vec3s bsdf, vec3s light, float pdf_bsdf, float pdf_sky)
{
    light = glms_vec3_mul(bsdf, light);
    float weight = power_heuristic(pdf_sky, pdf_bsdf);
    return glms_vec3_scale(light, weight);
}