//! Parameter initialization functions

use crate::Result;
use rand::prelude::*;
use rand_distr::{Distribution, Normal, Uniform};
use rtx_tensor::Tensor;

/// Non-linearity types for initialization
#[derive(Debug, Clone, Copy)]
pub enum NonLinearity {
    /// Linear activation
    Linear,
    /// ReLU activation
    ReLU,
}

impl NonLinearity {
    /// Get the gain factor for this non-linearity
    pub fn gain(self) -> f32 {
        match self {
            Self::Linear => 1.0,
            Self::ReLU => (2.0_f32).sqrt(),
        }
    }
}

/// Initialize tensor with uniform distribution
pub fn uniform(tensor: &mut Tensor, low: f32, high: f32) -> Result<()> {
    let shape = tensor.shape().dims().to_vec();
    let device = tensor.device().clone();
    let numel = tensor.numel();

    let mut rng = thread_rng();
    let dist = Uniform::new(low, high);
    let data: Vec<f32> = (0..numel).map(|_| dist.sample(&mut rng)).collect();

    *tensor = Tensor::from_data(data, shape, &device)?;
    Ok(())
}

/// Initialize tensor with normal distribution
pub fn normal(tensor: &mut Tensor, mean: f32, std: f32) -> Result<()> {
    let shape = tensor.shape().dims().to_vec();
    let device = tensor.device().clone();
    let numel = tensor.numel();

    let mut rng = thread_rng();
    let dist = Normal::new(mean, std).map_err(|e| crate::Error::InvalidParameter(e.to_string()))?;
    let data: Vec<f32> = (0..numel).map(|_| dist.sample(&mut rng)).collect();

    *tensor = Tensor::from_data(data, shape, &device)?;
    Ok(())
}

/// Initialize tensor with zeros
pub fn zeros(tensor: &mut Tensor) -> Result<()> {
    *tensor = Tensor::zeros_like(tensor)?;
    Ok(())
}

/// Initialize tensor with ones
pub fn ones(tensor: &mut Tensor) -> Result<()> {
    *tensor = Tensor::ones_like(tensor)?;
    Ok(())
}

/// Kaiming uniform initialization
pub fn kaiming_uniform(
    tensor: &mut Tensor,
    fan_in: usize,
    nonlinearity: NonLinearity,
) -> Result<()> {
    let gain = nonlinearity.gain();
    let bound = gain * (3.0 / fan_in as f32).sqrt();
    uniform(tensor, -bound, bound)
}

/// Kaiming normal initialization
pub fn kaiming_normal(
    tensor: &mut Tensor,
    fan_in: usize,
    nonlinearity: NonLinearity,
) -> Result<()> {
    let gain = nonlinearity.gain();
    let std = gain / (fan_in as f32).sqrt();
    normal(tensor, 0.0, std)
}

/// Xavier uniform initialization
pub fn xavier_uniform(tensor: &mut Tensor, fan_in: usize, fan_out: usize) -> Result<()> {
    let bound = (6.0 / (fan_in + fan_out) as f32).sqrt();
    uniform(tensor, -bound, bound)
}

/// Xavier normal initialization
pub fn xavier_normal(tensor: &mut Tensor, fan_in: usize, fan_out: usize) -> Result<()> {
    let std = (2.0 / (fan_in + fan_out) as f32).sqrt();
    normal(tensor, 0.0, std)
}