/// SetTransformer — the core SleepFM foundation model.
///
/// Python architecture:
///   1. patch_embedding (Tokenizer): raw signal → per-patch tokens
///   2. spatial_pooling (AttentionPooling): pool across channels per time step
///   3. positional_encoding + layer_norm
///   4. transformer_encoder (N layers): temporal modelling
///   5. temporal_pooling (AttentionPooling): pool across time → single embedding
///
/// Forward(x [B, C, T], mask [B, C]):
///   tokens = patch_embedding(x)         → [B, C, S, E]
///   rearrange to [(B*S), C, E]
///   spatial = spatial_pooling(tokens, mask) → [(B*S), E] → [B, S, E]
///   x = positional_encoding(spatial)
///   x = layer_norm(x)
///   x = transformer_encoder(x)
///   embedding = x.clone()
///   x = temporal_pooling(x)              → [B, E]
///   return (x, embedding)

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

use crate::model::tokenizer::Tokenizer;
use crate::model::attention_pooling::AttentionPooling;
use crate::model::positional_encoding::{positional_encoding, add_positional_encoding};
use crate::model::transformer_layer::TransformerEncoderLayer;

#[derive(Module, Debug)]
pub struct SetTransformer<B: Backend> {
    pub patch_embedding: Tokenizer<B>,
    pub spatial_pooling: AttentionPooling<B>,
    pub pe_buffer: Tensor<B, 3>,  // [1, max_seq_len, embed_dim]
    pub layer_norm: LayerNorm<B>,
    pub transformer_encoder: Vec<TransformerEncoderLayer<B>>,
    pub temporal_pooling: AttentionPooling<B>,
    pub embed_dim: usize,
}

impl<B: Backend> SetTransformer<B> {
    pub fn new(
        patch_size: usize,
        embed_dim: usize,
        num_heads: usize,
        num_layers: usize,
        pooling_head: usize,
        max_seq_length: usize,
        device: &B::Device,
    ) -> Self {
        let patch_embedding = Tokenizer::new(patch_size, embed_dim, device);
        let spatial_pooling = AttentionPooling::new(embed_dim, pooling_head, device);

        let pe_buffer = positional_encoding(max_seq_length, embed_dim, device);
        let layer_norm = LayerNormConfig::new(embed_dim).init(device);

        // TransformerEncoder: num_layers of TransformerEncoderLayer
        // SleepFM transformer_encoder uses norm_first=True (pre-norm)
        let transformer_encoder = (0..num_layers)
            .map(|_| TransformerEncoderLayer::new(
                embed_dim, num_heads, 2048, true, device,
            ))
            .collect();

        let temporal_pooling = AttentionPooling::new(embed_dim, pooling_head, device);

        Self {
            patch_embedding,
            spatial_pooling,
            pe_buffer,
            layer_norm,
            transformer_encoder,
            temporal_pooling,
            embed_dim,
        }
    }

    /// Forward pass.
    ///
    /// x: [B, C, T] — raw multichannel signal
    /// mask: [B, C] — channel mask (true = padding, to be ignored)
    ///
    /// Returns: (pooled [B, E], embeddings [B, S, E])
    pub fn forward(
        &self,
        x: Tensor<B, 3>,
        mask: Tensor<B, 2, Bool>,
    ) -> (Tensor<B, 2>, Tensor<B, 3>) {
        // 1. Patch embedding: [B, C, T] → [B, C, S, E]
        let tokens = self.patch_embedding.forward(x);
        let [b, _c, s, e] = tokens.dims();

        // 2. Rearrange: [B, C, S, E] → [(B*S), C, E]
        let c = tokens.dims()[1];
        let tokens = tokens.swap_dims(1, 2)  // [B, S, C, E]
            .reshape([b * s, c, e]);

        // Expand mask: [B, C] → [B, S, C] → [(B*S), C]
        let mask_expanded = mask.clone()
            .unsqueeze_dim::<3>(1)   // [B, 1, C]
            .expand([b, s, mask.dims()[1]])
            .reshape([b * s, mask.dims()[1]]);

        // 3. Spatial pooling: [(B*S), C, E] → [(B*S), E] → [B, S, E]
        let x = self.spatial_pooling.forward(tokens, Some(mask_expanded));
        let x = x.reshape([b, s, e]);

        // 4. Positional encoding + LayerNorm
        let x = add_positional_encoding(x, &self.pe_buffer);
        let x = self.layer_norm.forward(x);

        // 5. Transformer encoder
        let mut x = x;
        for layer in &self.transformer_encoder {
            x = layer.forward(x, None);
        }

        // 6. Clone embeddings before temporal pooling
        let embeddings = x.clone();

        // 7. Temporal pooling: [B, S, E] → [B, E]
        let pooled = self.temporal_pooling.forward(x, None);

        (pooled, embeddings)
    }

    /// Generate embeddings only (no temporal pooling).
    ///
    /// Returns: [B, S, E] — per-patch embeddings after transformer encoder.
    pub fn embed(
        &self,
        x: Tensor<B, 3>,
        mask: Tensor<B, 2, Bool>,
    ) -> Tensor<B, 3> {
        let (_pooled, embeddings) = self.forward(x, mask);
        embeddings
    }
}
