//! N8: Meta-Q network — predicts collapse probability K epochs ahead.
//!
//! Small MLP (7→32→16→1 sigmoid) trained online via BCE loss on historical
//! (health_components, future_collapse) pairs. Output is logged as a leading
//! indicator — does NOT feed into LearningHealth until validated.

use std::collections::VecDeque;

pub const META_Q_INPUT_DIM: usize = 7;
pub const META_Q_H1: usize = 32;
pub const META_Q_H2: usize = 16;
pub const META_Q_LOOKAHEAD: usize = 5;

#[derive(Debug, Clone)]
pub struct MetaQInput {
    pub q_gap_ema: f32,
    pub log_grad_norm: f32,
    pub atom_util: f32,
    pub ens_disagreement: f32,
    pub regime_stability: f32,
    pub spectral_gap_norm: f32,
    pub grad_consistency: f32,
}

impl MetaQInput {
    pub fn to_vec(&self) -> [f32; META_Q_INPUT_DIM] {
        [
            self.q_gap_ema,
            self.log_grad_norm,
            self.atom_util,
            self.ens_disagreement,
            self.regime_stability,
            self.spectral_gap_norm,
            self.grad_consistency,
        ]
    }
}

/// Small MLP on host (CPU). n_params ≈ 7×32 + 32×16 + 16×1 + biases = ~800.
#[derive(Debug)]
pub struct MetaQNetwork {
    w1: Vec<f32>, b1: Vec<f32>,
    w2: Vec<f32>, b2: Vec<f32>,
    w3: Vec<f32>, b3: Vec<f32>,
    m_w1: Vec<f32>, v_w1: Vec<f32>,
    m_w2: Vec<f32>, v_w2: Vec<f32>,
    m_w3: Vec<f32>, v_w3: Vec<f32>,
    lr: f32,
    step: u32,
    buffer: VecDeque<(MetaQInput, f32)>,
    buffer_cap: usize,
}

impl MetaQNetwork {
    pub fn new() -> Self {
        // Use rand with a fixed seed for determinism of initial weights.
        // fastrand is already a dep in several places; fall back to a simple
        // deterministic LCG if fastrand is unavailable.
        let mut rng = SimpleLcg::new(0xDEADBEEF);
        let xavier = |rng: &mut SimpleLcg, fan_in: usize, fan_out: usize| -> f32 {
            let range = (6.0 / (fan_in + fan_out) as f32).sqrt();
            (rng.next_f32() * 2.0 - 1.0) * range
        };
        let w1: Vec<f32> = (0..META_Q_INPUT_DIM * META_Q_H1)
            .map(|_| xavier(&mut rng, META_Q_INPUT_DIM, META_Q_H1)).collect();
        let w2: Vec<f32> = (0..META_Q_H1 * META_Q_H2)
            .map(|_| xavier(&mut rng, META_Q_H1, META_Q_H2)).collect();
        let w3: Vec<f32> = (0..META_Q_H2)
            .map(|_| xavier(&mut rng, META_Q_H2, 1)).collect();
        Self {
            w1, b1: vec![0.0; META_Q_H1],
            w2, b2: vec![0.0; META_Q_H2],
            w3, b3: vec![0.0; 1],
            m_w1: vec![0.0; META_Q_INPUT_DIM * META_Q_H1], v_w1: vec![0.0; META_Q_INPUT_DIM * META_Q_H1],
            m_w2: vec![0.0; META_Q_H1 * META_Q_H2], v_w2: vec![0.0; META_Q_H1 * META_Q_H2],
            m_w3: vec![0.0; META_Q_H2], v_w3: vec![0.0; META_Q_H2],
            lr: 1e-3,
            step: 0,
            buffer: VecDeque::with_capacity(200),
            buffer_cap: 200,
        }
    }

    pub fn predict(&self, input: &MetaQInput) -> f32 {
        let x = input.to_vec();
        let mut h1 = vec![0.0f32; META_Q_H1];
        for j in 0..META_Q_H1 {
            let mut s = self.b1[j];
            for i in 0..META_Q_INPUT_DIM { s += x[i] * self.w1[i * META_Q_H1 + j]; }
            h1[j] = s.max(0.0);
        }
        let mut h2 = vec![0.0f32; META_Q_H2];
        for j in 0..META_Q_H2 {
            let mut s = self.b2[j];
            for i in 0..META_Q_H1 { s += h1[i] * self.w2[i * META_Q_H2 + j]; }
            h2[j] = s.max(0.0);
        }
        let mut logit = self.b3[0];
        for i in 0..META_Q_H2 { logit += h2[i] * self.w3[i]; }
        1.0 / (1.0 + (-logit).exp())
    }

    pub fn push_sample(&mut self, input: MetaQInput, label: f32) {
        if self.buffer.len() >= self.buffer_cap { self.buffer.pop_front(); }
        self.buffer.push_back((input, label));
    }

    pub fn train_step(&mut self) {
        if self.buffer.len() < 32 { return; }
        let mut rng = SimpleLcg::new(self.step as u64 ^ 0xAA55AA55);
        for _ in 0..32 {
            let idx = (rng.next_u32() as usize) % self.buffer.len();
            let (input, label) = self.buffer[idx].clone();
            self.train_one(&input, label);
        }
        self.step += 1;
    }

