/// Load pretrained SleepFM weights from a safetensors file.
///
/// PyTorch state_dict keys have `module.` prefix (from DataParallel).
/// Sequential indices for the Tokenizer:
///   Stage 0: 0(Conv), 1(BN), 3(LN)
///   Stage 1: 4(Conv), 5(BN), 7(LN)
///   Stage 2: 8(Conv), 9(BN), 11(LN)
///   Stage 3: 12(Conv), 13(BN), 15(LN)
///   Stage 4: 16(Conv), 17(BN), 19(LN)
///   Stage 5: 20(Conv), 21(BN), 23(LN)
///   Final linear: 26
///
/// TransformerEncoderLayer (fused MHA):
///   self_attn.in_proj_weight [3*D, D], in_proj_bias [3*D]
///   self_attn.out_proj.weight [D, D], out_proj.bias [D]
///   linear1.weight, linear1.bias
///   linear2.weight, linear2.bias
///   norm1.weight, norm1.bias
///   norm2.weight, norm2.bias

use std::collections::HashMap;
use burn::prelude::*;
use half::bf16;
use safetensors::SafeTensors;

use crate::config::ModelConfig;
use crate::model::set_transformer::SetTransformer;

// ── WeightMap ─────────────────────────────────────────────────────────────────

pub struct WeightMap {
    pub tensors: HashMap<String, (Vec<f32>, Vec<usize>)>,
}

impl WeightMap {
    pub fn from_file(path: &str) -> anyhow::Result<Self> {
        let bytes = std::fs::read(path)?;
        let st = SafeTensors::deserialize(&bytes)?;
        let mut tensors = HashMap::with_capacity(st.len());

        for (raw_key, view) in st.tensors() {
            // Strip `module.` prefix from DataParallel
            let key = raw_key
                .strip_prefix("module.")
                .unwrap_or(raw_key.as_str())
                .to_string();

            let shape: Vec<usize> = view.shape().to_vec();
            let data = view.data();

            let f32s: Vec<f32> = match view.dtype() {
                safetensors::Dtype::BF16 => data
                    .chunks_exact(2)
                    .map(|b| bf16::from_le_bytes([b[0], b[1]]).to_f32())
                    .collect(),
                safetensors::Dtype::F32 => data
                    .chunks_exact(4)
                    .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
                    .collect(),
                safetensors::Dtype::I64 => {
                    // Skip non-float tensors like num_batches_tracked
                    continue;
                }
                other => anyhow::bail!("unsupported dtype {:?} for key {key}", other),
            };

            tensors.insert(key, (f32s, shape));
        }

        Ok(Self { tensors })
    }

    pub fn take<B: Backend, const N: usize>(
        &mut self,
        key: &str,
        device: &B::Device,
    ) -> anyhow::Result<Tensor<B, N>> {
        let (data, shape) = self.tensors.remove(key)
            .ok_or_else(|| anyhow::anyhow!("weight key not found: {key}"))?;
        if shape.len() != N {
            anyhow::bail!("rank mismatch for {key}: expected {N}, got {}", shape.len());
        }
        Ok(Tensor::<B, N>::from_data(TensorData::new(data, shape), device))
    }

    pub fn has(&self, key: &str) -> bool {
        self.tensors.contains_key(key)
    }

    pub fn print_keys(&self) {
        let mut keys: Vec<&str> = self.tensors.keys().map(String::as_str).collect();
        keys.sort();
        for k in keys {
            let (_, s) = &self.tensors[k];
            println!("  {k:80}  {s:?}");
        }
    }
}

// ── Weight setters ────────────────────────────────────────────────────────────

fn set_linear_wb<B: Backend>(linear: &mut burn::nn::Linear<B>, w: Tensor<B, 2>, b: Tensor<B, 1>) {
    linear.weight = linear.weight.clone().map(|_| w.transpose());
    if let Some(ref bias) = linear.bias {
        linear.bias = Some(bias.clone().map(|_| b));
    }
}

fn set_layernorm<B: Backend>(ln: &mut burn::nn::LayerNorm<B>, w: Tensor<B, 1>, b: Tensor<B, 1>) {
    ln.gamma = ln.gamma.clone().map(|_| w);
    if let Some(ref beta) = ln.beta {
        ln.beta = Some(beta.clone().map(|_| b));
    }
}

fn set_layernorm2d<B: Backend>(
    ln: &mut crate::model::layernorm2d::LayerNorm2d<B>,
    w: Tensor<B, 2>,
    b: Tensor<B, 2>,
) {
    ln.gamma = ln.gamma.clone().map(|_| w);
    ln.beta = ln.beta.clone().map(|_| b);
}

