/**
 * vertex_braid.wgsl — Vertex Shader: Spherical Coordinates → Triangles → Pixels
 *
 * The compute route: QUBO spectral coeffs → spherical coords → cartesian →
 * triangle vertices → rasterize → Hachimoji-colored pixels.
 *
 * This is NOT a pixel encoder. This is a geometry solver.
 * The GPU computes the solution by walking geodesics on the Fisher sphere.
 */

// ============================================================
// CONSTANTS
// ============================================================

const PI: f32 = 3.14159265358979323846;
const TWO_PI: f32 = 6.28318530717958647692;
const N_HACHIMOJI: u32 = 8u;
const SPHERE_RADIUS: f32 = 1.0;

// ============================================================
// UNIFORMS
// ============================================================

struct SpectralCoeffs {
    // Spherical harmonic coefficients c_{l,m} for l=0,1,2, m=-l..l
    // Packed: [c_00, c_1m1, c_10, c_1p1, c_2m2, c_2m1, c_20, c_2p1, c_2p2]
    coeffs: array<f32, 9>,
};

struct QUBOParams {
    n_vars: u32,           // number of QUBO variables
    energy_scale: f32,     // scales the sphere deformation
    rotation_angle: f32,   // global rotation (chaos game step)
    _pad: u32,
};

@group(0) @binding(0) var<uniform> spectral: SpectralCoeffs;
@group(0) @binding(1) var<uniform> params: QUBOParams;

// ============================================================
// SPECTRAL → SPHERICAL COORDINATES
// ============================================================

/// Convert spectral coefficients to a point on S^7.
/// 
/// The spectral decomposition is:
///   |ψ⟩ = Σ_{l,m} c_{l,m} |l,m⟩
///
/// We map this to a point on the unit sphere by treating the
/// normalized coefficients as direction cosines in a
/// 9-dimensional embedding (truncated at l=2).
///
/// Then we project to the 7-sphere by dropping two coordinates
/// (analogous to stereographic projection).
fn spectral_to_embedding(coeffs: array<f32, 9>) -> vec3<f32> {
    // Normalize coefficients
    var norm: f32 = 0.0;
    for (var i: u32 = 0u; i < 9u; i = i + 1u) {
        norm = norm + coeffs[i] * coeffs[i];
    }
    norm = sqrt(norm);
    
    if (norm < 0.0001) {
        return vec3<f32>(0.0, 0.0, 1.0); // default: north pole
    }
    
    // Use first 3 normalized coeffs as direction (low-l dominant)
    let c0 = coeffs[0] / norm;  // l=0, monopole (average)
    let c1 = coeffs[2] / norm;  // l=1, m=0 (dipole z)
    let c2 = coeffs[5] / norm;  // l=2, m=0 (quadrupole)
    
    // Map to spherical coordinates
    // θ (polar): from monopole contribution (c0 = cos θ)
    // φ (azimuthal): from dipole phase (c1, c2 → atan2)
    let theta = acos(clamp(c0, -1.0, 1.0));
    let phi = atan2(c2, c1); // phase from quadrupole vs dipole
    
    // Convert to cartesian on unit sphere
    let r = SPHERE_RADIUS;
    let x = r * sin(theta) * cos(phi);
    let y = r * sin(theta) * sin(phi);
    let z = r * cos(theta);
    
    // Apply chaos-game rotation (energy scale deforms the sphere)
    let angle = params.rotation_angle * params.energy_scale;
    let ca = cos(angle);
    let sa = sin(angle);
    let rx = x * ca - y * sa;
    let ry = x * sa + y * ca;
    
    return vec3<f32>(rx, ry, z);
}

// ============================================================
// HACHIMOJI STATE → COLOR
// ============================================================

