/// nn.TransformerEncoderLayer equivalent (Burn 0.20.1)
///
/// Supports both pre-norm (norm_first=True) and post-norm (norm_first=False).
///
/// PyTorch stores self-attention as fused in_proj_weight [3*D, D] / in_proj_bias [3*D].
///
/// Architecture:
///   self_attn: MultiheadAttention (fused QKV)
///   linear1:   Linear(D, ffn_dim)
///   linear2:   Linear(ffn_dim, D)
///   norm1:     LayerNorm(D)
///   norm2:     LayerNorm(D)

use burn::prelude::*;
use burn::nn::{Linear, LinearConfig, LayerNorm, LayerNormConfig};

/// Fused multi-head attention with a single in_proj (Q,K,V packed).
#[derive(Module, Debug)]
pub struct FusedMultiheadAttention<B: Backend> {
    pub in_proj: Linear<B>,   // [D, 3*D]  (burn stores transposed)
    pub out_proj: Linear<B>,  // [D, D]
    pub num_heads: usize,
    pub head_dim: usize,
}

impl<B: Backend> FusedMultiheadAttention<B> {
    pub fn new(d_model: usize, num_heads: usize, device: &B::Device) -> Self {
        let in_proj = LinearConfig::new(d_model, 3 * d_model).with_bias(true).init(device);
        let out_proj = LinearConfig::new(d_model, d_model).with_bias(true).init(device);
        Self {
            in_proj,
            out_proj,
            num_heads,
            head_dim: d_model / num_heads,
        }
    }

    /// Forward: x [B, S, D], key_padding_mask [B, S] (true = ignore)
    /// Returns: [B, S, D]
    pub fn forward(
        &self,
        x: Tensor<B, 3>,
        key_padding_mask: Option<Tensor<B, 2, Bool>>,
    ) -> Tensor<B, 3> {
        let [b, s, d] = x.dims();
        let h = self.num_heads;
        let hd = self.head_dim;

        // QKV projection: [B, S, D] → [B, S, 3*D]
        let qkv = self.in_proj.forward(x);

        // Split into Q, K, V each [B, S, D]
        let q = qkv.clone().slice([0..b, 0..s, 0..d]);
        let k = qkv.clone().slice([0..b, 0..s, d..(2 * d)]);
        let v = qkv.slice([0..b, 0..s, (2 * d)..(3 * d)]);

        // Reshape to [B, H, S, Hd]
        let q = q.reshape([b, s, h, hd]).swap_dims(1, 2);
        let k = k.reshape([b, s, h, hd]).swap_dims(1, 2);
        let v = v.reshape([b, s, h, hd]).swap_dims(1, 2);

        // Scaled dot-product attention
        let scale = (hd as f64).sqrt();
        let scores = q.matmul(k.swap_dims(2, 3)) / scale;  // [B, H, S, S]

        // Apply key_padding_mask: mask out padding positions in key dimension
        let scores = if let Some(ref mask) = key_padding_mask {
            // mask: [B, S] where true = padding
            // Expand to [B, 1, 1, S] for broadcasting over [B, H, Sq, Sk]
            let mask_4d = mask.clone()
                .reshape([b, 1, 1, s])
                .expand([b, h, s, s]);
            scores.mask_fill(mask_4d, -1e9)
        } else {
            scores
        };

        let attn = burn::tensor::activation::softmax(scores, 3);
        let out = attn.matmul(v);  // [B, H, S, Hd]

        // Reshape back: [B, H, S, Hd] → [B, S, D]
        let out = out.swap_dims(1, 2).reshape([b, s, d]);
        self.out_proj.forward(out)
    }
}

/// TransformerEncoderLayer with fused MHA.
#[derive(Module, Debug)]
pub struct TransformerEncoderLayer<B: Backend> {
    pub self_attn: FusedMultiheadAttention<B>,
    pub linear1: Linear<B>,
    pub linear2: Linear<B>,
    pub norm1: LayerNorm<B>,
    pub norm2: LayerNorm<B>,
    pub norm_first: bool,
}

impl<B: Backend> TransformerEncoderLayer<B> {
    pub fn new(
        d_model: usize,
        num_heads: usize,
        ffn_dim: usize,
        norm_first: bool,
        device: &B::Device,
    ) -> Self {
        Self {
            self_attn: FusedMultiheadAttention::new(d_model, num_heads, device),
            linear1: LinearConfig::new(d_model, ffn_dim).with_bias(true).init(device),
            linear2: LinearConfig::new(ffn_dim, d_model).with_bias(true).init(device),
            norm1: LayerNormConfig::new(d_model).init(device),
            norm2: LayerNormConfig::new(d_model).init(device),
            norm_first,
        }
    }

    /// Forward: x [B, S, D]
    pub fn forward(
        &self,
        x: Tensor<B, 3>,
        key_padding_mask: Option<Tensor<B, 2, Bool>>,
    ) -> Tensor<B, 3> {
        if self.norm_first {
            // Pre-norm
            let x2 = self.norm1.forward(x.clone());
            let x2 = self.self_attn.forward(x2, key_padding_mask);
            let x = x + x2;

            let x2 = self.norm2.forward(x.clone());
            let x2 = self.linear1.forward(x2);
            let x2 = burn::tensor::activation::relu(x2);
            let x2 = self.linear2.forward(x2);
            x + x2
        } else {
            // Post-norm (PyTorch default, used by SleepFM spatial/temporal pooling)
            let x2 = self.self_attn.forward(x.clone(), key_padding_mask);
            let x = self.norm1.forward(x + x2);

            let x2 = self.linear1.forward(x.clone());
            let x2 = burn::tensor::activation::relu(x2);
            let x2 = self.linear2.forward(x2);
            self.norm2.forward(x + x2)
        }
    }
}
