/// AttentionPooling — single TransformerEncoderLayer followed by masked mean pooling.
///
/// Python:
///   transformer_layer = nn.TransformerEncoderLayer(
///       d_model=input_dim, nhead=num_heads, dropout=dropout,
///       batch_first=True
///   )
///   # forward: run transformer, then masked mean over seq dim
///
/// Note: PyTorch TransformerEncoderLayer default is norm_first=False (post-norm).
/// The SleepFM model uses the default, so this is post-norm.

use burn::prelude::*;
use crate::model::transformer_layer::TransformerEncoderLayer;

#[derive(Module, Debug)]
pub struct AttentionPooling<B: Backend> {
    pub transformer_layer: TransformerEncoderLayer<B>,
}

impl<B: Backend> AttentionPooling<B> {
    pub fn new(input_dim: usize, num_heads: usize, device: &B::Device) -> Self {
        // PyTorch default: dim_feedforward=2048, norm_first=False
        let transformer_layer = TransformerEncoderLayer::new(
            input_dim, num_heads, 2048, false, device,
        );
        Self { transformer_layer }
    }

    /// Forward pass with optional key_padding_mask.
    ///
    /// x: [B, S, D]
    /// mask: [B, S] where true = padding (to be ignored)
    ///
    /// Returns: [B, D] (pooled output)
    pub fn forward(
        &self,
        x: Tensor<B, 3>,
        mask: Option<Tensor<B, 2, Bool>>,
    ) -> Tensor<B, 2> {
        let [_batch_size, _seq_len, _d] = x.dims();

        match mask {
            Some(ref m) => {
                let [b, s, d] = x.dims();
                // If all positions masked (seq_len=1 edge case), just mean
                if s == 1 {
                    return x.reshape([b, d]);
                }

                let out = self.transformer_layer.forward(x, Some(m.clone()));

                // Invert mask: true for valid → false for padding
                let valid = m.clone().bool_not();  // [B, S]
                let valid_f = valid.float();  // [B, S]

                // Expand for element-wise multiply: [B, S] → [B, S, D]
                let valid_expanded = valid_f.clone().unsqueeze_dim::<3>(2).expand([b, s, d]);

                // Masked sum along seq dim
                let masked_out = out * valid_expanded;
                // Sum over seq dim: [B, S, D] → [B, D]
                let sum = masked_out.sum_dim(1).reshape([b, d]);

                // Count valid positions per batch: [B, S] → sum → [B, 1] → expand → [B, D]
                let count = valid_f.sum_dim(1).reshape([b, 1]).expand([b, d]).clamp_min(1.0);

                sum / count
            }
            None => {
                let [b, _s, d] = x.dims();
                let out = self.transformer_layer.forward(x, None);
                out.sum_dim(1).reshape([b, d]) / _s as f64
            }
        }
    }
}
