//! Weight initialization for GPU tensors.
//!
//! Xavier (Glorot) and Kaiming (He) initialization generate values on CPU
//! and upload to GPU. For typical ML weight sizes (thousands to millions of
//! elements) the CPU generation + HtoD transfer is negligible compared to
//! a full training step, and avoids the complexity of a GPU RNG kernel.

use std::sync::Arc;

use cudarc::driver::{CudaSlice, CudaStream};

use crate::MLError;

/// Initialize a `CudaSlice<f32>` with Xavier uniform values.
///
/// Samples from U(-limit, limit) where `limit = sqrt(6 / (fan_in + fan_out))`.
/// Shape: `fan_out * fan_in` elements (row-major, matching `[out, in]` convention).
pub fn xavier_uniform(
    fan_in: usize,
    fan_out: usize,
    stream: &Arc<CudaStream>,
) -> Result<CudaSlice<f32>, MLError> {
    let n = fan_out * fan_in;
    let limit = (6.0_f64 / (fan_in + fan_out) as f64).sqrt();
    let host = generate_uniform(n, -limit, limit);
    upload_to_gpu(&host, stream)
}

/// Initialize a `CudaSlice<f32>` with Kaiming uniform values (for ReLU networks).
///
/// Samples from U(-limit, limit) where `limit = sqrt(6 / fan_in)`.
/// Shape: `fan_out * fan_in` elements.
pub fn kaiming_uniform(
    fan_in: usize,
    fan_out: usize,
    stream: &Arc<CudaStream>,
) -> Result<CudaSlice<f32>, MLError> {
    let n = fan_out * fan_in;
    let limit = (6.0_f64 / fan_in as f64).sqrt();
    let host = generate_uniform(n, -limit, limit);
    upload_to_gpu(&host, stream)
}

/// Initialize a bias vector: U(-bound, bound) where `bound = 1/sqrt(fan_in)`.
pub fn bias_uniform(
    fan_in: usize,
    fan_out: usize,
    stream: &Arc<CudaStream>,
) -> Result<CudaSlice<f32>, MLError> {
    let bound = 1.0 / (fan_in as f64).sqrt();
    let host = generate_uniform(fan_out, -bound, bound);
    upload_to_gpu(&host, stream)
}

/// Initialize a `CudaSlice<f32>` with zeros.
pub fn zeros(
    n: usize,
    stream: &Arc<CudaStream>,
) -> Result<CudaSlice<f32>, MLError> {
    stream.alloc_zeros::<f32>(n).map_err(|e| {
        MLError::ModelError(format!("init::zeros({n}): {e}"))
    })
}

/// Initialize a `CudaSlice<f32>` with near-zero Xavier values.
///
/// 0.01x Xavier scale — useful for distributional output layers where
/// near-uniform softmax is desired at initialization.
pub fn near_zero_xavier(
    fan_in: usize,
    fan_out: usize,
    stream: &Arc<CudaStream>,
) -> Result<CudaSlice<f32>, MLError> {
    let n = fan_out * fan_in;
    let limit = 0.01 * (6.0_f64 / (fan_in + fan_out) as f64).sqrt();
    let host = generate_uniform(n, -limit, limit);
    upload_to_gpu(&host, stream)
}

// ── Internal helpers ──────────────────────────────────────────────────────

/// Generate `n` uniform random f32 values in `[lo, hi)`.
///
/// Uses a simple xoshiro256++ PRNG seeded from thread_rng for speed.
/// This runs on CPU — acceptable for one-time initialization.
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;

    // Seed from time + thread id for reasonable uniqueness
    let mut hasher = DefaultHasher::new();
    SystemTime::now().hash(&mut hasher);
    std::thread::current().id().hash(&mut hasher);
    let seed = hasher.finish();

    // Simple xoshiro256++ state seeded from the hash
    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 {
        // xoshiro256++ step
        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);

        // Convert to f64 in [0, 1) then scale to [lo, hi)
        let f = (result >> 11) as f64 / (1_u64 << 53) as f64;
        out.push((lo + f * range) as f32);
    }
    out
}

/// Upload host f32 slice to a new CudaSlice on the given stream.
///
/// Public for use by downstream crates that need to initialize GPU parameters
/// from host data (e.g., RMSNorm weights initialized to ones).
pub fn upload_to_gpu(
    host: &[f32],
    stream: &Arc<CudaStream>,
) -> Result<CudaSlice<f32>, MLError> {
    let mut buf = stream.alloc_zeros::<f32>(host.len()).map_err(|e| {
        MLError::ModelError(format!("init upload alloc: {e}"))
    })?;
    stream.memcpy_htod(host, &mut buf).map_err(|e| {
        MLError::ModelError(format!("init upload HtoD: {e}"))
    })?;
    Ok(buf)
}

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

    fn make_stream() -> Arc<CudaStream> {
        cudarc::driver::CudaContext::new(0)
                .expect("CUDA required")
                .new_stream()
                .expect("fork stream")
    }

    #[test]
    fn test_xavier_shape() {
        let stream = make_stream();
        let w = xavier_uniform(64, 128, &stream).unwrap();
        assert_eq!(w.len(), 128 * 64);
    }

    #[test]
    fn test_xavier_statistics() {
        let stream = make_stream();
        let fan_in = 256;
        let fan_out = 256;
        let w = xavier_uniform(fan_in, fan_out, &stream).unwrap();
        let mut host = vec![0.0_f32; w.len()];
        stream.memcpy_dtoh(&w, &mut host).unwrap(); // test-only readback

        let mean: f32 = host.iter().sum::<f32>() / host.len() as f32;
        let variance: f32 = host.iter().map(|x| (x - mean) * (x - mean)).sum::<f32>()
            / host.len() as f32;
        let expected_var = 2.0 / (fan_in + fan_out) as f32;

        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}"
        );
    }

    #[test]
    fn test_kaiming_shape() {
        let stream = make_stream();
        let w = kaiming_uniform(128, 64, &stream).unwrap();
        assert_eq!(w.len(), 64 * 128);
    }

    #[test]
    fn test_bias_shape() {
        let stream = make_stream();
        let b = bias_uniform(64, 128, &stream).unwrap();
        assert_eq!(b.len(), 128);
    }

    #[test]
    fn test_zeros() {
        let stream = make_stream();
        let z = zeros(100, &stream).unwrap();
        let mut host = vec![1.0_f32; 100];
        stream.memcpy_dtoh(&z, &mut host).unwrap(); // test-only readback
        assert!(host.iter().all(|&x| x == 0.0));
    }
}
