//! Xavier/Glorot initialization helpers (CPU-side).
//!
//! GPU-side initialization is handled by `cuda_autograd::init` which directly
//! allocates and fills GPU buffers. This module provides CPU-side convenience
//! functions used by downstream crates that still operate on host data.
//!
//! Reference: Glorot & Bengio (2010)

/// Compute the Xavier uniform limit: `sqrt(6 / (fan_in + fan_out))`.
pub fn xavier_limit(fan_in: usize, fan_out: usize) -> f64 {
    (6.0 / (fan_in + fan_out) as f64).sqrt()
}

/// Compute the Kaiming uniform limit: `sqrt(6 / fan_in)`.
pub fn kaiming_limit(fan_in: usize) -> f64 {
    (6.0 / fan_in as f64).sqrt()
}

/// Generate `n` uniform random f32 values in `[lo, hi)`.
///
/// Uses a simple xoshiro256++ PRNG seeded from system time + thread id.
pub fn generate_uniform(n: usize, lo: f64, hi: f64) -> Vec<f32> {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};
    use std::time::SystemTime;

    let mut hasher = DefaultHasher::new();
    SystemTime::now().hash(&mut hasher);
    std::thread::current().id().hash(&mut hasher);
    let seed = hasher.finish();

    let mut s = [
        seed,
        seed.wrapping_mul(6364136223846793005).wrapping_add(1),
        seed.wrapping_mul(1442695040888963407).wrapping_add(3),
        seed.wrapping_mul(2891336453).wrapping_add(7),
    ];

    let range = hi - lo;
    let mut out = Vec::with_capacity(n);
    for _ in 0..n {
        let result = s[0].wrapping_add(s[3]).rotate_left(23).wrapping_add(s[0]);
        let t = s[1] << 17;
        s[2] ^= s[0];
        s[3] ^= s[1];
        s[1] ^= s[2];
        s[0] ^= s[3];
        s[2] ^= t;
        s[3] = s[3].rotate_left(45);

        let f = (result >> 11) as f64 / (1_u64 << 53) as f64;
        out.push((lo + f * range) as f32);
    }
    out
}

/// Generate Xavier uniform weights on CPU: shape `[fan_out, fan_in]`.
pub fn xavier_uniform_cpu(fan_in: usize, fan_out: usize) -> Vec<f32> {
    let limit = xavier_limit(fan_in, fan_out);
    generate_uniform(fan_out * fan_in, -limit, limit)
}

/// Verify Xavier initialization statistics for testing.
///
/// Returns `(mean, variance, expected_variance)`.
pub fn verify_xavier_stats_cpu(
    weights: &[f32],
    fan_in: usize,
    fan_out: usize,
) -> (f32, f32, f32) {
    let n = weights.len() as f32;
    let mean: f32 = weights.iter().sum::<f32>() / n;
    let variance: f32 = weights.iter().map(|x| (x - mean) * (x - mean)).sum::<f32>() / n;
    let expected_var = 2.0 / (fan_in + fan_out) as f32;
    (mean, variance, expected_var)
}

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

    #[test]
    fn test_xavier_uniform_cpu_shape() {
        let weights = xavier_uniform_cpu(64, 128);
        assert_eq!(weights.len(), 128 * 64);
    }

    #[test]
    fn test_xavier_uniform_cpu_statistics() {
        let fan_in = 256;
        let fan_out = 256;
        let weights = xavier_uniform_cpu(fan_in, fan_out);

        let (mean, variance, expected_var) = verify_xavier_stats_cpu(&weights, fan_in, fan_out);

        assert!(mean.abs() < 0.02, "Mean should be ~0, got {mean}");
        assert!(
            (variance - expected_var).abs() / expected_var < 0.2,
            "Variance {variance} should be ~{expected_var}"
        );
    }
}
