amari_gpu/

shaders.rs

//! WebGPU compute shader library for mathematical operations
//!
//! This module contains optimized WGSL compute shaders for various mathematical
//! domains including tropical algebra, automatic differentiation, and fusion systems.

use std::collections::HashMap;

/// Shader library managing all mathematical compute shaders
pub struct ShaderLibrary {
    shaders: HashMap<String, &'static str>,
}

impl ShaderLibrary {
    /// Create new shader library with all mathematical shaders
    pub fn new() -> Self {
        let mut shaders = HashMap::new();

        // Tropical algebra shaders
        shaders.insert(
            "tropical_matrix_multiply".to_string(),
            TROPICAL_MATRIX_MULTIPLY,
        );
        shaders.insert("tropical_vector_add".to_string(), TROPICAL_VECTOR_ADD);
        shaders.insert(
            "tropical_neural_network".to_string(),
            TROPICAL_NEURAL_NETWORK,
        );

        // Dual number shaders
        shaders.insert("dual_forward_ad".to_string(), DUAL_FORWARD_AD);
        shaders.insert("dual_batch_gradient".to_string(), DUAL_BATCH_GRADIENT);
        shaders.insert("dual_chain_rule".to_string(), DUAL_CHAIN_RULE);

        // Fusion system shaders
        shaders.insert("tropical_dual_clifford".to_string(), TROPICAL_DUAL_CLIFFORD);
        shaders.insert("fusion_attention".to_string(), FUSION_ATTENTION);

        // Information geometry shaders
        shaders.insert("fisher_information".to_string(), FISHER_INFORMATION);
        shaders.insert("kl_divergence_batch".to_string(), KL_DIVERGENCE_BATCH);

        // Cellular automata shaders
        shaders.insert("ca_evolution".to_string(), CA_EVOLUTION);
        shaders.insert("ca_self_assembly".to_string(), CA_SELF_ASSEMBLY);
        shaders.insert("rule_application".to_string(), RULE_APPLICATION);
        shaders.insert("energy_calculation".to_string(), ENERGY_CALCULATION);
        shaders.insert("neighbor_extraction".to_string(), NEIGHBOR_EXTRACTION);

        // Enumerative geometry shaders
        shaders.insert("intersection_theory".to_string(), INTERSECTION_THEORY);
        shaders.insert("schubert_calculus".to_string(), SCHUBERT_CALCULUS);

        Self { shaders }
    }

    /// Get shader source by name
    pub fn get_shader(&self, name: &str) -> Option<&'static str> {
        self.shaders.get(name).copied()
    }

    /// List all available shaders
    pub fn list_shaders(&self) -> Vec<String> {
        self.shaders.keys().cloned().collect()
    }
}

impl Default for ShaderLibrary {
    fn default() -> Self {
        Self::new()
    }
}

/// Tropical algebra shader collection
pub const TROPICAL_SHADERS: &[(&str, &str)] = &[
    ("tropical_matrix_multiply", TROPICAL_MATRIX_MULTIPLY),
    ("tropical_vector_add", TROPICAL_VECTOR_ADD),
    ("tropical_neural_network", TROPICAL_NEURAL_NETWORK),
];

/// Dual number shader collection
pub const DUAL_SHADERS: &[(&str, &str)] = &[
    ("dual_forward_ad", DUAL_FORWARD_AD),
    ("dual_batch_gradient", DUAL_BATCH_GRADIENT),
    ("dual_chain_rule", DUAL_CHAIN_RULE),
];

/// Fusion system shader collection
pub const FUSION_SHADERS: &[(&str, &str)] = &[
    ("tropical_dual_clifford", TROPICAL_DUAL_CLIFFORD),
    ("fusion_attention", FUSION_ATTENTION),
];

// =====================================================================
// TROPICAL ALGEBRA SHADERS
// =====================================================================

/// Tropical (max-plus) matrix multiplication: C = A ⊗ B where ⊗ is tropical product
const TROPICAL_MATRIX_MULTIPLY: &str = r#"
@group(0) @binding(0) var<storage, read> matrix_a: array<f32>;
@group(0) @binding(1) var<storage, read> matrix_b: array<f32>;
@group(0) @binding(2) var<storage, read_write> result: array<f32>;
@group(0) @binding(3) var<storage, read> dimensions: array<u32>; // [M, N, K]

