rust-lattice/src/lib.rs

use wasm_bindgen::prelude::*;
use rand::{Rng, RngCore, SeedableRng, rngs::StdRng};
use std::f32::consts::PI;

const Q: i32 = 8380417; // Modulus: 2^23 - 2^13 + 1
const N: usize = 1024;    // Polynomial degree
const ROOT_OF_UNITY: i64 = 1753; // 2n-th root of unity modulo q
const INV_N: i64 = 8372225; // n^-1 mod q

#[wasm_bindgen]
pub struct SecureRNG {
    rng: StdRng,
}

#[wasm_bindgen]
impl SecureRNG {
    #[wasm_bindgen(constructor)]
    pub fn new() -> Self {
        Self {
            rng: StdRng::from_entropy(),
        }
    }

    pub fn next_float(&mut self) -> f32 {
        self.rng.gen::<f32>()
    }

    pub fn next_uint32(&mut self) -> u32 {
        self.rng.next_u32()
    }

    pub fn gaussian(&mut self, sigma: f32) -> i32 {
        let u: f32 = self.rng.gen_range(0.0001..1.0);
        let v: f32 = self.rng.gen_range(0.0001..1.0);
        
        ((-2.0 * u.ln()).sqrt() * (2.0 * PI * v).cos() * sigma).round() as i32
    }

    pub fn cbd(&mut self, eta: i32) -> i32 {
        let mut res = 0;
        for _ in 0..eta {
            let t1 = (self.rng.next_u32() & 1) as i32;
            let t2 = (self.rng.next_u32() & 1) as i32;
            res += t1 - t2;
        }
        res
    }
}

pub struct Cell {
    pub gx: f32,
    pub gy: f32,
    pub size_mult: f32,
    pub stretch_x: f32,
    pub stretch_y: f32,
    pub jitter_x: f32,
    pub jitter_y: f32,
    pub drift_phase: f32,
    pub drift_speed: f32,
    pub alpha: f32,
    pub rotation: f32,
    pub shape_offsets: Vec<f32>,
    pub hue_shift: f32,
    pub layer: u8,
    pub edge_style: f32,
    pub eye_warp: f32,
    pub mouth_warp: f32,
    pub asymmetry: f32,
    pub phase_lag: f32,
    pub texture_shift: f32,
    pub lwe_index: usize,
    pub persistent_offset_x: f32,
    pub persistent_offset_y: f32,
    pub gaze_mismatch: f32,
    pub stutter_frequency: f32,
    pub tissue_inertia: f32,
    pub blink_lag: f32,
    pub is_neural_frozen: bool,
    pub freeze_duration: f32,
}

#[wasm_bindgen]
pub struct LatticeEngine {
    cells: Vec<Cell>,
    b_field: Vec<f32>,
    target_b_field: Vec<f32>,
    color_b_field: Vec<f32>,
    target_color_b_field: Vec<f32>,
    hidden_manifold: Vec<f32>,
    latent_secret: Vec<i32>,
    drift_field_a: Vec<i32>,
    ntt_psi: Vec<i32>,
    ntt_psi_inv: Vec<i32>,
    rng: SecureRNG,
    frame_count: usize,
}

#[wasm_bindgen]
impl LatticeEngine {
    #[wasm_bindgen(constructor)]
    pub fn new() -> Self {
        let mut rng = SecureRNG::new();
        let mut hidden_manifold = vec![0.0; N];
        let mut latent_secret = vec![0; N];
        for i in 0..N {
            hidden_manifold[i] = rng.next_float() - 0.5;
            latent_secret[i] = rng.gaussian(2.0);
        }

        let mut ntt_psi = vec![0; N];
        let mut ntt_psi_inv = vec![0; N];
        let psi = Self::mod_pow(ROOT_OF_UNITY as i64, 1);
        let psi_inv = Self::mod_pow(psi, (Q - 2) as i64);

        for i in 0..N {
            ntt_psi[i] = Self::mod_pow(psi, Self::bit_reverse(i, 10) as i64) as i32;
            ntt_psi_inv[i] = Self::mod_pow(psi_inv, Self::bit_reverse(i, 10) as i64) as i32;
        }

        Self {
            cells: Vec::new(),
            b_field: vec![0.0; N],
            target_b_field: vec![0.0; N],
            color_b_field: vec![0.0; N],
            target_color_b_field: vec![0.0; N],
            hidden_manifold,
            latent_secret,
            drift_field_a: vec![0; N],
            ntt_psi,
            ntt_psi_inv,
            rng,
            frame_count: 0,
        }
    }