    fn train_one(&mut self, input: &MetaQInput, label: f32) {
        let x = input.to_vec();
        let mut h1_pre = vec![0.0f32; META_Q_H1];
        let mut h1 = vec![0.0f32; META_Q_H1];
        for j in 0..META_Q_H1 {
            let mut s = self.b1[j];
            for i in 0..META_Q_INPUT_DIM { s += x[i] * self.w1[i * META_Q_H1 + j]; }
            h1_pre[j] = s; h1[j] = s.max(0.0);
        }
        let mut h2_pre = vec![0.0f32; META_Q_H2];
        let mut h2 = vec![0.0f32; META_Q_H2];
        for j in 0..META_Q_H2 {
            let mut s = self.b2[j];
            for i in 0..META_Q_H1 { s += h1[i] * self.w2[i * META_Q_H2 + j]; }
            h2_pre[j] = s; h2[j] = s.max(0.0);
        }
        let mut logit = self.b3[0];
        for i in 0..META_Q_H2 { logit += h2[i] * self.w3[i]; }
        let pred = 1.0 / (1.0 + (-logit).exp());

        let d_logit = pred - label;
        let mut d_w3 = vec![0.0f32; META_Q_H2];
        for i in 0..META_Q_H2 { d_w3[i] = d_logit * h2[i]; }
        let d_b3 = d_logit;
        let d_h2: Vec<f32> = (0..META_Q_H2).map(|i| d_logit * self.w3[i] * (if h2_pre[i] > 0.0 { 1.0 } else { 0.0 })).collect();
        let mut d_w2 = vec![0.0f32; META_Q_H1 * META_Q_H2];
        let d_b2 = d_h2.clone();
        for i in 0..META_Q_H1 {
            for j in 0..META_Q_H2 { d_w2[i * META_Q_H2 + j] = d_h2[j] * h1[i]; }
        }
        let d_h1: Vec<f32> = (0..META_Q_H1).map(|i| {
            let mut s = 0.0f32;
            for j in 0..META_Q_H2 { s += d_h2[j] * self.w2[i * META_Q_H2 + j]; }
            s * (if h1_pre[i] > 0.0 { 1.0 } else { 0.0 })
        }).collect();
        let mut d_w1 = vec![0.0f32; META_Q_INPUT_DIM * META_Q_H1];
        let d_b1 = d_h1.clone();
        for i in 0..META_Q_INPUT_DIM {
            for j in 0..META_Q_H1 { d_w1[i * META_Q_H1 + j] = d_h1[j] * x[i]; }
        }
        self.adam_update(d_w1, d_w2, d_w3, d_b1, d_b2, d_b3);
    }

    fn adam_update(&mut self, d_w1: Vec<f32>, d_w2: Vec<f32>, d_w3: Vec<f32>, d_b1: Vec<f32>, d_b2: Vec<f32>, d_b3: f32) {
        let beta1 = 0.9f32; let beta2 = 0.999f32; let eps = 1e-8f32;
        let t = (self.step + 1) as f32;
        let bc1 = 1.0 - beta1.powf(t); let bc2 = 1.0 - beta2.powf(t);
        let lr = self.lr;

        let adam = |p: &mut [f32], m: &mut [f32], v: &mut [f32], g: &[f32]| {
            for i in 0..p.len() {
                m[i] = beta1 * m[i] + (1.0 - beta1) * g[i];
                v[i] = beta2 * v[i] + (1.0 - beta2) * g[i] * g[i];
                let mh = m[i] / bc1;
                let vh = v[i] / bc2;
                p[i] -= lr * mh / (vh.sqrt() + eps);
            }
        };
        adam(&mut self.w1, &mut self.m_w1, &mut self.v_w1, &d_w1);
        adam(&mut self.w2, &mut self.m_w2, &mut self.v_w2, &d_w2);
        adam(&mut self.w3, &mut self.m_w3, &mut self.v_w3, &d_w3);
        for i in 0..self.b1.len() { self.b1[i] -= lr * d_b1[i]; }
        for i in 0..self.b2.len() { self.b2[i] -= lr * d_b2[i]; }
        self.b3[0] -= lr * d_b3;
    }
}

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

/// Deterministic LCG for reproducible weight init and minibatch sampling.
#[derive(Debug)]
struct SimpleLcg { state: u64 }

impl SimpleLcg {
    fn new(seed: u64) -> Self { Self { state: seed.wrapping_add(1) } }
    fn next_u32(&mut self) -> u32 {
        self.state = self.state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
        (self.state >> 33) as u32
    }
    fn next_f32(&mut self) -> f32 {
        (self.next_u32() as f32) / (u32::MAX as f32)
    }
}

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

    #[test]
    fn test_predict_returns_probability() {
        let net = MetaQNetwork::new();
        let input = MetaQInput {
            q_gap_ema: 0.3, log_grad_norm: 2.0, atom_util: 0.5,
            ens_disagreement: 0.1, regime_stability: 0.7,
            spectral_gap_norm: 0.6, grad_consistency: 0.3,
        };
        let p = net.predict(&input);
        assert!((0.0..=1.0).contains(&p), "probability out of range: {}", p);
    }

    #[test]
    fn test_training_reduces_loss() {
        let mut net = MetaQNetwork::new();
        let collapse_input = MetaQInput {
            q_gap_ema: 0.01, log_grad_norm: 10.0, atom_util: 0.1,
            ens_disagreement: 0.0, regime_stability: 0.5,
            spectral_gap_norm: 0.1, grad_consistency: -0.3,
        };
        for _ in 0..100 { net.push_sample(collapse_input.clone(), 1.0); }
        for _ in 0..50 { net.train_step(); }
        let p = net.predict(&collapse_input);
        assert!(p > 0.5, "Trained net should predict >0.5 for collapse state, got {}", p);
    }
}
