/// SleepFM encoder — produce PSG embeddings.
///
/// The SetTransformer produces:
/// - Pooled embedding: [B, E] (temporal-pooled single vector)
/// - Per-patch embeddings: [B, S, E] (one per time patch)

use std::path::Path;
use std::time::Instant;

use anyhow::Context;
use burn::prelude::*;

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

/// Per-modality embedding produced by SleepFM.
pub struct ModalityEmbedding {
    /// Pooled output: [E] — single embedding vector.
    pub pooled: Vec<f32>,
    /// Per-patch outputs: [S, E] — temporal embeddings.
    pub patches: Vec<f32>,
    /// Shape of patches: [S, E]
    pub patches_shape: Vec<usize>,
    /// Embedding dimension.
    pub embed_dim: usize,
}

/// Collection of per-modality embeddings.
pub struct EncodingResult {
    pub embeddings: Vec<ModalityEmbedding>,
    pub ms_load: f64,
    pub ms_encode: f64,
}

// ── SleepFmEncoder ────────────────────────────────────────────────────────────

/// SleepFM encoder for PSG signal processing.
pub struct SleepFmEncoder<B: Backend> {
    model: SetTransformer<B>,
    pub model_cfg: ModelConfig,
    device: B::Device,
}

impl<B: Backend> SleepFmEncoder<B> {
    /// Load model from config.json and safetensors weights.
    pub fn load(
        config_path: &Path,
        weights_path: &Path,
        device: B::Device,
    ) -> anyhow::Result<(Self, f64)> {
        let cfg_str = std::fs::read_to_string(config_path)
            .with_context(|| format!("config: {}", config_path.display()))?;
        let model_cfg: ModelConfig = serde_json::from_str(&cfg_str)
            .context("parsing model config")?;

        let t = Instant::now();
        let model = load_model::<B>(
            &model_cfg,
            weights_path.to_str().context("weights path not valid UTF-8")?,
            &device,
        )?;
        let ms = t.elapsed().as_secs_f64() * 1000.0;

        Ok((Self { model, model_cfg, device }, ms))
    }

    pub fn describe(&self) -> String {
        let c = &self.model_cfg;
        format!(
            "SleepFM  embed_dim={}  heads={}  layers={}  pooling_head={}  patch_size={}",
            c.embed_dim, c.num_heads, c.num_layers, c.pooling_head, c.patch_size,
        )
    }

    /// Run inference on a single modality.
    ///
    /// signal: [1, C, T] — multichannel signal for one modality
    /// mask: [1, C] — channel mask (true = padding)
    pub fn run(
        &self,
        signal: Tensor<B, 3>,
        mask: Tensor<B, 2, Bool>,
    ) -> anyhow::Result<ModalityEmbedding> {
        let (pooled, patches) = self.model.forward(signal, mask);

        let embed_dim = self.model_cfg.embed_dim;
        let patches_shape = patches.dims().to_vec();

        let pooled_vec = pooled.squeeze::<1>()
            .into_data()
            .to_vec::<f32>()
            .map_err(|e| anyhow::anyhow!("pooled→vec: {e:?}"))?;

        let patches_vec = patches.squeeze::<2>()
            .into_data()
            .to_vec::<f32>()
            .map_err(|e| anyhow::anyhow!("patches→vec: {e:?}"))?;

        Ok(ModalityEmbedding {
            pooled: pooled_vec,
            patches: patches_vec,
            patches_shape: patches_shape[1..].to_vec(),
            embed_dim,
        })
    }

    pub fn device(&self) -> &B::Device {
        &self.device
    }
}