fn set_conv1d_wb<B: Backend>(conv: &mut burn::nn::conv::Conv1d<B>, w: Tensor<B, 3>, b: Tensor<B, 1>) {
    conv.weight = conv.weight.clone().map(|_| w);
    if let Some(ref bias) = conv.bias {
        conv.bias = Some(bias.clone().map(|_| b));
    }
}

fn set_batchnorm_params<B: Backend>(
    bn: &mut burn::nn::BatchNorm<B>,
    gamma: Tensor<B, 1>,
    beta: Tensor<B, 1>,
    running_mean: Tensor<B, 1>,
    running_var: Tensor<B, 1>,
) {
    bn.gamma = bn.gamma.clone().map(|_| gamma);
    bn.beta = bn.beta.clone().map(|_| beta);
    bn.running_mean = burn::module::RunningState::new(running_mean);
    bn.running_var = burn::module::RunningState::new(running_var);
}

// ── Fused MHA loader ──────────────────────────────────────────────────────────

fn load_fused_mha<B: Backend>(
    wm: &mut WeightMap,
    mha: &mut crate::model::transformer_layer::FusedMultiheadAttention<B>,
    prefix: &str,
    device: &B::Device,
) -> anyhow::Result<()> {
    if let (Ok(w), Ok(b)) = (
        wm.take::<B, 2>(&format!("{prefix}.in_proj_weight"), device),
        wm.take::<B, 1>(&format!("{prefix}.in_proj_bias"), device),
    ) {
        set_linear_wb(&mut mha.in_proj, w, b);
    }
    if let (Ok(w), Ok(b)) = (
        wm.take::<B, 2>(&format!("{prefix}.out_proj.weight"), device),
        wm.take::<B, 1>(&format!("{prefix}.out_proj.bias"), device),
    ) {
        set_linear_wb(&mut mha.out_proj, w, b);
    }
    Ok(())
}

fn load_transformer_encoder_layer<B: Backend>(
    wm: &mut WeightMap,
    layer: &mut crate::model::transformer_layer::TransformerEncoderLayer<B>,
    prefix: &str,
    device: &B::Device,
) -> anyhow::Result<()> {
    // self_attn
    load_fused_mha(wm, &mut layer.self_attn, &format!("{prefix}.self_attn"), device)?;
    // norm1
    if let (Ok(w), Ok(b)) = (
        wm.take::<B, 1>(&format!("{prefix}.norm1.weight"), device),
        wm.take::<B, 1>(&format!("{prefix}.norm1.bias"), device),
    ) {
        set_layernorm(&mut layer.norm1, w, b);
    }
    // norm2
    if let (Ok(w), Ok(b)) = (
        wm.take::<B, 1>(&format!("{prefix}.norm2.weight"), device),
        wm.take::<B, 1>(&format!("{prefix}.norm2.bias"), device),
    ) {
        set_layernorm(&mut layer.norm2, w, b);
    }
    // linear1
    if let (Ok(w), Ok(b)) = (
        wm.take::<B, 2>(&format!("{prefix}.linear1.weight"), device),
        wm.take::<B, 1>(&format!("{prefix}.linear1.bias"), device),
    ) {
        set_linear_wb(&mut layer.linear1, w, b);
    }
    // linear2
    if let (Ok(w), Ok(b)) = (
        wm.take::<B, 2>(&format!("{prefix}.linear2.weight"), device),
        wm.take::<B, 1>(&format!("{prefix}.linear2.bias"), device),
    ) {
        set_linear_wb(&mut layer.linear2, w, b);
    }
    Ok(())
}

// ── Full model loader ─────────────────────────────────────────────────────────

/// Load the SetTransformer model from a safetensors file.
pub fn load_model<B: Backend>(
    cfg: &ModelConfig,
    weights_path: &str,
    device: &B::Device,
) -> anyhow::Result<SetTransformer<B>> {
    let mut wm = WeightMap::from_file(weights_path)?;
    eprintln!("Loading {} weight tensors...", wm.tensors.len());
    load_model_from_wm(cfg, &mut wm, device)
}

pub fn load_model_from_wm<B: Backend>(
    cfg: &ModelConfig,
    wm: &mut WeightMap,
    device: &B::Device,
) -> anyhow::Result<SetTransformer<B>> {
    let mut model = SetTransformer::new(
        cfg.patch_size,
        cfg.embed_dim,
        cfg.num_heads,
        cfg.num_layers,
        cfg.pooling_head,
        cfg.max_seq_length,
        device,
    );

    load_weights(wm, &mut model, device)?;
    Ok(model)
}