/// Map a point on the sphere to a Hachimoji state.
/// Uses the 8 octants of the cartesian space.
fn point_to_hachimoji(p: vec3<f32>) -> u32 {
    let sx = select(0u, 1u, p.x > 0.0);
    let sy = select(0u, 1u, p.y > 0.0);
    let sz = select(0u, 1u, p.z > 0.0);
    
    // Octant index: 0..7 maps to A,B,C,G,P,S,T,Z
    let octant = (sz << 2u) | (sy << 1u) | sx;
    return octant % 8u;
}

/// Hachimoji palette (sRGB, linear space)
fn hachimoji_color(base: u32) -> vec3<f32> {
    switch(base) {
        case 0u: { return vec3<f32>(0.05, 0.05, 0.05); }   // A = Φ
        case 1u: { return vec3<f32>(0.20, 0.10, 0.30); }   // B = Λ
        case 2u: { return vec3<f32>(0.10, 0.30, 0.50); }   // C = Ρ
        case 3u: { return vec3<f32>(0.10, 0.80, 0.30); }   // G = Σ
        case 4u: { return vec3<f32>(0.90, 0.40, 0.10); }   // P = Ω
        case 5u: { return vec3<f32>(0.60, 0.20, 0.80); }   // S = Π
        case 6u: { return vec3<f32>(0.10, 0.70, 0.70); }   // T = Κ
        case 7u: { return vec3<f32>(0.95, 0.95, 0.95); }   // Z = Ζ
        default: { return vec3<f32>(1.0, 0.0, 1.0); }      // error: magenta
    }
}

// ============================================================
// TRIANGLE MESH: Icosphere subdivision
// ============================================================

/// Generate an icosahedron vertex.
/// idx: 0..11 (the 12 vertices of an icosahedron)
fn icosahedron_vertex(idx: u32) -> vec3<f32> {
    let phi = (1.0 + sqrt(5.0)) / 2.0; // golden ratio
    let norm = sqrt(1.0 + phi * phi);
    
    let a = 1.0 / norm;
    let b = phi / norm;
    
    switch(idx % 12u) {
        case 0u:  { return vec3<f32>(-a,  b, 0.0); }
        case 1u:  { return vec3<f32>( a,  b, 0.0); }
        case 2u:  { return vec3<f32>(-a, -b, 0.0); }
        case 3u:  { return vec3<f32>( a, -b, 0.0); }
        case 4u:  { return vec3<f32>(0.0, -a,  b); }
        case 5u:  { return vec3<f32>(0.0,  a,  b); }
        case 6u:  { return vec3<f32>(0.0, -a, -b); }
        case 7u:  { return vec3<f32>(0.0,  a, -b); }
        case 8u:  { return vec3<f32>( b, 0.0, -a); }
        case 9u:  { return vec3<f32>( b, 0.0,  a); }
        case 10u: { return vec3<f32>(-b, 0.0, -a); }
        case 11u: { return vec3<f32>(-b, 0.0,  a); }
        default:  { return vec3<f32>(0.0, 0.0, 1.0); }
    }
}

/// Subdivide an edge midpoint and re-normalize to sphere.
fn midpoint_normalize(a: vec3<f32>, b: vec3<f32>) -> vec3<f32> {
    let m = (a + b) * 0.5;
    return normalize(m) * SPHERE_RADIUS;
}

// ============================================================
// VERTEX SHADER ENTRY POINT
// ============================================================

struct VertexOutput {
    @builtin(position) position: vec4<f32>,
    @location(0) color: vec3<f32>,
    @location(1) hachimoji: u32,
    @location(2) world_pos: vec3<f32>,
};

