#include "Algorithm/MicrofacetGGX.h"

// Trowbridge-Reitz GGX Normal Distribution Function
float ggx_distribution(float n_dot_h, float roughness)
{
    float a = roughness * roughness;
    float a2 = a * a;
    float n_dot_h2 = n_dot_h * n_dot_h;

    float nom = a2;
    float denom = (n_dot_h2 * (a2 - 1.0f) + 1.0f);
    denom = PI * denom * denom;

    return nom / fmaxf(denom, 1e-20f);
}

float ggx_g1(float n_dot_v, float roughness)
{
    if (n_dot_v <= 0.0f)
    {
        return 0.0f;
    }

    // Our GGX D() uses alpha = roughness^2, so keep the same convention here.
    float alpha = roughness * roughness;
    float alpha2 = alpha * alpha;

    float n2 = n_dot_v * n_dot_v;
    float denom = n_dot_v + sqrtf(alpha2 + (1.0f - alpha2) * n2);
    return (2.0f * n_dot_v) / fmaxf(denom, FLT_EPSILON);
}

float ggx_g_smith(float n_dot_v, float n_dot_l, float roughness)
{
    return ggx_g1(n_dot_v, roughness) * ggx_g1(n_dot_l, roughness);
}

float ggx_visibility(float n_dot_v, float n_dot_l, float roughness)
{
    float alpha = roughness * roughness;
    float alpha2 = alpha * alpha;

    float ggx_v = n_dot_l * sqrtf(alpha2 + (1.0f - alpha2) * n_dot_v * n_dot_v);
    float ggx_l = n_dot_v * sqrtf(alpha2 + (1.0f - alpha2) * n_dot_l * n_dot_l);

    return 0.5f / fmaxf(ggx_v + ggx_l, FLT_EPSILON);
}

vec3s ggx_sample_vndf(vec3s n, vec3s v, float roughness, float u1, float u2)
{
    // Build local frame around n.
    vec3s tangent, bitangent;
    create_orthonormal_basis(n, &tangent, &bitangent);

    // View direction in local coordinates.
    vec3s v_local = (vec3s){glms_vec3_dot(v, tangent), glms_vec3_dot(v, bitangent), glms_vec3_dot(v, n)};
    v_local = glms_vec3_normalize(v_local);

    // Stretch view.
    float alpha = roughness * roughness;
    vec3s v_h = (vec3s){alpha * v_local.x, alpha * v_local.y, v_local.z};
    v_h = glms_vec3_normalize(v_h);

    // Orthonormal basis around v_h.
    vec3s t1;
    if (v_h.z < 0.9999f)
    {
        t1 = glms_vec3_normalize(glms_vec3_cross((vec3s){0.0f, 0.0f, 1.0f}, v_h));
    }
    else
    {
        t1 = (vec3s){1.0f, 0.0f, 0.0f};
    }
    vec3s t2 = glms_vec3_cross(v_h, t1);

    // Sample a point on a disk.
    float r = sqrtf(u1);
    float phi = TWO_PI * u2;
    float t1p = r * cosf(phi);
    float t2p = r * sinf(phi);

    // Warp to the hemisphere.
    float s = 0.5f * (1.0f + v_h.z);
    t2p = (1.0f - s) * sqrtf(fmaxf(0.0f, 1.0f - t1p * t1p)) + s * t2p;

    vec3s n_h = glms_vec3_add(
        glms_vec3_add(glms_vec3_scale(t1, t1p), glms_vec3_scale(t2, t2p)),
        glms_vec3_scale(v_h, sqrtf(fmaxf(0.0f, 1.0f - t1p * t1p - t2p * t2p))));

    // Unstretch.
    vec3s h_local = (vec3s){alpha * n_h.x, alpha * n_h.y, fmaxf(0.0f, n_h.z)};
    h_local = glms_vec3_normalize(h_local);

    // Back to world.
    vec3s h_world = glms_vec3_add(
        glms_vec3_add(glms_vec3_scale(tangent, h_local.x), glms_vec3_scale(bitangent, h_local.y)),
        glms_vec3_scale(n, h_local.z));

    return glms_vec3_normalize(h_world);
}