@compute @workgroup_size(16, 16)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let M = dimensions[0];
    let N = dimensions[1];
    let K = dimensions[2];

    let row = global_id.x;
    let col = global_id.y;

    if (row >= M || col >= N) {
        return;
    }

    // Tropical matrix multiplication: (A ⊗ B)[i,j] = max_k(A[i,k] + B[k,j])
    var max_val = -3.4028235e+38; // -infinity in tropical algebra

    for (var k = 0u; k < K; k = k + 1u) {
        let a_val = matrix_a[row * K + k];
        let b_val = matrix_b[k * N + col];

        // Tropical multiplication: a ⊗ b = a + b
        let tropical_product = a_val + b_val;

        // Tropical addition: max operation
        if (tropical_product > max_val) {
            max_val = tropical_product;
        }
    }

    result[row * N + col] = max_val;
}
"#;

/// Tropical vector addition: c = a ⊕ b where ⊕ is max operation
const TROPICAL_VECTOR_ADD: &str = r#"
@group(0) @binding(0) var<storage, read> vector_a: array<f32>;
@group(0) @binding(1) var<storage, read> vector_b: array<f32>;
@group(0) @binding(2) var<storage, read_write> result: array<f32>;

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let idx = global_id.x;

    if (idx >= arrayLength(&vector_a)) {
        return;
    }

    // Tropical addition: a ⊕ b = max(a, b)
    result[idx] = max(vector_a[idx], vector_b[idx]);
}
"#;

/// Tropical neural network layer computation
const TROPICAL_NEURAL_NETWORK: &str = r#"
@group(0) @binding(0) var<storage, read> input: array<f32>;
@group(0) @binding(1) var<storage, read> weights: array<f32>;
@group(0) @binding(2) var<storage, read> bias: array<f32>;
@group(0) @binding(3) var<storage, read_write> output: array<f32>;
@group(0) @binding(4) var<storage, read> dimensions: array<u32>; // [batch_size, input_size, output_size]

@compute @workgroup_size(16, 16)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let batch_idx = global_id.x;
    let output_idx = global_id.y;

    let batch_size = dimensions[0];
    let input_size = dimensions[1];
    let output_size = dimensions[2];

    if (batch_idx >= batch_size || output_idx >= output_size) {
        return;
    }

    // Tropical neural network: max-plus linear transformation
    var max_val = -3.4028235e+38; // -infinity

    for (var i = 0u; i < input_size; i = i + 1u) {
        let input_val = input[batch_idx * input_size + i];
        let weight_val = weights[i * output_size + output_idx];

        // Tropical multiplication: input ⊗ weight = input + weight
        let product = input_val + weight_val;

        // Tropical addition: max operation
        if (product > max_val) {
            max_val = product;
        }
    }

    // Add bias (tropical addition = max)
    let bias_val = bias[output_idx];
    let final_result = max(max_val, bias_val);

    output[batch_idx * output_size + output_idx] = final_result;
}
"#;

// =====================================================================
// DUAL NUMBER SHADERS (AUTOMATIC DIFFERENTIATION)
// =====================================================================

/// Forward-mode automatic differentiation for dual numbers
const DUAL_FORWARD_AD: &str = r#"
struct DualNumber {
    real: f32,
    dual: f32, // derivative part
}

@group(0) @binding(0) var<storage, read> input_dual: array<DualNumber>;
@group(0) @binding(1) var<storage, read> operation_params: array<f32>;
@group(0) @binding(2) var<storage, read_write> output_dual: array<DualNumber>;

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let idx = global_id.x;

    if (idx >= arrayLength(&input_dual)) {
        return;
    }

    let x = input_dual[idx];
    let op_type = u32(operation_params[0]); // Operation type

    var result: DualNumber;

    // Forward-mode AD for different operations
    switch (op_type) {
        case 0u: { // sin(x): (sin(x), cos(x) * dx)
            result.real = sin(x.real);
            result.dual = cos(x.real) * x.dual;
        }
        case 1u: { // exp(x): (exp(x), exp(x) * dx)
            let exp_val = exp(x.real);
            result.real = exp_val;
            result.dual = exp_val * x.dual;
        }
        case 2u: { // x^2: (x^2, 2x * dx)
            result.real = x.real * x.real;
            result.dual = 2.0 * x.real * x.dual;
        }
        case 3u: { // log(x): (log(x), (1/x) * dx)
            result.real = log(x.real);
            result.dual = x.dual / x.real;
        }
        default: { // identity
            result = x;
        }
    }

    output_dual[idx] = result;
}
"#;

