/**
 * dna_surface.wgsl — WebGPU Compute + Render: 8×8 Hachimoji Surface
 *
 * The piece de resistance:
 *   1. Compute shader: braid sort finds optimal QUBO solution
 *   2. Render: solution → 8×8 pixel canvas (Hachimoji color map)
 *   3. Hidden surface: rendered off-screen, captured as image
 *
 * Each pixel = one variable in the solution.
 * Color encodes the Hachimoji state:
 *   x=0 → dark (A-state, black)
 *   x=1 → bright (G-state, Hachimoji green)
 *
 * The 8×8 grid is the eigenvalue fingerprint of the QUBO solution.
 */

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

const N_BASES: u32 = 8u;
const GRID_SIZE: u32 = 8u;  // 8×8 pixel surface

// Hachimoji color palette (sRGB)
// Each base has a canonical color
const COLOR_A: vec4<f32> = vec4<f32>(0.05, 0.05, 0.05, 1.0);  // near black
const COLOR_B: vec4<f32> = vec4<f32>(0.20, 0.10, 0.30, 1.0);  // deep purple
const COLOR_C: vec4<f32> = vec4<f32>(0.10, 0.30, 0.50, 1.0);  // ocean blue
const COLOR_G: vec4<f32> = vec4<f32>(0.10, 0.80, 0.30, 1.0);  // hachimoji green
const COLOR_P: vec4<f32> = vec4<f32>(0.90, 0.40, 0.10, 1.0);  // plasma orange
const COLOR_S: vec4<f32> = vec4<f32>(0.60, 0.20, 0.80, 1.0);  // spectral violet
const COLOR_T: vec4<f32> = vec4<f32>(0.10, 0.70, 0.70, 1.0);  // teal
const COLOR_Z: vec4<f32> = vec4<f32>(0.95, 0.95, 0.95, 1.0);  // near white

// ============================================================
// BINDINGS
// ============================================================

@group(0) @binding(0) var<storage, read> solution: array<u32>;     // QUBO solution (binary vector)
@group(0) @binding(1) var<storage, read> energies: array<f32>;     // QUBO energies
@group(0) @binding(2) var<uniform> params: SurfaceParams;          // parameters
@group(0) @binding(3) var output_texture: texture_storage_2d<rgba8unorm, write>;

struct SurfaceParams {
    n_vars: u32,          // number of variables (max 64)
    optimal_energy: f32,  // energy of the solution
    grid_size: u32,       // 8
    _pad: u32,
};

// ============================================================
// HACHIMOJI STATE MAPPING
// ============================================================

/// Map a variable value to a Hachimoji base index.
/// x=0 → base A (index 0)
/// x=1 → base G (index 3) — the "high energy" base
fn value_to_base(value: u32) -> u32 {
    if (value == 0u) {
        return 0u;  // A
    }
    return 3u;  // G
}

/// Map a Hachimoji base index to a color.
fn base_to_color(base: u32) -> vec4<f32> {
    switch (base) {
        case 0u: { return COLOR_A; }
        case 1u: { return COLOR_B; }
        case 2u: { return COLOR_C; }
        case 3u: { return COLOR_G; }
        case 4u: { return COLOR_P; }
        case 5u: { return COLOR_S; }
        case 6u: { return COLOR_T; }
        case 7u: { return COLOR_Z; }
        default: { return COLOR_A; }
    }
}

/// Map a variable index to a grid position (row, col).
/// Variables are laid out in row-major order on the 8×8 grid.
fn var_to_grid(var_index: u32) -> vec2<u32> {
    let row = var_index / GRID_SIZE;
    let col = var_index % GRID_SIZE;
    return vec2<u32>(col, row);
}

// ============================================================
// KERNEL: RENDER SURFACE
// ============================================================

/// Render the QUBO solution as an 8×8 Hachimoji surface.
/// Each thread computes one pixel.
@compute @workgroup_size(8, 8)
fn render_surface(@builtin(global_invocation_id) gid: vec3<u32>) {
    let x = gid.x;
    let y = gid.y;

    if (x >= GRID_SIZE || y >= GRID_SIZE) {
        return;
    }

    let var_index = y * GRID_SIZE + x;

    if (var_index >= params.n_vars) {
        // Out of range: render as void (black)
        textureStore(output_texture, vec2<u32>(x, y), COLOR_A);
        return;
    }

    // Get variable value from solution
    let value = solution[var_index];

    // Map to Hachimoji base
    let base = value_to_base(value);

    // Map to color
    var color = base_to_color(base);

    // Modulate brightness by energy (lower energy = brighter)
    // This makes the optimal solution visually distinct
    let energy_factor = clamp(1.0 - abs(params.optimal_energy) * 0.01, 0.3, 1.0);
    color = vec4<f32>(color.rgb * energy_factor, color.a);

    // Write pixel
    textureStore(output_texture, vec2<u32>(x, y), color);
}

// ============================================================
// KERNEL: RENDER ENERGY HEATMAP
// ============================================================

/// Alternative: render as energy heatmap.
/// Each pixel's brightness = contribution to total energy.
@compute @workgroup_size(8, 8)
fn render_energy_heatmap(@builtin(global_invocation_id) gid: vec3<u32>) {
    let x = gid.x;
    let y = gid.y;

    if (x >= GRID_SIZE || y >= GRID_SIZE) {
        return;
    }

    let var_index = y * GRID_SIZE + x;

    if (var_index >= params.n_vars) {
        textureStore(output_texture, vec2<u32>(x, y), COLOR_A);
        return;
    }

    // Energy contribution: if x_i=1, color by Q_ii
    // For now, use a simple heatmap: 0=dark, 1=bright
    let value = solution[var_index];

    // Heatmap: value 0 = dark blue, value 1 = hot yellow
    var color: vec4<f32>;
    if (value == 0u) {
        color = vec4<f32>(0.05, 0.05, 0.20, 1.0);  // cold
    } else {
        color = vec4<f32>(0.95, 0.85, 0.10, 1.0);  // hot
    }

    textureStore(output_texture, vec2<u32>(x, y), color);
}