/// Tokenizer — 6-stage 1D CNN that converts raw signal patches into embeddings.
///
/// Python architecture (nn.Sequential):
///   Stage i (i=0..5):
///     Conv1d(in_ch, out_ch, kernel=5, stride=2, padding=2)
///     BatchNorm1d(out_ch)
///     ELU()
///     LayerNorm([out_ch, input_size / 2^(i+1)])
///
///   Then: AdaptiveAvgPool1d(1), Flatten(), Linear(128, output_size)
///
/// Channel progression: 1 → 4 → 8 → 16 → 32 → 64 → 128
/// Spatial progression: input_size → /2 → /4 → /8 → /16 → /32 → /64
///
/// Weight key mapping from PyTorch state_dict (Sequential indices):
///   Stage 0: tokenizer.0 (Conv), tokenizer.1 (BN), tokenizer.3 (LN)
///   Stage 1: tokenizer.4 (Conv), tokenizer.5 (BN), tokenizer.7 (LN)
///   Stage 2: tokenizer.8 (Conv), tokenizer.9 (BN), tokenizer.11 (LN)
///   Stage 3: tokenizer.12 (Conv), tokenizer.13 (BN), tokenizer.15 (LN)
///   Stage 4: tokenizer.16 (Conv), tokenizer.17 (BN), tokenizer.19 (LN)
///   Stage 5: tokenizer.20 (Conv), tokenizer.21 (BN), tokenizer.23 (LN)
///   Final:   tokenizer.26 (Linear)

use burn::prelude::*;
use burn::nn::{
    Linear, LinearConfig,
    conv::{Conv1d, Conv1dConfig},
    BatchNorm, BatchNormConfig,
};

use crate::model::layernorm2d::LayerNorm2d;

/// One stage of the tokenizer CNN.
#[derive(Module, Debug)]
pub struct TokenizerStage<B: Backend> {
    pub conv: Conv1d<B>,
    pub bn: BatchNorm<B>,
    pub ln: LayerNorm2d<B>,
}

impl<B: Backend> TokenizerStage<B> {
    pub fn new(
        in_channels: usize,
        out_channels: usize,
        ln_spatial: usize,  // spatial dim after conv
        device: &B::Device,
    ) -> Self {
        let conv = Conv1dConfig::new(in_channels, out_channels, 5)
            .with_stride(2)
            .with_padding(burn::nn::PaddingConfig1d::Explicit(2))
            .with_bias(true)
            .init(device);
        let bn = BatchNormConfig::new(out_channels).init(device);
        // LayerNorm over [out_channels, ln_spatial] — normalises jointly over last 2 dims
        let ln = LayerNorm2d::new(out_channels, ln_spatial, 1e-5, device);
        Self { conv, bn, ln }
    }

    /// Forward: [N, C_in, T] → [N, C_out, T/2]
    pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
        let x = self.conv.forward(x);
        let x = self.bn.forward(x);
        let x = elu(x);
        self.ln.forward(x)
    }
}

/// ELU activation: max(0,x) + min(0, exp(x)-1)
fn elu<B: Backend>(x: Tensor<B, 3>) -> Tensor<B, 3> {
    let zeros = Tensor::zeros_like(&x);
    let positive = x.clone().max_pair(zeros.clone());
    let negative = x.min_pair(zeros).exp().sub_scalar(1.0);
    positive + negative
}

/// Full Tokenizer: 6 Conv stages + AdaptiveAvgPool1d + Linear
#[derive(Module, Debug)]
pub struct Tokenizer<B: Backend> {
    pub stages: Vec<TokenizerStage<B>>,
    pub linear: Linear<B>,
    pub input_size: usize,
    pub output_size: usize,
}

impl<B: Backend> Tokenizer<B> {
    pub fn new(input_size: usize, output_size: usize, device: &B::Device) -> Self {
        let channels = [1, 4, 8, 16, 32, 64, 128];
        let mut stages = Vec::with_capacity(6);
        let mut spatial = input_size;
        for i in 0..6 {
            spatial /= 2;
            stages.push(TokenizerStage::new(
                channels[i], channels[i + 1], spatial, device,
            ));
        }
        let linear = LinearConfig::new(128, output_size).with_bias(true).init(device);
        Self { stages, linear, input_size, output_size }
    }

    /// Forward pass.
    ///
    /// Input:  [B, C, T]  where T = num_channels * num_patches * input_size
    /// Output: [B, C, num_patches, output_size]
    ///
    /// Python logic:
    ///   x: [B, C, T] → view [B, C, num_patches, input_size]
    ///   reshape to [B*C*num_patches, 1, input_size]
    ///   run through CNN stages
    ///   adaptive_avg_pool1d(1) → flatten → linear
    ///   reshape to [B, C, num_patches, output_size]
    pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 4> {
        let [b, c, t] = x.dims();
        let num_patches = t / self.input_size;

        // [B, C, T] → [B*C*num_patches, 1, input_size]
        let x = x.reshape([b * c * num_patches, 1, self.input_size]);

        // Run through 6 CNN stages
        let mut x = x;
        for stage in &self.stages {
            x = stage.forward(x);
        }

        // AdaptiveAvgPool1d(1): [N, 128, spatial] → mean over spatial → [N, 128]
        let [n, ch, sp] = x.dims();
        let x = x.sum_dim(2).reshape([n, ch]) / (sp as f64);

        // Linear: [N, 128] → [N, output_size]
        let x = self.linear.forward(x);

        // Reshape to [B, C, num_patches, output_size]
        x.reshape([b, c, num_patches, self.output_size])
    }
}