/// Batch gradient computation for multiple functions
const DUAL_BATCH_GRADIENT: &str = r#"
struct DualNumber {
    real: f32,
    dual: f32,
}

@group(0) @binding(0) var<storage, read> input_batch: array<DualNumber>;
@group(0) @binding(1) var<storage, read> function_params: array<f32>;
@group(0) @binding(2) var<storage, read_write> gradients: array<f32>;
@group(0) @binding(3) var<storage, read> batch_info: array<u32>; // [batch_size, function_dim]

@compute @workgroup_size(16, 16)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let batch_idx = global_id.x;
    let var_idx = global_id.y;

    let batch_size = batch_info[0];
    let function_dim = batch_info[1];

    if (batch_idx >= batch_size || var_idx >= function_dim) {
        return;
    }

    let input_idx = batch_idx * function_dim + var_idx;
    let x = input_batch[input_idx];

    // Compute gradient of composite function f(g(x)) where g is parameterized
    let param_idx = var_idx % 4u; // Assume up to 4 parameters per function
    let param = function_params[param_idx];

    // Example: f(x) = param * x^2 + sin(x), gradient = 2 * param * x + cos(x)
    let gradient = 2.0 * param * x.real + cos(x.real);

    gradients[input_idx] = gradient * x.dual;
}
"#;

/// Chain rule implementation for complex function compositions
const DUAL_CHAIN_RULE: &str = r#"
struct DualNumber {
    real: f32,
    dual: f32,
}

@group(0) @binding(0) var<storage, read> inner_function: array<DualNumber>;
@group(0) @binding(1) var<storage, read> outer_params: array<f32>;
@group(0) @binding(2) var<storage, read_write> composed_result: array<DualNumber>;

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let idx = global_id.x;

    if (idx >= arrayLength(&inner_function)) {
        return;
    }

    let u = inner_function[idx]; // u = g(x), du/dx
    let outer_type = u32(outer_params[0]);

    var result: DualNumber;

    // Chain rule: d/dx[f(g(x))] = f'(g(x)) * g'(x) = f'(u) * du/dx
    switch (outer_type) {
        case 0u: { // f(u) = sin(u)
            result.real = sin(u.real);
            result.dual = cos(u.real) * u.dual; // cos(u) * du/dx
        }
        case 1u: { // f(u) = u^3
            result.real = u.real * u.real * u.real;
            result.dual = 3.0 * u.real * u.real * u.dual; // 3u^2 * du/dx
        }
        case 2u: { // f(u) = exp(u)
            let exp_u = exp(u.real);
            result.real = exp_u;
            result.dual = exp_u * u.dual; // exp(u) * du/dx
        }
        default: { // f(u) = u (identity)
            result = u;
        }
    }

    composed_result[idx] = result;
}
"#;

// =====================================================================
// FUSION SYSTEM SHADERS
// =====================================================================

/// TropicalDualClifford operations for LLM evaluation
const TROPICAL_DUAL_CLIFFORD: &str = r#"
struct TropicalNumber {
    value: f32, // Tropical number value
}

struct DualNumber {
    real: f32,
    dual: f32,
}

struct Multivector {
    coeffs: array<f32, 8>, // 3D Clifford algebra: 8 basis elements
}

struct TropicalDualClifford {
    tropical: TropicalNumber,
    dual: DualNumber,
    clifford: Multivector,
}