fn load_weights<B: Backend>(
    wm: &mut WeightMap,
    model: &mut SetTransformer<B>,
    device: &B::Device,
) -> anyhow::Result<()> {
    // ── Tokenizer stages ────────────────────────────────────────────────────
    // Sequential indices: Conv, BN, ELU(skip), LN
    let conv_indices = [0, 4, 8, 12, 16, 20];
    let bn_indices   = [1, 5, 9, 13, 17, 21];
    let ln_indices   = [3, 7, 11, 15, 19, 23];

    for (stage_idx, ((conv_i, bn_i), ln_i)) in
        conv_indices.iter().zip(bn_indices.iter()).zip(ln_indices.iter()).enumerate()
    {
        let stage = &mut model.patch_embedding.stages[stage_idx];

        // Conv1d
        if let (Ok(w), Ok(b)) = (
            wm.take::<B, 3>(&format!("patch_embedding.tokenizer.{conv_i}.weight"), device),
            wm.take::<B, 1>(&format!("patch_embedding.tokenizer.{conv_i}.bias"), device),
        ) {
            set_conv1d_wb(&mut stage.conv, w, b);
        }

        // BatchNorm
        if let (Ok(gamma), Ok(beta), Ok(rm), Ok(rv)) = (
            wm.take::<B, 1>(&format!("patch_embedding.tokenizer.{bn_i}.weight"), device),
            wm.take::<B, 1>(&format!("patch_embedding.tokenizer.{bn_i}.bias"), device),
            wm.take::<B, 1>(&format!("patch_embedding.tokenizer.{bn_i}.running_mean"), device),
            wm.take::<B, 1>(&format!("patch_embedding.tokenizer.{bn_i}.running_var"), device),
        ) {
            set_batchnorm_params(&mut stage.bn, gamma, beta, rm, rv);
        }

        // LayerNorm2d — PyTorch nn.LayerNorm([C, T]) with weight/bias shape [C, T]
        if let (Ok(w), Ok(b)) = (
            wm.take::<B, 2>(&format!("patch_embedding.tokenizer.{ln_i}.weight"), device),
            wm.take::<B, 2>(&format!("patch_embedding.tokenizer.{ln_i}.bias"), device),
        ) {
            set_layernorm2d(&mut stage.ln, w, b);
        }
    }

    // Tokenizer final linear (index 26)
    if let (Ok(w), Ok(b)) = (
        wm.take::<B, 2>("patch_embedding.tokenizer.26.weight", device),
        wm.take::<B, 1>("patch_embedding.tokenizer.26.bias", device),
    ) {
        set_linear_wb(&mut model.patch_embedding.linear, w, b);
    }

    // ── Spatial pooling (AttentionPooling with single TransformerEncoderLayer) ──
    load_transformer_encoder_layer(
        wm,
        &mut model.spatial_pooling.transformer_layer,
        "spatial_pooling.transformer_layer",
        device,
    )?;

    // ── Positional encoding (loaded as buffer, not parameter) ───────────────
    if let Ok(pe) = wm.take::<B, 3>("positional_encoding.pe", device) {
        model.pe_buffer = pe;
    }

    // ── LayerNorm ───────────────────────────────────────────────────────────
    if let (Ok(w), Ok(b)) = (
        wm.take::<B, 1>("layer_norm.weight", device),
        wm.take::<B, 1>("layer_norm.bias", device),
    ) {
        set_layernorm(&mut model.layer_norm, w, b);
    }

    // ── Transformer encoder layers ──────────────────────────────────────────
    for i in 0..model.transformer_encoder.len() {
        let prefix = format!("transformer_encoder.layers.{i}");
        load_transformer_encoder_layer(
            wm,
            &mut model.transformer_encoder[i],
            &prefix,
            device,
        )?;
    }

    // ── Temporal pooling ────────────────────────────────────────────────────
    load_transformer_encoder_layer(
        wm,
        &mut model.temporal_pooling.transformer_layer,
        "temporal_pooling.transformer_layer",
        device,
    )?;

    // Report unused keys (ignore num_batches_tracked — integer tensors skipped during load)
    let unused: Vec<&str> = wm.tensors.keys()
        .map(String::as_str)
        .filter(|k| !k.ends_with("num_batches_tracked"))
        .collect();
    if !unused.is_empty() {
        eprintln!("Warning: {} unused weight keys:", unused.len());
        for k in &unused {
            eprintln!("  {k}");
        }
    }

    Ok(())
}
