/// LayerNorm2d — normalises over the last 2 dimensions jointly.
///
/// Matches PyTorch `nn.LayerNorm([C, T])` applied to input `[N, C, T]`.
///
/// PyTorch computes mean/var over dims (1, 2) jointly, then applies
/// learnable weight [C, T] and bias [C, T] element-wise.

use burn::prelude::*;
use burn::module::Param;

#[derive(Module, Debug)]
pub struct LayerNorm2d<B: Backend> {
    /// Learnable scale [C, T].
    pub gamma: Param<Tensor<B, 2>>,
    /// Learnable bias [C, T].
    pub beta: Param<Tensor<B, 2>>,
    pub eps: f64,
    pub channels: usize,
    pub spatial: usize,
}

impl<B: Backend> LayerNorm2d<B> {
    pub fn new(channels: usize, spatial: usize, eps: f64, device: &B::Device) -> Self {
        let gamma = Param::from_tensor(Tensor::ones([channels, spatial], device));
        let beta = Param::from_tensor(Tensor::zeros([channels, spatial], device));
        Self { gamma, beta, eps, channels, spatial }
    }

    /// Forward: [N, C, T] → [N, C, T]
    ///
    /// Computes mean/var over (C, T) jointly per sample, then applies γ and β.
    pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
        let [n, c, t] = x.dims();
        debug_assert_eq!(c, self.channels);
        debug_assert_eq!(t, self.spatial);

        let num_elements = (c * t) as f64;

        // Flatten last 2 dims: [N, C*T]
        let flat = x.clone().reshape([n, c * t]);

        // Mean per sample: [N, 1]
        let mean = flat.clone().sum_dim(1).reshape([n, 1]) / num_elements;

        // Var per sample: [N, 1]
        let diff = flat - mean.clone().expand([n, c * t]);
        let var = diff.clone().powf_scalar(2.0).sum_dim(1).reshape([n, 1]) / num_elements;

        // Normalise: [N, C*T]
        let inv_std = (var + self.eps).sqrt().recip();
        let normed = diff * inv_std.expand([n, c * t]);

        // Reshape back to [N, C, T] and apply affine
        let normed = normed.reshape([n, c, t]);
        let gamma = self.gamma.val().unsqueeze_dim::<3>(0).expand([n, c, t]);
        let beta = self.beta.val().unsqueeze_dim::<3>(0).expand([n, c, t]);

        normed * gamma + beta
    }
}