@group(0) @binding(0) var<storage, read> input_batch: array<TropicalDualClifford>;
@group(0) @binding(1) var<storage, read> operation_params: array<f32>;
@group(0) @binding(2) var<storage, read_write> output_batch: array<TropicalDualClifford>;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let idx = global_id.x;

    if (idx >= arrayLength(&input_batch)) {
        return;
    }

    let tdc = input_batch[idx];
    let op_type = u32(operation_params[0]);

    var result: TropicalDualClifford;

    switch (op_type) {
        case 0u: { // LLM attention computation
            // Combine tropical path selection with dual gradients and geometric transformations
            result.tropical.value = max(tdc.tropical.value, operation_params[1]);
            result.dual.real = tdc.dual.real * operation_params[2];
            result.dual.dual = tdc.dual.dual * operation_params[2];

            // Geometric rotation in Clifford algebra
            let angle = operation_params[3];
            let cos_half = cos(angle * 0.5);
            let sin_half = sin(angle * 0.5);

            // Simple rotation around e12 plane
            result.clifford.coeffs[0] = cos_half * tdc.clifford.coeffs[0]; // scalar
            result.clifford.coeffs[1] = tdc.clifford.coeffs[1]; // e1
            result.clifford.coeffs[2] = tdc.clifford.coeffs[2]; // e2
            result.clifford.coeffs[3] = tdc.clifford.coeffs[3]; // e3
            result.clifford.coeffs[4] = sin_half * tdc.clifford.coeffs[0]; // e12
            result.clifford.coeffs[5] = tdc.clifford.coeffs[5]; // e13
            result.clifford.coeffs[6] = tdc.clifford.coeffs[6]; // e23
            result.clifford.coeffs[7] = tdc.clifford.coeffs[7]; // e123
        }
        default: {
            result = tdc;
        }
    }

    output_batch[idx] = result;
}
"#;

/// Fusion attention mechanism using tropical algebra
const FUSION_ATTENTION: &str = r#"
@group(0) @binding(0) var<storage, read> queries: array<f32>;
@group(0) @binding(1) var<storage, read> keys: array<f32>;
@group(0) @binding(2) var<storage, read> values: array<f32>;
@group(0) @binding(3) var<storage, read_write> attention_output: array<f32>;
@group(0) @binding(4) var<storage, read> dimensions: array<u32>; // [seq_len, d_model]

@compute @workgroup_size(16, 16)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let seq_pos = global_id.x;
    let feature_idx = global_id.y;

    let seq_len = dimensions[0];
    let d_model = dimensions[1];

    if (seq_pos >= seq_len || feature_idx >= d_model) {
        return;
    }

    // Tropical attention: use max-plus algebra instead of softmax
    var max_score = -3.4028235e+38; // -infinity
    var best_key_idx = 0u;

    // Find the key with maximum tropical attention score
    for (var key_idx = 0u; key_idx < seq_len; key_idx = key_idx + 1u) {
        var score = -3.4028235e+38;

        // Compute tropical dot product: sum becomes max, product becomes sum
        for (var d = 0u; d < d_model; d = d + 1u) {
            let q = queries[seq_pos * d_model + d];
            let k = keys[key_idx * d_model + d];

            // Tropical multiplication: q ⊗ k = q + k
            let tropical_product = q + k;

            // Tropical sum: max operation
            if (tropical_product > score) {
                score = tropical_product;
            }
        }

        if (score > max_score) {
            max_score = score;
            best_key_idx = key_idx;
        }
    }

    // Tropical attention: select value from best key (winner-takes-all)
    attention_output[seq_pos * d_model + feature_idx] =
        values[best_key_idx * d_model + feature_idx];
}
"#;

// =====================================================================
// INFORMATION GEOMETRY SHADERS
// =====================================================================

/// Fisher information matrix computation
const FISHER_INFORMATION: &str = r#"
@group(0) @binding(0) var<storage, read> probability_params: array<f32>;
@group(0) @binding(1) var<storage, read> data_points: array<f32>;
@group(0) @binding(2) var<storage, read_write> fisher_matrix: array<f32>;
@group(0) @binding(3) var<storage, read> dimensions: array<u32>; // [n_params, n_data]