    fn mod_pow(base: i64, exp: i64) -> i64 {
        let mut res = 1i64;
        let mut b = base % Q as i64;
        let mut e = exp;
        while e > 0 {
            if e % 2 == 1 {
                res = (res * b) % Q as i64;
            }
            b = (b * b) % Q as i64;
            e /= 2;
        }
        res
    }

    fn bit_reverse(mut x: usize, bits: usize) -> usize {
        let mut res = 0;
        for _ in 0..bits {
            res = (res << 1) | (x & 1);
            x >>= 1;
        }
        res
    }

    pub fn init_lattice(&mut self, dim: usize, rlwe_n: usize) {
        self.cells.clear();
        for gy in 0..dim {
            for gx in 0..dim {
                let rand_val = self.rng.next_float();
                let mut size_mult = 0.5 + self.rng.next_float() * 0.3;
                let mut layer = 0;
                let edge_style = self.rng.next_float();

                if rand_val > 0.5 && rand_val <= 0.85 {
                    size_mult = 0.9 + self.rng.next_float() * 0.4;
                    layer = 1;
                } else if rand_val > 0.85 {
                    size_mult = 1.3 + self.rng.next_float() * 0.5;
                    layer = 2;
                }

                let shape_count = 4 + (self.rng.next_float() * 3.0) as usize;
                let mut shape_offsets = Vec::with_capacity(shape_count);
                for _ in 0..shape_count {
                    shape_offsets.push(0.6 + self.rng.next_float() * 0.8);
                }

                self.cells.push(Cell {
                    gx: gx as f32,
                    gy: gy as f32,
                    size_mult,
                    stretch_x: 0.6 + self.rng.next_float() * 0.8,
                    stretch_y: 0.6 + self.rng.next_float() * 0.8,
                    jitter_x: (self.rng.next_float() - 0.5) * 1.5,
                    jitter_y: (self.rng.next_float() - 0.5) * 1.5,
                    drift_phase: self.rng.next_float() * PI * 2.0,
                    drift_speed: 0.05 + self.rng.next_float() * 0.15,
                    alpha: 0.5 + self.rng.next_float() * 0.45,
                    rotation: self.rng.next_float() * PI * 2.0,
                    shape_offsets,
                    hue_shift: (self.rng.next_float() - 0.5) * 16.0,
                    layer,
                    edge_style,
                    eye_warp: (self.rng.next_float() - 0.5) * 0.8,
                    mouth_warp: (self.rng.next_float() - 0.5) * 0.6,
                    asymmetry: (self.rng.next_float() - 0.5) * 0.4,
                    phase_lag: self.rng.next_float() * PI * 2.0,
                    texture_shift: (self.rng.next_float() - 0.5) * 8.0,
                    lwe_index: (self.rng.next_float() * rlwe_n as f32) as usize,
                    persistent_offset_x: (self.rng.next_float() - 0.5) * 4.0,
                    persistent_offset_y: (self.rng.next_float() - 0.5) * 4.0,
                    gaze_mismatch: 0.9 + self.rng.next_float() * 0.2,
                    stutter_frequency: 0.015 + self.rng.next_float() * 0.05,
                    tissue_inertia: 0.8 + self.rng.next_float() * 0.6,
                    blink_lag: self.rng.next_float() * 120.0,
                    is_neural_frozen: false,
                    freeze_duration: 0.0,
                });
            }
        }
    }

    fn run_forward_ntt(a: &mut [i32], psi: &[i32]) {
        let mut k = 1;
        for len in (1..N).rev().step_by(1) {
            // Simplified for brevity, in a real implementation we'd use the full Cooley-Tukey
            // But I'll mirror the TS implementation exactly.
        }
        // Mirroring the TS logic exactly is safer for "exact same behavior"
    }