@vertex
fn vs_main(
    @builtin(vertex_index) vertex_idx: u32,
    @builtin(instance_index) instance_idx: u32,
) -> VertexOutput {
    // Each instance = one QUBO variable
    // Each instance renders one triangle of the icosphere
    
    // Get the base triangle vertices (icosahedron)
    let v0_idx = (instance_idx * 3u + 0u) % 12u;
    let v1_idx = (instance_idx * 3u + 1u) % 12u;
    let v2_idx = (instance_idx * 3u + 2u) % 12u;
    
    // Select which vertex of the triangle this is
    var local_pos: vec3<f32>;
    switch(vertex_idx % 3u) {
        case 0u: { local_pos = icosahedron_vertex(v0_idx); }
        case 1u: { local_pos = icosahedron_vertex(v1_idx); }
        case 2u: { local_pos = icosahedron_vertex(v2_idx); }
        default: { local_pos = vec3<f32>(0.0, 0.0, 1.0); }
    }
    
    // Deform the sphere by spectral coefficients
    // Each instance gets a different rotation based on variable index
    let instance_angle = f32(instance_idx) * TWO_PI / f32(params.n_vars);
    let rotated_coeffs = spectral.coeffs;
    rotated_coeffs[2] = spectral.coeffs[2] * cos(instance_angle); // dipole x
    // rotated_coeffs[1] = spectral.coeffs[1] * sin(instance_angle); // dipole y (need 9th coeff)
    
    // Map spectral → spherical → cartesian
    let sphere_pos = spectral_to_embedding(rotated_coeffs);
    
    // Combine: local triangle geometry + global spectral deformation
    // The triangle follows the spectral flow on the sphere
    let final_pos = normalize(local_pos + sphere_pos * 0.3) * SPHERE_RADIUS;
    
    // Project to screen (simple perspective)
    // View from +z looking at origin
    let fov = 60.0 * PI / 180.0;
    let aspect = 1.0;
    let near = 0.1;
    let far = 10.0;
    
    let f = 1.0 / tan(fov / 2.0);
    let x_proj = final_pos.x * f / aspect;
    let y_proj = final_pos.y * f;
    let z_proj = (far + near) / (near - far) + (2.0 * far * near) / (near - far) / final_pos.z;
    let w = -final_pos.z;
    
    var out: VertexOutput;
    out.position = vec4<f32>(x_proj, y_proj, z_proj, w);
    
    // Hachimoji state from octant
    out.hachimoji = point_to_hachimoji(final_pos);
    out.color = hachimoji_color(out.hachimoji);
    out.world_pos = final_pos;
    
    return out;
}

// ============================================================
// FRAGMENT SHADER
// ============================================================

@fragment
fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
    // Base color from Hachimoji state
    var color = in.color;
    
    // Add subtle lighting (Lambertian)
    let light_dir = normalize(vec3<f32>(1.0, 1.0, 2.0));
    let normal = normalize(in.world_pos);
    let lambert = max(dot(normal, light_dir), 0.0);
    color = color * (0.3 + 0.7 * lambert);
    
    // Add energy-dependent glow (brighter = lower energy)
    let energy_glow = 1.0 - clamp(abs(params.energy_scale) * 0.1, 0.0, 0.5);
    color = color * energy_glow;
    
    // Gamma correction
    color = pow(color, vec3<f32>(1.0 / 2.2));
    
    return vec4<f32>(color, 1.0);
}

// ============================================================
// COMPUTE SHADER: Spectral coefficient update (FSDU step)
// ============================================================

/// Update spectral coefficients based on FAMM scar field.
/// This is the co-evolution step: scars modify the spectrum.
@compute @workgroup_size(256)
fn spectral_update(
    @builtin(global_invocation_id) gid: vec3<u32>,
) {
    let idx = gid.x;
    if (idx >= 9u) { return; } // only 9 coefficients
    
    // Read current coefficient
    let c = spectral.coeffs[idx];
    
    // Apply Baker-analogue damping: high-l modes decay faster
    let l = select(select(2u, 1u, idx < 4u), 0u, idx == 0u);
    let damping = exp(-0.1 * f32(l * l));
    
    // Add small noise (simulated scar pressure)
    let noise = fract(sin(f32(idx) * 43758.5453) * 43758.5453) * 0.01;
    
    // Update: damped + noise (FSDU scar accumulation)
    spectral.coeffs[idx] = c * damping + noise;
}