@compute @workgroup_size(16, 16)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let param_i = global_id.x;
    let param_j = global_id.y;

    let n_params = dimensions[0];
    let n_data = dimensions[1];

    if (param_i >= n_params || param_j >= n_params) {
        return;
    }

    // Fisher Information Matrix: I[i,j] = E[∂²log p(x|θ)/∂θᵢ∂θⱼ]
    var fisher_element = 0.0;

    for (var data_idx = 0u; data_idx < n_data; data_idx = data_idx + 1u) {
        let x = data_points[data_idx];

        // Gaussian log-likelihood example: log p(x|μ,σ) = -½log(2πσ²) - (x-μ)²/(2σ²)
        let mu = probability_params[0];
        let sigma = probability_params[1];
        let sigma_sq = sigma * sigma;

        var d2_log_p = 0.0;

        if (param_i == 0u && param_j == 0u) { // ∂²/∂μ²
            d2_log_p = -1.0 / sigma_sq;
        } else if (param_i == 1u && param_j == 1u) { // ∂²/∂σ²
            let diff = x - mu;
            d2_log_p = -1.0 / sigma_sq + 3.0 * diff * diff / (sigma_sq * sigma_sq);
        } else if ((param_i == 0u && param_j == 1u) || (param_i == 1u && param_j == 0u)) { // ∂²/∂μ∂σ
            let diff = x - mu;
            d2_log_p = 2.0 * diff / (sigma_sq * sigma);
        }

        fisher_element += -d2_log_p; // Fisher = -E[Hessian of log-likelihood]
    }

    fisher_matrix[param_i * n_params + param_j] = fisher_element / f32(n_data);
}
"#;

/// Batch KL divergence computation
const KL_DIVERGENCE_BATCH: &str = r#"
@group(0) @binding(0) var<storage, read> distribution_p: array<f32>;
@group(0) @binding(1) var<storage, read> distribution_q: array<f32>;
@group(0) @binding(2) var<storage, read_write> kl_divergences: array<f32>;
@group(0) @binding(3) var<storage, read> batch_info: array<u32>; // [batch_size, dist_size]

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let batch_idx = global_id.x;
    let batch_size = batch_info[0];
    let dist_size = batch_info[1];

    if (batch_idx >= batch_size) {
        return;
    }

    // KL divergence: D_KL(P||Q) = Σ P(x) log(P(x)/Q(x))
    var kl_div = 0.0;

    for (var i = 0u; i < dist_size; i = i + 1u) {
        let p_i = distribution_p[batch_idx * dist_size + i];
        let q_i = distribution_q[batch_idx * dist_size + i];

        if (p_i > 1e-10 && q_i > 1e-10) { // Avoid log(0)
            kl_div += p_i * log(p_i / q_i);
        }
    }

    kl_divergences[batch_idx] = kl_div;
}
"#;

// =====================================================================
// CELLULAR AUTOMATA SHADERS
// =====================================================================

/// Cellular automata evolution step
pub const CA_EVOLUTION: &str = r#"
@group(0) @binding(0) var<storage, read> current_state: array<u32>;
@group(0) @binding(1) var<storage, read_write> next_state: array<u32>;
@group(0) @binding(2) var<storage, read> rules: array<u32>;
@group(0) @binding(3) var<storage, read> dimensions: array<u32>; // [width, height]

@compute @workgroup_size(16, 16)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let x = global_id.x;
    let y = global_id.y;

    let width = dimensions[0];
    let height = dimensions[1];

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

    let idx = y * width + x;
    let current_cell = current_state[idx];

    // Count alive neighbors (Moore neighborhood)
    var alive_neighbors = 0u;

    for (var dy = 0u; dy < 3u; dy = dy + 1u) {
        for (var dx = 0u; dx < 3u; dx = dx + 1u) {
            if (dx == 1u && dy == 1u) { continue; } // Skip center cell

            let nx = (x + dx + width - 1u) % width; // Wrap around
            let ny = (y + dy + height - 1u) % height;
            let neighbor_idx = ny * width + nx;

            if (current_state[neighbor_idx] == 1u) {
                alive_neighbors = alive_neighbors + 1u;
            }
        }
    }

    // Conway's Game of Life rules (can be customized via rules buffer)
    var new_state = 0u;

    if (current_cell == 1u) { // Currently alive
        if (alive_neighbors == 2u || alive_neighbors == 3u) {
            new_state = 1u; // Survive
        }
    } else { // Currently dead
        if (alive_neighbors == 3u) {
            new_state = 1u; // Birth
        }
    }

    next_state[idx] = new_state;
}
"#;