    // Since I want to be 100% accurate, I'll copy the logic I just refined in TS.
    pub fn rotate_field(&mut self, error_width: i32) {
        let mut b_final_warp = vec![0i32; N];
        let mut b_final_color = vec![0i32; N];

        for i in 0..N {
            self.drift_field_a[i] = (self.rng.next_uint32() % Q as u32) as i32;
        }

        self.poly_multiply_ntt(&self.drift_field_a, &self.latent_secret, &mut b_final_warp);

        let mut a2 = vec![0i32; N];
        for i in 0..N {
            a2[i] = (self.rng.next_uint32() % Q as u32) as i32;
        }
        self.poly_multiply_ntt(&a2, &self.latent_secret, &mut b_final_color);

        for i in 0..N {
            let error_w = self.rng.cbd(error_width);
            let val_w = (b_final_warp[i] + error_w + Q) % Q;
            let shell_w = (val_w as f32 / Q as f32) - 0.5;
            self.target_b_field[i] = shell_w * 0.95 + self.hidden_manifold[i] * 0.05;

            let error_c = self.rng.cbd(error_width);
            let val_c = (b_final_color[i] + error_c + Q) % Q;
            let shell_c = (val_c as f32 / Q as f32) - 0.5;
            self.target_color_b_field[i] = shell_c * 0.95 + self.hidden_manifold[i] * 0.05;
        }
    }

    fn poly_multiply_ntt(&self, a: &[i32], s: &[i32], res: &mut [i32]) {
        let mut a_ntt = a.to_vec();
        let mut s_ntt = s.to_vec();
        self.forward_ntt(&mut a_ntt);
        self.forward_ntt(&mut s_ntt);
        for i in 0..N {
            res[i] = ((a_ntt[i] as i64 * s_ntt[i] as i64) % Q as i64) as i32;
        }
        self.inverse_ntt(res);
    }

    fn forward_ntt(&self, a: &mut [i32]) {
        let mut k = 1;
        let mut len = N / 2;
        while len >= 1 {
            let mut start = 0;
            while start < N {
                let zeta = self.ntt_psi[k] as i64;
                k += 1;
                for j in start..start + len {
                    let t = (zeta * a[j + len] as i64) % Q as i64;
                    a[j + len] = ((a[j] as i64 - t + Q as i64) % Q as i64) as i32;
                    a[j] = ((a[j] as i64 + t) % Q as i64) as i32;
                }
                start += 2 * len;
            }
            len /= 2;
        }
    }

    fn inverse_ntt(&self, a: &mut [i32]) {
        let mut k = N - 1;
        let mut len = 1;
        while len < N {
            let mut start = 0;
            while start < N {
                let zeta = self.ntt_psi_inv[k] as i64;
                k -= 1;
                for j in start..start + len {
                    let u = a[j] as i64;
                    let v = (a[j + len] as i64 * zeta) % Q as i64;
                    a[j] = ((u + v) % Q as i64) as i32;
                    a[j + len] = ((u - v + Q as i64) % Q as i64) as i32;
                }
                start += 2 * len;
            }
            len *= 2;
        }
        for i in 0..N {
            a[i] = ((a[i] as i64 * INV_N) % Q as i64) as i32;
        }
    }

    pub fn update(&mut self, alpha_smooth: f32) {
        for i in 0..N {
            let d_w = self.target_b_field[i] - self.b_field[i];
            let d_c = self.target_color_b_field[i] - self.color_b_field[i];
            self.b_field[i] += d_w.tanh() * alpha_smooth;
            self.color_b_field[i] += d_c.tanh() * alpha_smooth;
        }

        // Neural Stutter and Freezing logic in Rust
        for cell in &mut self.cells {
            if cell.is_neural_frozen {
                cell.freeze_duration -= 16.6;
                if cell.freeze_duration <= 0.0 {
                    cell.is_neural_frozen = false;
                }
            } else if self.rng.next_float() < 0.0005 {
                cell.is_neural_frozen = true;
                cell.freeze_duration = 100.0 + self.rng.next_float() * 300.0;
            }
        }

        self.frame_count += 1;
    }

    pub fn get_cells_raw(&self) -> Vec<f32> {
        let mut data = Vec::with_capacity(self.cells.len() * 10);
        for cell in &self.cells {
            data.push(cell.gx);
            data.push(cell.gy);
            data.push(cell.lwe_index as f32);
            data.push(if cell.is_neural_frozen { 1.0 } else { 0.0 });
            data.push(cell.stutter_frequency);
            data.push(cell.tissue_inertia);
            data.push(cell.blink_lag);
            data.push(cell.persistent_offset_x);
            data.push(cell.persistent_offset_y);
            data.push(cell.gaze_mismatch);
        }
        data
    }

    pub fn get_stats(&self) -> String {
        format!("{{ \"frame\": {}, \"active_cells\": {}, \"engine\": \"Rust/Wasm\" }}", 
            self.frame_count, self.cells.len())
    }
}