use rand::{rng, RngExt};
use rand::distr::Uniform;
use serde::{Serialize, Deserialize};

use crate::activations::{Activation, activate, derivative};
use crate::error::{TuphiumError, TuphiumResult};
use crate::layers::module::Module;
use crate::linalg::Matrix;
use crate::optimizers::Optimizer;

// Xavier Glorot initialization (out_features x in_features)
fn glorot(in_features: usize, out_features: usize) -> TuphiumResult<Matrix> {
    let limit = (6.0 / (in_features as f32 + out_features as f32)).sqrt();
    let mut rng = rng();
    let dist = Uniform::new(-limit, limit).map_err(|error| TuphiumError::NumericDomain {
        operation: "Layer::new",
        input: limit,
        reason: match error {
            _ => "invalid Glorot initialization range",
        },
    })?;
    let data = (0..in_features * out_features)
        .map(|_| rng.sample(dist))
        .collect();
    Matrix::from_data(out_features, in_features, data)
}

// Kaiming He initialization (out_features x in_features)
fn he(in_features: usize, out_features: usize) -> TuphiumResult<Matrix> {
    let limit = (6.0 / in_features as f32).sqrt();
    let mut rng = rng();
    let dist = Uniform::new(-limit, limit).map_err(|error| TuphiumError::NumericDomain {
        operation: "Layer::new",
        input: limit,
        reason: match error {
            _ => "invalid He initialization range",
        },
    })?;
    let data = (0..in_features * out_features)
        .map(|_| rng.sample(dist))
        .collect();
    Matrix::from_data(out_features, in_features, data)
}

#[derive(Serialize, Deserialize)]
pub struct Layer {
    pub weights: Matrix, // out x in
    pub bias: Vec<f32>,  // out
    pub activation: Activation,

    #[serde(skip)]
    last_input: Vec<f32>,
    #[serde(skip)]
    last_output: Vec<f32>,
    #[serde(skip)]
    grad_weights: Matrix,
    #[serde(skip)]
    grad_bias: Vec<f32>,
    #[serde(skip)]
    has_forward: bool,
}

impl Layer {
    /// Builds a dense layer mapping `in_features` inputs to `out_features` outputs.
    ///
    /// Weights are initialized according to the activation: Glorot/Xavier for
    /// `Sigmoid` and Kaiming/He otherwise (e.g. `ReLU`). Biases start at zero.
    ///
    /// # Examples
    ///
    /// ```
    /// use tuphium::activations::Activation;
    /// use tuphium::layers::Layer;
    ///
    /// let layer = Layer::new(4, 8, Activation::ReLU).unwrap();
    /// assert_eq!(layer.weights.rows, 8);
    /// assert_eq!(layer.weights.cols, 4);
    /// ```
    pub fn new(
        in_features: usize,
        out_features: usize,
        activation: Activation,
    ) -> TuphiumResult<Self> {
        if in_features == 0 {
            return Err(TuphiumError::InvalidDimension {
                operation: "Layer::new",
                dimension: "input",
                value: in_features,
                reason: "must be greater than zero",
            });
        }
        if out_features == 0 {
            return Err(TuphiumError::InvalidDimension {
                operation: "Layer::new",
                dimension: "output",
                value: out_features,
                reason: "must be greater than zero",
            });
        }
        let weights = match activation {
            Activation::Sigmoid => glorot(in_features, out_features)?,
            _ => he(in_features, out_features)?, // ReLU / Tanh
        };
        Ok(Self {
            weights,
            bias: vec![0.0; out_features],
            activation,
            last_input: vec![0.0; in_features],
            last_output: vec![0.0; out_features],
            grad_weights: Matrix::zeros(out_features, in_features)?,
            grad_bias: vec![0.0; out_features],
            has_forward: false,
        })
    }

    /// Rebuild transient (#[serde(skip)]) buffers after deserialization.
    pub fn rebuild_buffers(&mut self) {
        let out = self.weights.rows;
        let inp = self.weights.cols;
        self.last_input = vec![0.0; inp];
        self.last_output = vec![0.0; out];
        self.grad_weights =
            Matrix::zeros(out, inp).expect("Layer::rebuild_buffers: valid stored dimensions");
        self.grad_bias = vec![0.0; out];
        self.has_forward = false;
    }