/// Rule application for geometric algebra cellular automata
pub const RULE_APPLICATION: &str = r#"
struct GpuCellData {
    scalar: f32,
    e1: f32,
    e2: f32,
    e3: f32,
    e12: f32,
    e13: f32,
    e23: f32,
    e123: f32,
    generation: f32,
    neighborhood_size: f32,
    rule_type: f32,
    boundary_condition: f32,
    padding: array<f32, 4>,
}

struct GpuRuleConfig {
    rule_type: f32,
    threshold: f32,
    damping_factor: f32,
    energy_conservation: f32,
    time_step: f32,
    spatial_scale: f32,
    geometric_weight: f32,
    nonlinear_factor: f32,
    boundary_type: f32,
    neighborhood_radius: f32,
    evolution_speed: f32,
    stability_factor: f32,
    padding: array<f32, 4>,
}

@group(0) @binding(0) var<storage, read> cells: array<GpuCellData>;
@group(0) @binding(1) var<storage, read> rules: array<GpuRuleConfig>;
@group(0) @binding(2) var<storage, read_write> output: array<GpuCellData>;

@compute @workgroup_size(256)
fn rule_application_main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let idx = global_id.x;

    if (idx >= arrayLength(&cells)) {
        return;
    }

    let cell = cells[idx];
    let rule = rules[0]; // Use first rule for now

    var new_cell = cell;

    // Apply damping factor
    new_cell.scalar = cell.scalar * (1.0 - rule.damping_factor);
    new_cell.e1 = cell.e1 * (1.0 - rule.damping_factor);
    new_cell.e2 = cell.e2 * (1.0 - rule.damping_factor);
    new_cell.e3 = cell.e3 * (1.0 - rule.damping_factor);
    new_cell.e12 = cell.e12 * (1.0 - rule.damping_factor);
    new_cell.e13 = cell.e13 * (1.0 - rule.damping_factor);
    new_cell.e23 = cell.e23 * (1.0 - rule.damping_factor);
    new_cell.e123 = cell.e123 * (1.0 - rule.damping_factor);

    // Apply threshold
    if (abs(new_cell.scalar) < rule.threshold) {
        new_cell.scalar = 0.0;
    }

    output[idx] = new_cell;
}
"#;

/// Energy calculation for cellular automata
pub const ENERGY_CALCULATION: &str = r#"
struct GpuCellData {
    scalar: f32,
    e1: f32,
    e2: f32,
    e3: f32,
    e12: f32,
    e13: f32,
    e23: f32,
    e123: f32,
    generation: f32,
    neighborhood_size: f32,
    rule_type: f32,
    boundary_condition: f32,
    padding: array<f32, 4>,
}

@group(0) @binding(0) var<storage, read> cells: array<GpuCellData>;
@group(0) @binding(1) var<storage, read_write> total_energy: array<f32>;

@compute @workgroup_size(1)
fn energy_calculation_main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    var energy = 0.0;

    // Sum the squared magnitudes of all multivector components
    for (var i = 0u; i < arrayLength(&cells); i = i + 1u) {
        let cell = cells[i];
        energy += cell.scalar * cell.scalar;
        energy += cell.e1 * cell.e1;
        energy += cell.e2 * cell.e2;
        energy += cell.e3 * cell.e3;
        energy += cell.e12 * cell.e12;
        energy += cell.e13 * cell.e13;
        energy += cell.e23 * cell.e23;
        energy += cell.e123 * cell.e123;
    }

    total_energy[0] = energy;
}
"#;

/// Neighbor extraction for cellular automata
pub const NEIGHBOR_EXTRACTION: &str = r#"
struct GpuCellData {
    scalar: f32,
    e1: f32,
    e2: f32,
    e3: f32,
    e12: f32,
    e13: f32,
    e23: f32,
    e123: f32,
    generation: f32,
    neighborhood_size: f32,
    rule_type: f32,
    boundary_condition: f32,
    padding: array<f32, 4>,
}

@group(0) @binding(0) var<storage, read> cells: array<GpuCellData>;
@group(0) @binding(1) var<uniform> params: array<f32, 4>; // [width, height, total_cells, padding]
@group(0) @binding(2) var<storage, read_write> neighborhoods: array<GpuCellData>;

@compute @workgroup_size(256)
fn neighbor_extraction_main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let idx = global_id.x;
    let width = u32(params[0]);
    let height = u32(params[1]);
    let total_cells = u32(params[2]);

    if (idx >= total_cells) {
        return;
    }

    // Calculate 2D position from linear index
    let x = idx % width;
    let y = idx / width;

    // Moore neighborhood: 8 neighbors
    let offsets = array<vec2<i32>, 8>(
        vec2<i32>(-1, -1), vec2<i32>(0, -1), vec2<i32>(1, -1),
        vec2<i32>(-1,  0),                   vec2<i32>(1,  0),
        vec2<i32>(-1,  1), vec2<i32>(0,  1), vec2<i32>(1,  1)
    );

    // Extract neighbors with wrapping boundaries
    for (var i = 0u; i < 8u; i = i + 1u) {
        let offset = offsets[i];
        let nx = (i32(x) + offset.x + i32(width)) % i32(width);
        let ny = (i32(y) + offset.y + i32(height)) % i32(height);
        let neighbor_idx = u32(ny) * width + u32(nx);

        // Store neighbor in output array
        neighborhoods[idx * 8u + i] = cells[neighbor_idx];
    }
}
"#;

/// Self-assembly pattern formation
const CA_SELF_ASSEMBLY: &str = r#"
@group(0) @binding(0) var<storage, read> particles: array<f32>; // [x, y, type, energy]
@group(0) @binding(1) var<storage, read_write> new_particles: array<f32>;
@group(0) @binding(2) var<storage, read> assembly_rules: array<f32>;
@group(0) @binding(3) var<storage, read> simulation_params: array<u32>; // [n_particles, grid_size]

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let particle_idx = global_id.x;
    let n_particles = simulation_params[0];
    let grid_size = simulation_params[1];

    if (particle_idx >= n_particles) {
        return;
    }

    let base_idx = particle_idx * 4u;
    let x = particles[base_idx];
    let y = particles[base_idx + 1u];
    let particle_type = particles[base_idx + 2u];
    let energy = particles[base_idx + 3u];

    // Self-assembly based on local interactions
    var new_x = x;
    var new_y = y;
    var new_energy = energy;

    // Calculate forces from nearby particles
    var force_x = 0.0;
    var force_y = 0.0;

    for (var other_idx = 0u; other_idx < n_particles; other_idx = other_idx + 1u) {
        if (other_idx == particle_idx) { continue; }

        let other_base = other_idx * 4u;
        let other_x = particles[other_base];
        let other_y = particles[other_base + 1u];
        let other_type = particles[other_base + 2u];

        let dx = other_x - x;
        let dy = other_y - y;
        let distance = sqrt(dx * dx + dy * dy);

        if (distance < 5.0 && distance > 0.1) { // Interaction range
            let interaction_strength = assembly_rules[u32(particle_type) * 4u + u32(other_type)];

            // Attractive/repulsive force based on particle types
            let force_magnitude = interaction_strength / (distance * distance);
            force_x += force_magnitude * dx / distance;
            force_y += force_magnitude * dy / distance;
        }
    }

    // Update position based on forces
    new_x += force_x * 0.1; // time step
    new_y += force_y * 0.1;

    // Keep within bounds
    new_x = clamp(new_x, 0.0, f32(grid_size));
    new_y = clamp(new_y, 0.0, f32(grid_size));

    // Energy dissipation
    new_energy = energy * 0.99;

    new_particles[base_idx] = new_x;
    new_particles[base_idx + 1u] = new_y;
    new_particles[base_idx + 2u] = particle_type;
    new_particles[base_idx + 3u] = new_energy;
}
"#;

// =====================================================================
// ENUMERATIVE GEOMETRY SHADERS
// =====================================================================

/// Intersection theory computations
pub const INTERSECTION_THEORY: &str = r#"
struct RationalNumber {
    numerator: i32,
    denominator: i32,
}

@group(0) @binding(0) var<storage, read> chow_class_a: array<RationalNumber>;
@group(0) @binding(1) var<storage, read> chow_class_b: array<RationalNumber>;
@group(0) @binding(2) var<storage, read_write> intersection_result: array<RationalNumber>;
@group(0) @binding(3) var<storage, read> geometry_params: array<u32>; // [dimension, degree_a, degree_b]