    /// Apply an optimizer step to weights and bias using accumulated gradients.
    pub fn step(&mut self, opt: &mut dyn Optimizer) -> TuphiumResult<()> {
        // read accumulated gradients from GPU before optimizer reads grad_weights.data
        self.grad_weights.sync_from_device()?;
        opt.update(&mut self.weights.data, &self.grad_weights.data)?;
        opt.update(&mut self.bias, &self.grad_bias)?;
        // mark weights dirty so next forward/backward re-uploads to GPU
        self.weights.mark_dirty();
        Ok(())
    }

    fn validate_finite_slice(operation: &'static str, values: &[f32]) -> TuphiumResult<()> {
        if let Some(&value) = values.iter().find(|value| !value.is_finite()) {
            return Err(TuphiumError::NonFiniteValue { operation, value });
        }
        Ok(())
    }
}

impl Module for Layer {
    /// Computes `activation(W·x + b)`, caching the input and output for backward.
    fn forward(&mut self, x: &[f32]) -> TuphiumResult<Vec<f32>> {
        if x.len() != self.weights.cols {
            return Err(TuphiumError::ShapeMismatch {
                operation: "Layer::forward",
                expected: vec![self.weights.cols],
                actual: vec![x.len()],
            });
        }
        Self::validate_finite_slice("Layer::forward", x)?;
        self.last_input = x.to_vec();
        let mut z = self.weights.matvec(x)?; // W·x
        for i in 0..z.len() {
            z[i] += self.bias[i];
        } // + b
        let h = activate(z, &self.activation);
        Self::validate_finite_slice("Layer::forward", &h)?;
        self.last_output = h.clone();
        self.has_forward = true;
        Ok(h)
    }

    /// Accumulates weight and bias gradients from `grad_output` and returns the
    /// gradient with respect to the layer's input (`Wᵀ · dz`).
    ///
    /// Relies on the input/output cached by the preceding [`Module::forward`].
    fn backward(&mut self, grad_output: &[f32]) -> TuphiumResult<Vec<f32>> {
        if !self.has_forward {
            return Err(TuphiumError::InvalidState {
                operation: "Layer::backward",
                state: "forward must be called before backward",
            });
        }
        if grad_output.len() != self.bias.len() {
            return Err(TuphiumError::ShapeMismatch {
                operation: "Layer::backward",
                expected: vec![self.bias.len()],
                actual: vec![grad_output.len()],
            });
        }
        Self::validate_finite_slice("Layer::backward", grad_output)?;
        // dL/dz = grad_output ⊙ act'(h)
        let dz: Vec<f32> = (0..self.bias.len())
            .map(|i| grad_output[i] * derivative(self.last_output[i], &self.activation))
            .collect();
        Self::validate_finite_slice("Layer::backward", &dz)?;

        // accumulate gradients (caller zero_grad's before the optimizer step)
        self.grad_weights.add_outer(&dz, &self.last_input)?; // dz ⊗ x
        for i in 0..self.grad_bias.len() {
            self.grad_bias[i] += dz[i];
        }

        // dL/d(input) = Wᵀ · dz
        self.weights.t_matvec(&dz)
    }

    fn params(&mut self) -> Vec<&mut f32> {
        let mut out: Vec<&mut f32> = self.weights.data.iter_mut().collect();
        out.extend(self.bias.iter_mut());
        out
    }

    fn grads(&self) -> Vec<f32> {
        let mut out = self.grad_weights.data.clone();
        out.extend_from_slice(&self.grad_bias);
        out
    }

    fn zero_grad(&mut self) {
        self.grad_weights.fill(0.0);
        for g in self.grad_bias.iter_mut() {
            *g = 0.0;
        }
    }

    fn step(&mut self, opt: &mut dyn Optimizer) -> TuphiumResult<()> {
        Layer::step(self, opt)
    }
}