fn gcd(a: u32, b: u32) -> u32 {
    if (b == 0u) { return a; }
    return gcd(b, a % b);
}

fn add_rationals(a: RationalNumber, b: RationalNumber) -> RationalNumber {
    let num = a.numerator * b.denominator + b.numerator * a.denominator;
    let den = a.denominator * b.denominator;
    let g = gcd(u32(abs(num)), u32(abs(den)));

    return RationalNumber(num / i32(g), den / i32(g));
}

fn multiply_rationals(a: RationalNumber, b: RationalNumber) -> RationalNumber {
    let num = a.numerator * b.numerator;
    let den = a.denominator * b.denominator;
    let g = gcd(u32(abs(num)), u32(abs(den)));

    return RationalNumber(num / i32(g), den / i32(g));
}

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let idx = global_id.x;

    if (idx >= arrayLength(&chow_class_a)) {
        return;
    }

    let dimension = geometry_params[0];
    let degree_a = geometry_params[1];
    let degree_b = geometry_params[2];

    // Intersection product in Chow ring: A · B
    // For simplicity, implement as pointwise multiplication for this example
    let a = chow_class_a[idx];
    let b = chow_class_b[idx];

    // Check degree compatibility (degree_a + degree_b ≤ dimension)
    if (degree_a + degree_b <= dimension) {
        intersection_result[idx] = multiply_rationals(a, b);
    } else {
        intersection_result[idx] = RationalNumber(0, 1); // Zero class
    }
}
"#;

/// Schubert calculus computations
const SCHUBERT_CALCULUS: &str = r#"
@group(0) @binding(0) var<storage, read> partition_a: array<u32>;
@group(0) @binding(1) var<storage, read> partition_b: array<u32>;
@group(0) @binding(2) var<storage, read_write> littlewood_coeff: array<u32>;
@group(0) @binding(3) var<storage, read> grassmann_params: array<u32>; // [n, k]

@compute @workgroup_size(128)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let coeff_idx = global_id.x;

    let n = grassmann_params[0]; // Dimension of ambient space
    let k = grassmann_params[1]; // Dimension of subspaces

    if (coeff_idx >= arrayLength(&littlewood_coeff)) {
        return;
    }

    // Schubert calculus: compute Littlewood-Richardson coefficients
    // This is a simplified version - full LR coefficients require more complex algorithms

    let max_parts = min(arrayLength(&partition_a), arrayLength(&partition_b));
    var coefficient = 0u;

    // Simplified intersection number computation
    for (var i = 0u; i < max_parts; i = i + 1u) {
        let part_a = partition_a[i];
        let part_b = partition_b[i];

        // Check compatibility with Grassmannian Gr(k, n)
        if (part_a <= n - k && part_b <= n - k) {
            coefficient += part_a * part_b;
        }
    }

    littlewood_coeff[coeff_idx] = coefficient;
}
"#;

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_shader_library_creation() {
        let library = ShaderLibrary::new();
        let shaders = library.list_shaders();

        // Should have shaders for all mathematical domains
        assert!(shaders.contains(&"tropical_matrix_multiply".to_string()));
        assert!(shaders.contains(&"dual_forward_ad".to_string()));
        assert!(shaders.contains(&"tropical_dual_clifford".to_string()));
        assert!(shaders.contains(&"fisher_information".to_string()));
        assert!(shaders.contains(&"ca_evolution".to_string()));
        assert!(shaders.contains(&"intersection_theory".to_string()));
    }

    #[test]
    fn test_shader_retrieval() {
        let library = ShaderLibrary::new();

        let shader = library.get_shader("tropical_matrix_multiply");
        assert!(shader.is_some());
        assert!(shader.unwrap().contains("@compute"));
        assert!(shader.unwrap().contains("tropical"));
    }

    #[test]
    fn test_shader_constants() {
        assert_eq!(TROPICAL_SHADERS.len(), 3);
        assert_eq!(DUAL_SHADERS.len(), 3);
        assert_eq!(FUSION_SHADERS.len(), 2);
    }
}