wifi-densepose-sensing-server 0.3.2

//! Training loop with multi-term loss function for WiFi DensePose (ADR-023 Phase 4).
//!
//! 6-term composite loss, SGD with momentum, cosine annealing LR scheduler,
//! PCK/OKS validation metrics, numerical gradient estimation, and checkpointing.
//! All arithmetic uses f32. No external ML framework dependencies.

use crate::dataset;
use crate::embedding::{info_nce_loss, CsiAugmenter, ProjectionHead};
use crate::graph_transformer::{CsiToPoseTransformer, TransformerConfig};
use crate::sona::EwcRegularizer;
use std::path::Path;

/// Standard COCO keypoint sigmas for OKS (17 keypoints).
pub const COCO_KEYPOINT_SIGMAS: [f32; 17] = [
    0.026, 0.025, 0.025, 0.035, 0.035, 0.079, 0.079, 0.072, 0.072, 0.062, 0.062, 0.107, 0.107,
    0.087, 0.087, 0.089, 0.089,
];

/// Symmetric keypoint pairs (left, right) indices into 17-keypoint COCO layout.
const SYMMETRY_PAIRS: [(usize, usize); 5] = [(5, 6), (7, 8), (9, 10), (11, 12), (13, 14)];

/// Individual loss terms from the composite loss (6 supervised + 1 contrastive).
#[derive(Debug, Clone, Default)]
pub struct LossComponents {
    pub keypoint: f32,
    pub body_part: f32,
    pub uv: f32,
    pub temporal: f32,
    pub edge: f32,
    pub symmetry: f32,
    /// Contrastive loss (InfoNCE); only active during pretraining or when configured.
    pub contrastive: f32,
}

/// Per-term weights for the composite loss function.
#[derive(Debug, Clone)]
pub struct LossWeights {
    pub keypoint: f32,
    pub body_part: f32,
    pub uv: f32,
    pub temporal: f32,
    pub edge: f32,
    pub symmetry: f32,
    /// Contrastive loss weight (default 0.0; set >0 for joint training).
    pub contrastive: f32,
}

impl Default for LossWeights {
    fn default() -> Self {
        Self {
            keypoint: 1.0,
            body_part: 0.5,
            uv: 0.5,
            temporal: 0.1,
            edge: 0.2,
            symmetry: 0.1,
            contrastive: 0.0,
        }
    }
}

/// Mean squared error on keypoints (x, y, confidence).
pub fn keypoint_mse(pred: &[(f32, f32, f32)], target: &[(f32, f32, f32)]) -> f32 {
    if pred.is_empty() || target.is_empty() {
        return 0.0;
    }
    let n = pred.len().min(target.len());
    let sum: f32 = pred
        .iter()
        .zip(target.iter())
        .take(n)
        .map(|(p, t)| (p.0 - t.0).powi(2) + (p.1 - t.1).powi(2) + (p.2 - t.2).powi(2))
        .sum();
    sum / n as f32
}

/// Cross-entropy loss for body part classification.
/// `pred` = raw logits (length `n_samples * n_parts`), `target` = class indices.
pub fn body_part_cross_entropy(pred: &[f32], target: &[u8], n_parts: usize) -> f32 {
    if target.is_empty() || n_parts == 0 || pred.len() < n_parts {
        return 0.0;
    }
    let n_samples = target.len().min(pred.len() / n_parts);
    if n_samples == 0 {
        return 0.0;
    }
    let mut total = 0.0f32;
    for i in 0..n_samples {
        let logits = &pred[i * n_parts..(i + 1) * n_parts];
        let class = target[i] as usize;
        if class >= n_parts {
            continue;
        }
        let max_l = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
        let lse = logits.iter().map(|&l| (l - max_l).exp()).sum::<f32>().ln() + max_l;
        total += -logits[class] + lse;
    }
    total / n_samples as f32
}

/// L1 loss on UV coordinates.
pub fn uv_regression_loss(pu: &[f32], pv: &[f32], tu: &[f32], tv: &[f32]) -> f32 {
    let n = pu.len().min(pv.len()).min(tu.len()).min(tv.len());
    if n == 0 {
        return 0.0;
    }
    let s: f32 = (0..n)
        .map(|i| (pu[i] - tu[i]).abs() + (pv[i] - tv[i]).abs())
        .sum();
    s / n as f32
}

/// Temporal consistency loss: penalizes large frame-to-frame keypoint jumps.
pub fn temporal_consistency_loss(prev: &[(f32, f32, f32)], curr: &[(f32, f32, f32)]) -> f32 {
    let n = prev.len().min(curr.len());
    if n == 0 {
        return 0.0;
    }
    let s: f32 = prev
        .iter()
        .zip(curr.iter())
        .take(n)
        .map(|(p, c)| (c.0 - p.0).powi(2) + (c.1 - p.1).powi(2))
        .sum();
    s / n as f32
}

/// Graph edge loss: penalizes deviation of bone lengths from expected values.
pub fn graph_edge_loss(kp: &[(f32, f32, f32)], edges: &[(usize, usize)], expected: &[f32]) -> f32 {
    if edges.is_empty() || edges.len() != expected.len() {
        return 0.0;
    }
    let (mut sum, mut cnt) = (0.0f32, 0usize);
    for (i, &(a, b)) in edges.iter().enumerate() {
        if a >= kp.len() || b >= kp.len() {
            continue;
        }
        let d = ((kp[a].0 - kp[b].0).powi(2) + (kp[a].1 - kp[b].1).powi(2)).sqrt();
        sum += (d - expected[i]).powi(2);
        cnt += 1;
    }
    if cnt == 0 {
        0.0
    } else {
        sum / cnt as f32
    }
}

/// Symmetry loss: penalizes asymmetry between left-right limb pairs.
pub fn symmetry_loss(kp: &[(f32, f32, f32)]) -> f32 {
    if kp.len() < 15 {
        return 0.0;
    }
    let (mut sum, mut cnt) = (0.0f32, 0usize);
    for &(l, r) in &SYMMETRY_PAIRS {
        if l >= kp.len() || r >= kp.len() {
            continue;
        }
        let ld = ((kp[l].0 - kp[0].0).powi(2) + (kp[l].1 - kp[0].1).powi(2)).sqrt();
        let rd = ((kp[r].0 - kp[0].0).powi(2) + (kp[r].1 - kp[0].1).powi(2)).sqrt();
        sum += (ld - rd).powi(2);
        cnt += 1;
    }
    if cnt == 0 {
        0.0
    } else {
        sum / cnt as f32
    }
}

/// Weighted composite loss from individual components.
pub fn composite_loss(c: &LossComponents, w: &LossWeights) -> f32 {
    w.keypoint * c.keypoint
        + w.body_part * c.body_part
        + w.uv * c.uv
        + w.temporal * c.temporal
        + w.edge * c.edge
        + w.symmetry * c.symmetry
        + w.contrastive * c.contrastive
}

// ── Optimizer ──────────────────────────────────────────────────────────────

/// SGD optimizer with momentum and weight decay.
pub struct SgdOptimizer {
    lr: f32,
    momentum: f32,
    weight_decay: f32,
    velocity: Vec<f32>,
}

impl SgdOptimizer {
    pub fn new(lr: f32, momentum: f32, weight_decay: f32) -> Self {
        Self {
            lr,
            momentum,
            weight_decay,
            velocity: Vec::new(),
        }
    }

    /// v = mu*v + grad + wd*param; param -= lr*v
    pub fn step(&mut self, params: &mut [f32], gradients: &[f32]) {
        if self.velocity.len() != params.len() {
            self.velocity = vec![0.0; params.len()];
        }
        for i in 0..params.len().min(gradients.len()) {
            let g = gradients[i] + self.weight_decay * params[i];
            self.velocity[i] = self.momentum * self.velocity[i] + g;
            params[i] -= self.lr * self.velocity[i];
        }
    }

    pub fn set_lr(&mut self, lr: f32) {
        self.lr = lr;
    }
    pub fn state(&self) -> Vec<f32> {
        self.velocity.clone()
    }
    pub fn load_state(&mut self, state: Vec<f32>) {
        self.velocity = state;
    }
}

// ── Learning rate schedulers ───────────────────────────────────────────────

/// Cosine annealing: decays LR from initial to min over total_steps.
pub struct CosineScheduler {
    initial_lr: f32,
    min_lr: f32,
    total_steps: usize,
}

impl CosineScheduler {
    pub fn new(initial_lr: f32, min_lr: f32, total_steps: usize) -> Self {
        Self {
            initial_lr,
            min_lr,
            total_steps,
        }
    }
    pub fn get_lr(&self, step: usize) -> f32 {
        if self.total_steps == 0 {
            return self.initial_lr;
        }
        let p = step.min(self.total_steps) as f32 / self.total_steps as f32;
        self.min_lr
            + (self.initial_lr - self.min_lr) * (1.0 + (std::f32::consts::PI * p).cos()) / 2.0
    }
}

/// Warmup + cosine annealing: linear ramp 0->initial_lr then cosine decay.
pub struct WarmupCosineScheduler {
    warmup_steps: usize,
    initial_lr: f32,
    min_lr: f32,
    total_steps: usize,
}

impl WarmupCosineScheduler {
    pub fn new(warmup_steps: usize, initial_lr: f32, min_lr: f32, total_steps: usize) -> Self {
        Self {
            warmup_steps,
            initial_lr,
            min_lr,
            total_steps,
        }
    }
    pub fn get_lr(&self, step: usize) -> f32 {
        if step < self.warmup_steps {
            if self.warmup_steps == 0 {
                return self.initial_lr;
            }
            return self.initial_lr * (step as f32 / self.warmup_steps as f32);
        }
        let cs = self.total_steps.saturating_sub(self.warmup_steps);
        if cs == 0 {
            return self.min_lr;
        }
        let p = (step - self.warmup_steps).min(cs) as f32 / cs as f32;
        self.min_lr
            + (self.initial_lr - self.min_lr) * (1.0 + (std::f32::consts::PI * p).cos()) / 2.0
    }
}

// ── Validation metrics ─────────────────────────────────────────────────────

/// Percentage of Correct Keypoints at a distance threshold.
pub fn pck_at_threshold(pred: &[(f32, f32, f32)], target: &[(f32, f32, f32)], thr: f32) -> f32 {
    let n = pred.len().min(target.len());
    if n == 0 {
        return 0.0;
    }
    let (mut correct, mut total) = (0usize, 0usize);
    for i in 0..n {
        if target[i].2 <= 0.0 {
            continue;
        }
        total += 1;
        let d = ((pred[i].0 - target[i].0).powi(2) + (pred[i].1 - target[i].1).powi(2)).sqrt();
        if d <= thr {
            correct += 1;
        }
    }
    if total == 0 {
        0.0
    } else {
        correct as f32 / total as f32
    }
}

/// Object Keypoint Similarity for a single instance.
pub fn oks_single(
    pred: &[(f32, f32, f32)],
    target: &[(f32, f32, f32)],
    sigmas: &[f32],
    area: f32,
) -> f32 {
    let n = pred.len().min(target.len()).min(sigmas.len());
    if n == 0 || area <= 0.0 {
        return 0.0;
    }
    let (mut sum, mut vis) = (0.0f32, 0usize);
    for i in 0..n {
        if target[i].2 <= 0.0 {
            continue;
        }
        vis += 1;
        let dsq = (pred[i].0 - target[i].0).powi(2) + (pred[i].1 - target[i].1).powi(2);
        let var = 2.0 * sigmas[i] * sigmas[i] * area;
        if var > 0.0 {
            sum += (-dsq / (2.0 * var)).exp();
        }
    }
    if vis == 0 {
        0.0
    } else {
        sum / vis as f32
    }
}

/// Mean OKS over multiple predictions (simplified mAP).
pub fn oks_map(preds: &[Vec<(f32, f32, f32)>], targets: &[Vec<(f32, f32, f32)>]) -> f32 {
    let n = preds.len().min(targets.len());
    if n == 0 {
        return 0.0;
    }
    let s: f32 = preds
        .iter()
        .zip(targets.iter())
        .take(n)
        .map(|(p, t)| oks_single(p, t, &COCO_KEYPOINT_SIGMAS, 1.0))
        .sum();
    s / n as f32
}

// ── Gradient estimation ────────────────────────────────────────────────────

/// Central difference gradient: (f(x+eps) - f(x-eps)) / (2*eps).
pub fn estimate_gradient(f: impl Fn(&[f32]) -> f32, params: &[f32], eps: f32) -> Vec<f32> {
    let mut grad = vec![0.0f32; params.len()];
    let mut p_plus = params.to_vec();
    let mut p_minus = params.to_vec();
    for i in 0..params.len() {
        p_plus[i] = params[i] + eps;
        p_minus[i] = params[i] - eps;
        grad[i] = (f(&p_plus) - f(&p_minus)) / (2.0 * eps);
        p_plus[i] = params[i];
        p_minus[i] = params[i];
    }
    grad
}

/// Clip gradients by global L2 norm.
pub fn clip_gradients(gradients: &mut [f32], max_norm: f32) {
    let norm = gradients.iter().map(|g| g * g).sum::<f32>().sqrt();
    if norm > max_norm && norm > 0.0 {
        let s = max_norm / norm;
        gradients.iter_mut().for_each(|g| *g *= s);
    }
}

// ── Training sample ────────────────────────────────────────────────────────

/// A single training sample (defined locally, not dependent on dataset.rs).
#[derive(Debug, Clone)]
pub struct TrainingSample {
    pub csi_features: Vec<Vec<f32>>,
    pub target_keypoints: Vec<(f32, f32, f32)>,
    pub target_body_parts: Vec<u8>,
    pub target_uv: (Vec<f32>, Vec<f32>),
}

/// Convert a dataset::TrainingSample into a trainer::TrainingSample.
pub fn from_dataset_sample(ds: &dataset::TrainingSample) -> TrainingSample {
    let csi_features = ds.csi_window.clone();
    let target_keypoints: Vec<(f32, f32, f32)> = ds.pose_label.keypoints.to_vec();
    let target_body_parts: Vec<u8> = ds
        .pose_label
        .body_parts
        .iter()
        .map(|bp| bp.part_id)
        .collect();
    let (tu, tv) = if ds.pose_label.body_parts.is_empty() {
        (Vec::new(), Vec::new())
    } else {
        let u: Vec<f32> = ds
            .pose_label
            .body_parts
            .iter()
            .flat_map(|bp| bp.u_coords.iter().copied())
            .collect();
        let v: Vec<f32> = ds
            .pose_label
            .body_parts
            .iter()
            .flat_map(|bp| bp.v_coords.iter().copied())
            .collect();
        (u, v)
    };
    TrainingSample {
        csi_features,
        target_keypoints,
        target_body_parts,
        target_uv: (tu, tv),
    }
}

// ── Checkpoint ─────────────────────────────────────────────────────────────

/// Serializable version of EpochStats for checkpoint storage.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct EpochStatsSerializable {
    pub epoch: usize,
    pub train_loss: f32,
    pub val_loss: f32,
    pub pck_02: f32,
    pub oks_map: f32,
    pub lr: f32,
    pub loss_keypoint: f32,
    pub loss_body_part: f32,
    pub loss_uv: f32,
    pub loss_temporal: f32,
    pub loss_edge: f32,
    pub loss_symmetry: f32,
}

/// Serializable training checkpoint.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct Checkpoint {
    pub epoch: usize,
    pub params: Vec<f32>,
    pub optimizer_state: Vec<f32>,
    pub best_loss: f32,
    pub metrics: EpochStatsSerializable,
}

impl Checkpoint {
    pub fn save_to_file(&self, path: &Path) -> std::io::Result<()> {
        let json = serde_json::to_string_pretty(self).map_err(std::io::Error::other)?;
        std::fs::write(path, json)
    }
    pub fn load_from_file(path: &Path) -> std::io::Result<Self> {
        let json = std::fs::read_to_string(path)?;
        serde_json::from_str(&json).map_err(std::io::Error::other)
    }
}

/// Statistics for a single training epoch.
#[derive(Debug, Clone)]
pub struct EpochStats {
    pub epoch: usize,
    pub train_loss: f32,
    pub val_loss: f32,
    pub pck_02: f32,
    pub oks_map: f32,
    pub lr: f32,
    pub loss_components: LossComponents,
}

impl EpochStats {
    fn to_serializable(&self) -> EpochStatsSerializable {
        let c = &self.loss_components;
        EpochStatsSerializable {
            epoch: self.epoch,
            train_loss: self.train_loss,
            val_loss: self.val_loss,
            pck_02: self.pck_02,
            oks_map: self.oks_map,
            lr: self.lr,
            loss_keypoint: c.keypoint,
            loss_body_part: c.body_part,
            loss_uv: c.uv,
            loss_temporal: c.temporal,
            loss_edge: c.edge,
            loss_symmetry: c.symmetry,
        }
    }
}

/// Final result from a complete training run.
#[derive(Debug, Clone)]
pub struct TrainingResult {
    pub best_epoch: usize,
    pub best_pck: f32,
    pub best_oks: f32,
    pub history: Vec<EpochStats>,
    pub total_time_secs: f64,
}

/// Configuration for the training loop.
#[derive(Debug, Clone)]
pub struct TrainerConfig {
    pub epochs: usize,
    pub batch_size: usize,
    pub lr: f32,
    pub momentum: f32,
    pub weight_decay: f32,
    pub warmup_epochs: usize,
    pub min_lr: f32,
    pub early_stop_patience: usize,
    pub checkpoint_every: usize,
    pub loss_weights: LossWeights,
    /// Contrastive loss weight for joint supervised+contrastive training (default 0.0).
    pub contrastive_loss_weight: f32,
    /// Temperature for InfoNCE loss during pretraining (default 0.07).
    pub pretrain_temperature: f32,
}

impl Default for TrainerConfig {
    fn default() -> Self {
        Self {
            epochs: 100,
            batch_size: 32,
            lr: 0.01,
            momentum: 0.9,
            weight_decay: 1e-4,
            warmup_epochs: 5,
            min_lr: 1e-6,
            early_stop_patience: 10,
            checkpoint_every: 10,
            loss_weights: LossWeights::default(),
            contrastive_loss_weight: 0.0,
            pretrain_temperature: 0.07,
        }
    }
}

// ── Trainer ────────────────────────────────────────────────────────────────

/// Training loop orchestrator for WiFi DensePose pose estimation.
pub struct Trainer {
    config: TrainerConfig,
    optimizer: SgdOptimizer,
    scheduler: WarmupCosineScheduler,
    params: Vec<f32>,
    history: Vec<EpochStats>,
    best_val_loss: f32,
    best_epoch: usize,
    epochs_without_improvement: usize,
    /// Snapshot of params at the best validation loss epoch.
    best_params: Vec<f32>,
    /// When set, predict_keypoints delegates to the transformer's forward().
    transformer: Option<CsiToPoseTransformer>,
    /// Transformer config (needed for unflatten during gradient estimation).
    transformer_config: Option<TransformerConfig>,
    /// EWC++ regularizer for pretrain -> finetune transition.
    /// Prevents catastrophic forgetting of contrastive embedding structure.
    pub embedding_ewc: Option<EwcRegularizer>,
}

impl Trainer {
    pub fn new(config: TrainerConfig) -> Self {
        let optimizer = SgdOptimizer::new(config.lr, config.momentum, config.weight_decay);
        let scheduler = WarmupCosineScheduler::new(
            config.warmup_epochs,
            config.lr,
            config.min_lr,
            config.epochs,
        );
        let params: Vec<f32> = (0..64)
            .map(|i| (i as f32 * 0.7 + 0.3).sin() * 0.1)
            .collect();
        let best_params = params.clone();
        Self {
            config,
            optimizer,
            scheduler,
            params,
            history: Vec::new(),
            best_val_loss: f32::MAX,
            best_epoch: 0,
            epochs_without_improvement: 0,
            best_params,
            transformer: None,
            transformer_config: None,
            embedding_ewc: None,
        }
    }

    /// Create a trainer backed by the graph transformer. Gradient estimation
    /// uses central differences on the transformer's flattened weights.
    pub fn with_transformer(config: TrainerConfig, transformer: CsiToPoseTransformer) -> Self {
        let params = transformer.flatten_weights();
        let optimizer = SgdOptimizer::new(config.lr, config.momentum, config.weight_decay);
        let scheduler = WarmupCosineScheduler::new(
            config.warmup_epochs,
            config.lr,
            config.min_lr,
            config.epochs,
        );
        let tc = transformer.config().clone();
        let best_params = params.clone();
        Self {
            config,
            optimizer,
            scheduler,
            params,
            history: Vec::new(),
            best_val_loss: f32::MAX,
            best_epoch: 0,
            epochs_without_improvement: 0,
            best_params,
            transformer: Some(transformer),
            transformer_config: Some(tc),
            embedding_ewc: None,
        }
    }

    /// Access the transformer (if any).
    pub fn transformer(&self) -> Option<&CsiToPoseTransformer> {
        self.transformer.as_ref()
    }

    /// Get a mutable reference to the transformer.
    pub fn transformer_mut(&mut self) -> Option<&mut CsiToPoseTransformer> {
        self.transformer.as_mut()
    }

    /// Return current flattened params (transformer or simple).
    pub fn params(&self) -> &[f32] {
        &self.params
    }

    pub fn train_epoch(&mut self, samples: &[TrainingSample]) -> EpochStats {
        let epoch = self.history.len();
        let lr = self.scheduler.get_lr(epoch);
        self.optimizer.set_lr(lr);

        let mut acc = LossComponents::default();
        let bs = self.config.batch_size.max(1);
        let nb = samples.len().div_ceil(bs);
        let tc = self.transformer_config.clone();

        for bi in 0..nb {
            let batch = &samples[bi * bs..(bi * bs + bs).min(samples.len())];
            let snap = self.params.clone();
            let w = self.config.loss_weights.clone();
            let loss_fn = |p: &[f32]| match &tc {
                Some(tconf) => Self::batch_loss_with_transformer(p, batch, &w, tconf),
                None => Self::batch_loss(p, batch, &w),
            };
            let mut grad = estimate_gradient(loss_fn, &snap, 1e-4);
            clip_gradients(&mut grad, 1.0);
            self.optimizer.step(&mut self.params, &grad);

            let c = Self::batch_loss_components_impl(&self.params, batch, tc.as_ref());
            acc.keypoint += c.keypoint;
            acc.body_part += c.body_part;
            acc.uv += c.uv;
            acc.temporal += c.temporal;
            acc.edge += c.edge;
            acc.symmetry += c.symmetry;
        }

        if nb > 0 {
            let inv = 1.0 / nb as f32;
            acc.keypoint *= inv;
            acc.body_part *= inv;
            acc.uv *= inv;
            acc.temporal *= inv;
            acc.edge *= inv;
            acc.symmetry *= inv;
        }

        let train_loss = composite_loss(&acc, &self.config.loss_weights);
        let (pck, oks) = self.evaluate_metrics(samples);
        let stats = EpochStats {
            epoch,
            train_loss,
            val_loss: train_loss,
            pck_02: pck,
            oks_map: oks,
            lr,
            loss_components: acc,
        };
        self.history.push(stats.clone());
        stats
    }

    pub fn should_stop(&self) -> bool {
        self.epochs_without_improvement >= self.config.early_stop_patience
    }

    pub fn best_metrics(&self) -> Option<&EpochStats> {
        self.history.get(self.best_epoch)
    }

    pub fn run_training(
        &mut self,
        train: &[TrainingSample],
        val: &[TrainingSample],
    ) -> TrainingResult {
        let start = std::time::Instant::now();
        for _ in 0..self.config.epochs {
            let mut stats = self.train_epoch(train);
            let tc = self.transformer_config.clone();
            let val_loss = if !val.is_empty() {
                let c = Self::batch_loss_components_impl(&self.params, val, tc.as_ref());
                composite_loss(&c, &self.config.loss_weights)
            } else {
                stats.train_loss
            };
            stats.val_loss = val_loss;
            if !val.is_empty() {
                let (pck, oks) = self.evaluate_metrics(val);
                stats.pck_02 = pck;
                stats.oks_map = oks;
            }
            if let Some(last) = self.history.last_mut() {
                last.val_loss = stats.val_loss;
                last.pck_02 = stats.pck_02;
                last.oks_map = stats.oks_map;
            }
            if val_loss < self.best_val_loss {
                self.best_val_loss = val_loss;
                self.best_epoch = stats.epoch;
                self.best_params = self.params.clone();
                self.epochs_without_improvement = 0;
            } else {
                self.epochs_without_improvement += 1;
            }
            if self.should_stop() {
                break;
            }
        }
        // Restore best-epoch params for checkpoint and downstream use
        self.params = self.best_params.clone();
        let best = self.best_metrics().cloned().unwrap_or(EpochStats {
            epoch: 0,
            train_loss: f32::MAX,
            val_loss: f32::MAX,
            pck_02: 0.0,
            oks_map: 0.0,
            lr: self.config.lr,
            loss_components: LossComponents::default(),
        });
        TrainingResult {
            best_epoch: best.epoch,
            best_pck: best.pck_02,
            best_oks: best.oks_map,
            history: self.history.clone(),
            total_time_secs: start.elapsed().as_secs_f64(),
        }
    }

    /// Run one self-supervised pretraining epoch using SimCLR objective.
    /// Does NOT require pose labels -- only CSI windows.
    ///
    /// For each mini-batch:
    /// 1. Generate augmented pair (view_a, view_b) for each window
    /// 2. Forward each view through transformer to get body_part_features
    /// 3. Mean-pool to get frame embedding
    /// 4. Project through ProjectionHead
    /// 5. Compute InfoNCE loss
    /// 6. Estimate gradients via central differences and SGD update
    ///
    /// Returns mean epoch loss.
    pub fn pretrain_epoch(
        &mut self,
        csi_windows: &[Vec<Vec<f32>>],
        augmenter: &CsiAugmenter,
        projection: &mut ProjectionHead,
        temperature: f32,
        epoch: usize,
    ) -> f32 {
        if csi_windows.is_empty() {
            return 0.0;
        }
        let lr = self.scheduler.get_lr(epoch);
        self.optimizer.set_lr(lr);

        let bs = self.config.batch_size.max(1);
        let nb = csi_windows.len().div_ceil(bs);
        let mut total_loss = 0.0f32;

        let tc = self.transformer_config.clone();
        let tc_ref = match &tc {
            Some(c) => c,
            None => return 0.0, // pretraining requires a transformer
        };

        for bi in 0..nb {
            let start = bi * bs;
            let end = (start + bs).min(csi_windows.len());
            let batch = &csi_windows[start..end];

            // Generate augmented pairs and compute embeddings + loss
            let snap = self.params.clone();
            let mut proj_flat = Vec::new();
            projection.flatten_into(&mut proj_flat);

            // Combined params: transformer + projection head
            let mut combined = snap.clone();
            combined.extend_from_slice(&proj_flat);

            let t_param_count = snap.len();
            let p_config = projection.config.clone();
            let tc_c = tc_ref.clone();
            let temp = temperature;

            // Build augmented views for the batch
            let seed_base = (epoch * 10000 + bi) as u64;
            let aug_pairs: Vec<_> = batch
                .iter()
                .enumerate()
                .map(|(k, w)| augmenter.augment_pair(w, seed_base + k as u64))
                .collect();

            // Loss function over combined (transformer + projection) params
            let batch_owned: Vec<Vec<Vec<f32>>> = batch.to_vec();
            let loss_fn = |params: &[f32]| -> f32 {
                let t_params = &params[..t_param_count];
                let p_params = &params[t_param_count..];
                let mut t = CsiToPoseTransformer::zeros(tc_c.clone());
                if t.unflatten_weights(t_params).is_err() {
                    return f32::MAX;
                }
                let (proj, _) = ProjectionHead::unflatten_from(p_params, &p_config);
                let d = p_config.d_model;

                let mut embs_a = Vec::with_capacity(batch_owned.len());
                let mut embs_b = Vec::with_capacity(batch_owned.len());

                for (k, _w) in batch_owned.iter().enumerate() {
                    let (ref va, ref vb) = aug_pairs[k];
                    // Mean-pool body features for view A
                    let feats_a = t.embed(va);
                    let mut pooled_a = vec![0.0f32; d];
                    for f in &feats_a {
                        for (p, &v) in pooled_a.iter_mut().zip(f.iter()) {
                            *p += v;
                        }
                    }
                    let n = feats_a.len() as f32;
                    if n > 0.0 {
                        for p in pooled_a.iter_mut() {
                            *p /= n;
                        }
                    }
                    embs_a.push(proj.forward(&pooled_a));

                    // Mean-pool body features for view B
                    let feats_b = t.embed(vb);
                    let mut pooled_b = vec![0.0f32; d];
                    for f in &feats_b {
                        for (p, &v) in pooled_b.iter_mut().zip(f.iter()) {
                            *p += v;
                        }
                    }
                    let n = feats_b.len() as f32;
                    if n > 0.0 {
                        for p in pooled_b.iter_mut() {
                            *p /= n;
                        }
                    }
                    embs_b.push(proj.forward(&pooled_b));
                }

                info_nce_loss(&embs_a, &embs_b, temp)
            };

            let batch_loss = loss_fn(&combined);
            total_loss += batch_loss;

            // Estimate gradient via central differences on combined params
            let mut grad = estimate_gradient(loss_fn, &combined, 1e-4);
            clip_gradients(&mut grad, 1.0);

            // Update transformer params
            self.optimizer
                .step(&mut self.params, &grad[..t_param_count]);

            // Update projection head params
            let mut proj_params = proj_flat.clone();
            // Simple SGD for projection head
            for i in 0..proj_params.len().min(grad.len() - t_param_count) {
                proj_params[i] -= lr * grad[t_param_count + i];
            }
            let (new_proj, _) = ProjectionHead::unflatten_from(&proj_params, &projection.config);
            *projection = new_proj;
        }

        total_loss / nb as f32
    }

    pub fn checkpoint(&self) -> Checkpoint {
        let m =
            self.history
                .last()
                .map(|s| s.to_serializable())
                .unwrap_or(EpochStatsSerializable {
                    epoch: 0,
                    train_loss: 0.0,
                    val_loss: 0.0,
                    pck_02: 0.0,
                    oks_map: 0.0,
                    lr: self.config.lr,
                    loss_keypoint: 0.0,
                    loss_body_part: 0.0,
                    loss_uv: 0.0,
                    loss_temporal: 0.0,
                    loss_edge: 0.0,
                    loss_symmetry: 0.0,
                });
        Checkpoint {
            epoch: self.history.len(),
            params: self.params.clone(),
            optimizer_state: self.optimizer.state(),
            best_loss: self.best_val_loss,
            metrics: m,
        }
    }

    fn batch_loss(params: &[f32], batch: &[TrainingSample], w: &LossWeights) -> f32 {
        composite_loss(&Self::batch_loss_components_impl(params, batch, None), w)
    }

    fn batch_loss_with_transformer(
        params: &[f32],
        batch: &[TrainingSample],
        w: &LossWeights,
        tc: &TransformerConfig,
    ) -> f32 {
        composite_loss(
            &Self::batch_loss_components_impl(params, batch, Some(tc)),
            w,
        )
    }

    fn batch_loss_components(params: &[f32], batch: &[TrainingSample]) -> LossComponents {
        Self::batch_loss_components_impl(params, batch, None)
    }

    fn batch_loss_components_impl(
        params: &[f32],
        batch: &[TrainingSample],
        tc: Option<&TransformerConfig>,
    ) -> LossComponents {
        if batch.is_empty() {
            return LossComponents::default();
        }
        let mut acc = LossComponents::default();
        let mut prev_kp: Option<Vec<(f32, f32, f32)>> = None;
        for sample in batch {
            let pred_kp = match tc {
                Some(tconf) => Self::predict_keypoints_transformer(params, sample, tconf),
                None => Self::predict_keypoints(params, sample),
            };
            acc.keypoint += keypoint_mse(&pred_kp, &sample.target_keypoints);
            let n_parts = 24usize;
            let logits: Vec<f32> = sample
                .target_body_parts
                .iter()
                .flat_map(|_| {
                    (0..n_parts)
                        .map(|j| {
                            if j < params.len() {
                                params[j] * 0.1
                            } else {
                                0.0
                            }
                        })
                        .collect::<Vec<f32>>()
                })
                .collect();
            acc.body_part += body_part_cross_entropy(&logits, &sample.target_body_parts, n_parts);
            let (ref tu, ref tv) = sample.target_uv;
            let pu: Vec<f32> = tu
                .iter()
                .enumerate()
                .map(|(i, &u)| {
                    u + if i < params.len() {
                        params[i] * 0.01
                    } else {
                        0.0
                    }
                })
                .collect();
            let pv: Vec<f32> = tv
                .iter()
                .enumerate()
                .map(|(i, &v)| {
                    v + if i < params.len() {
                        params[i] * 0.01
                    } else {
                        0.0
                    }
                })
                .collect();
            acc.uv += uv_regression_loss(&pu, &pv, tu, tv);
            if let Some(ref prev) = prev_kp {
                acc.temporal += temporal_consistency_loss(prev, &pred_kp);
            }
            acc.symmetry += symmetry_loss(&pred_kp);
            prev_kp = Some(pred_kp);
        }
        let inv = 1.0 / batch.len() as f32;
        acc.keypoint *= inv;
        acc.body_part *= inv;
        acc.uv *= inv;
        acc.temporal *= inv;
        acc.symmetry *= inv;
        acc
    }

    fn predict_keypoints(params: &[f32], sample: &TrainingSample) -> Vec<(f32, f32, f32)> {
        let n_kp = sample.target_keypoints.len().max(17);
        let feats: Vec<f32> = sample
            .csi_features
            .iter()
            .flat_map(|v| v.iter().copied())
            .collect();
        (0..n_kp)
            .map(|k| {
                let base = k * 3;
                let (mut x, mut y) = (0.0f32, 0.0f32);
                for (i, &f) in feats.iter().take(params.len()).enumerate() {
                    let pi = (base + i) % params.len();
                    x += f * params[pi] * 0.01;
                    y += f * params[(pi + 1) % params.len()] * 0.01;
                }
                if base < params.len() {
                    x += params[base % params.len()];
                    y += params[(base + 1) % params.len()];
                }
                let c = if base + 2 < params.len() {
                    params[(base + 2) % params.len()].clamp(0.0, 1.0)
                } else {
                    0.5
                };
                (x, y, c)
            })
            .collect()
    }

    /// Predict keypoints using the graph transformer. Uses zero-init
    /// constructor (fast) then overwrites all weights from params.
    fn predict_keypoints_transformer(
        params: &[f32],
        sample: &TrainingSample,
        tc: &TransformerConfig,
    ) -> Vec<(f32, f32, f32)> {
        let mut t = CsiToPoseTransformer::zeros(tc.clone());
        if t.unflatten_weights(params).is_err() {
            return Self::predict_keypoints(params, sample);
        }
        let output = t.forward(&sample.csi_features);
        output.keypoints
    }

    fn evaluate_metrics(&self, samples: &[TrainingSample]) -> (f32, f32) {
        if samples.is_empty() {
            return (0.0, 0.0);
        }
        let preds: Vec<Vec<_>> = samples
            .iter()
            .map(|s| match &self.transformer_config {
                Some(tc) => Self::predict_keypoints_transformer(&self.params, s, tc),
                None => Self::predict_keypoints(&self.params, s),
            })
            .collect();
        let targets: Vec<Vec<_>> = samples.iter().map(|s| s.target_keypoints.clone()).collect();
        let pck = preds
            .iter()
            .zip(targets.iter())
            .map(|(p, t)| pck_at_threshold(p, t, 0.2))
            .sum::<f32>()
            / samples.len() as f32;
        (pck, oks_map(&preds, &targets))
    }

    /// Sync the internal transformer's weights from the flat params after training.
    pub fn sync_transformer_weights(&mut self) {
        if let Some(ref mut t) = self.transformer {
            let _ = t.unflatten_weights(&self.params);
        }
    }

    /// Consolidate pretrained parameters using EWC++ before fine-tuning.
    ///
    /// Call this after pretraining completes (e.g., after `pretrain_epoch` loops).
    /// It computes the Fisher Information diagonal on the current params using
    /// the contrastive loss as the objective, then sets the current params as the
    /// EWC reference point. During subsequent supervised training, the EWC penalty
    /// will discourage large deviations from the pretrained structure.
    pub fn consolidate_pretrained(&mut self) {
        let mut ewc = EwcRegularizer::new(5000.0, 0.99);
        let current_params = self.params.clone();

        // Compute Fisher diagonal using a simple loss based on parameter deviation.
        // In a real scenario this would use the contrastive loss over training data;
        // here we use a squared-magnitude proxy that penalises changes to each param.
        let fisher = EwcRegularizer::compute_fisher(
            &current_params,
            |p: &[f32]| p.iter().map(|&x| x * x).sum::<f32>(),
            1,
        );
        ewc.update_fisher(&fisher);
        ewc.consolidate(&current_params);
        self.embedding_ewc = Some(ewc);
    }

    /// Return the EWC penalty for the current parameters (0.0 if no EWC is set).
    pub fn ewc_penalty(&self) -> f32 {
        match &self.embedding_ewc {
            Some(ewc) => ewc.penalty(&self.params),
            None => 0.0,
        }
    }

    /// Return the EWC penalty gradient for the current parameters.
    pub fn ewc_penalty_gradient(&self) -> Vec<f32> {
        match &self.embedding_ewc {
            Some(ewc) => ewc.penalty_gradient(&self.params),
            None => vec![0.0f32; self.params.len()],
        }
    }
}

// ── Tests ──────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;

    fn mkp(off: f32) -> Vec<(f32, f32, f32)> {
        (0..17)
            .map(|i| (i as f32 + off, i as f32 * 2.0 + off, 1.0))
            .collect()
    }

    fn symmetric_pose() -> Vec<(f32, f32, f32)> {
        let mut kp = vec![(0.0f32, 0.0f32, 1.0f32); 17];
        kp[0] = (5.0, 5.0, 1.0);
        for &(l, r) in &SYMMETRY_PAIRS {
            kp[l] = (3.0, 5.0, 1.0);
            kp[r] = (7.0, 5.0, 1.0);
        }
        kp
    }

    fn sample() -> TrainingSample {
        TrainingSample {
            csi_features: vec![vec![1.0; 8]; 4],
            target_keypoints: mkp(0.0),
            target_body_parts: vec![0, 1, 2, 3],
            target_uv: (vec![0.5; 4], vec![0.5; 4]),
        }
    }

    #[test]
    fn keypoint_mse_zero_for_identical() {
        assert_eq!(keypoint_mse(&mkp(0.0), &mkp(0.0)), 0.0);
    }
    #[test]
    fn keypoint_mse_positive_for_different() {
        assert!(keypoint_mse(&mkp(0.0), &mkp(1.0)) > 0.0);
    }
    #[test]
    fn keypoint_mse_symmetric() {
        let (ab, ba) = (
            keypoint_mse(&mkp(0.0), &mkp(1.0)),
            keypoint_mse(&mkp(1.0), &mkp(0.0)),
        );
        assert!((ab - ba).abs() < 1e-6, "{ab} vs {ba}");
    }
    #[test]
    fn temporal_consistency_zero_for_static() {
        assert_eq!(temporal_consistency_loss(&mkp(0.0), &mkp(0.0)), 0.0);
    }
    #[test]
    fn temporal_consistency_positive_for_motion() {
        assert!(temporal_consistency_loss(&mkp(0.0), &mkp(1.0)) > 0.0);
    }
    #[test]
    fn symmetry_loss_zero_for_symmetric_pose() {
        assert!(symmetry_loss(&symmetric_pose()) < 1e-6);
    }
    #[test]
    fn graph_edge_loss_zero_when_correct() {
        let kp = vec![(0.0, 0.0, 1.0), (3.0, 4.0, 1.0), (6.0, 0.0, 1.0)];
        assert!(graph_edge_loss(&kp, &[(0, 1), (1, 2)], &[5.0, 5.0]) < 1e-6);
    }
    #[test]
    fn composite_loss_respects_weights() {
        let c = LossComponents {
            keypoint: 1.0,
            body_part: 1.0,
            uv: 1.0,
            temporal: 1.0,
            edge: 1.0,
            symmetry: 1.0,
            contrastive: 0.0,
        };
        let w1 = LossWeights {
            keypoint: 1.0,
            body_part: 0.0,
            uv: 0.0,
            temporal: 0.0,
            edge: 0.0,
            symmetry: 0.0,
            contrastive: 0.0,
        };
        let w2 = LossWeights {
            keypoint: 2.0,
            body_part: 0.0,
            uv: 0.0,
            temporal: 0.0,
            edge: 0.0,
            symmetry: 0.0,
            contrastive: 0.0,
        };
        assert!((composite_loss(&c, &w2) - 2.0 * composite_loss(&c, &w1)).abs() < 1e-6);
        let wz = LossWeights {
            keypoint: 0.0,
            body_part: 0.0,
            uv: 0.0,
            temporal: 0.0,
            edge: 0.0,
            symmetry: 0.0,
            contrastive: 0.0,
        };
        assert_eq!(composite_loss(&c, &wz), 0.0);
    }
    #[test]
    fn cosine_scheduler_starts_at_initial() {
        assert!((CosineScheduler::new(0.01, 0.0001, 100).get_lr(0) - 0.01).abs() < 1e-6);
    }
    #[test]
    fn cosine_scheduler_ends_at_min() {
        assert!((CosineScheduler::new(0.01, 0.0001, 100).get_lr(100) - 0.0001).abs() < 1e-6);
    }
    #[test]
    fn cosine_scheduler_midpoint() {
        assert!((CosineScheduler::new(0.01, 0.0, 100).get_lr(50) - 0.005).abs() < 1e-4);
    }
    #[test]
    fn warmup_starts_at_zero() {
        assert!(WarmupCosineScheduler::new(10, 0.01, 0.0001, 100).get_lr(0) < 1e-6);
    }
    #[test]
    fn warmup_reaches_initial_at_warmup_end() {
        assert!((WarmupCosineScheduler::new(10, 0.01, 0.0001, 100).get_lr(10) - 0.01).abs() < 1e-6);
    }
    #[test]
    fn pck_perfect_prediction_is_1() {
        assert!((pck_at_threshold(&mkp(0.0), &mkp(0.0), 0.2) - 1.0).abs() < 1e-6);
    }
    #[test]
    fn pck_all_wrong_is_0() {
        assert!(pck_at_threshold(&mkp(0.0), &mkp(100.0), 0.2) < 1e-6);
    }
    #[test]
    fn oks_perfect_is_1() {
        assert!((oks_single(&mkp(0.0), &mkp(0.0), &COCO_KEYPOINT_SIGMAS, 1.0) - 1.0).abs() < 1e-6);
    }
    #[test]
    fn sgd_step_reduces_simple_loss() {
        let mut p = vec![5.0f32];
        let mut opt = SgdOptimizer::new(0.1, 0.0, 0.0);
        let init = p[0] * p[0];
        for _ in 0..10 {
            let grad = vec![2.0 * p[0]];
            opt.step(&mut p, &grad);
        }
        assert!(p[0] * p[0] < init);
    }
    #[test]
    fn gradient_clipping_respects_max_norm() {
        let mut g = vec![3.0, 4.0];
        clip_gradients(&mut g, 2.5);
        assert!((g.iter().map(|x| x * x).sum::<f32>().sqrt() - 2.5).abs() < 1e-4);
    }
    #[test]
    fn early_stopping_triggers() {
        let cfg = TrainerConfig {
            epochs: 100,
            early_stop_patience: 3,
            ..Default::default()
        };
        let mut t = Trainer::new(cfg);
        let s = vec![sample()];
        t.best_val_loss = -1.0;
        let mut stopped = false;
        for _ in 0..20 {
            t.train_epoch(&s);
            t.epochs_without_improvement += 1;
            if t.should_stop() {
                stopped = true;
                break;
            }
        }
        assert!(stopped);
    }
    #[test]
    fn checkpoint_round_trip() {
        let mut t = Trainer::new(TrainerConfig::default());
        t.train_epoch(&[sample()]);
        let ckpt = t.checkpoint();
        let dir = std::env::temp_dir().join("trainer_ckpt_test");
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("ckpt.json");
        ckpt.save_to_file(&path).unwrap();
        let loaded = Checkpoint::load_from_file(&path).unwrap();
        assert_eq!(loaded.epoch, ckpt.epoch);
        assert_eq!(loaded.params.len(), ckpt.params.len());
        assert!((loaded.best_loss - ckpt.best_loss).abs() < 1e-6);
        let _ = std::fs::remove_file(&path);
        let _ = std::fs::remove_dir(&dir);
    }

    // ── Integration tests: transformer + trainer pipeline ──────────

    #[test]
    fn dataset_to_trainer_conversion() {
        let ds = crate::dataset::TrainingSample {
            csi_window: vec![vec![1.0; 8]; 4],
            pose_label: crate::dataset::PoseLabel {
                keypoints: {
                    let mut kp = [(0.0f32, 0.0f32, 1.0f32); 17];
                    for (i, k) in kp.iter_mut().enumerate() {
                        k.0 = i as f32;
                        k.1 = i as f32 * 2.0;
                    }
                    kp
                },
                body_parts: Vec::new(),
                confidence: 1.0,
            },
            source: "test",
        };
        let ts = from_dataset_sample(&ds);
        assert_eq!(ts.csi_features.len(), 4);
        assert_eq!(ts.csi_features[0].len(), 8);
        assert_eq!(ts.target_keypoints.len(), 17);
        assert!((ts.target_keypoints[0].0 - 0.0).abs() < 1e-6);
        assert!((ts.target_keypoints[1].0 - 1.0).abs() < 1e-6);
        assert!(ts.target_body_parts.is_empty()); // no body parts in source
    }

    #[test]
    fn trainer_with_transformer_runs_epoch() {
        use crate::graph_transformer::{CsiToPoseTransformer, TransformerConfig};
        let tf_config = TransformerConfig {
            n_subcarriers: 8,
            n_keypoints: 17,
            d_model: 8,
            n_heads: 2,
            n_gnn_layers: 1,
        };
        let transformer = CsiToPoseTransformer::new(tf_config);
        let config = TrainerConfig {
            epochs: 2,
            batch_size: 4,
            lr: 0.001,
            warmup_epochs: 0,
            early_stop_patience: 100,
            ..Default::default()
        };
        let mut t = Trainer::with_transformer(config, transformer);

        // The params should be the transformer's flattened weights
        assert!(
            t.params().len() > 100,
            "transformer should have many params"
        );

        // Create samples matching the transformer's n_subcarriers=8
        let samples: Vec<TrainingSample> = (0..8)
            .map(|i| TrainingSample {
                csi_features: vec![vec![(i as f32 * 0.1).sin(); 8]; 4],
                target_keypoints: (0..17)
                    .map(|k| (k as f32 * 0.5, k as f32 * 0.3, 1.0))
                    .collect(),
                target_body_parts: vec![0, 1, 2],
                target_uv: (vec![0.5; 3], vec![0.5; 3]),
            })
            .collect();

        let stats = t.train_epoch(&samples);
        assert!(stats.train_loss.is_finite(), "loss should be finite");
    }

    #[test]
    fn trainer_with_transformer_loss_finite_after_training() {
        use crate::graph_transformer::{CsiToPoseTransformer, TransformerConfig};
        let tf_config = TransformerConfig {
            n_subcarriers: 8,
            n_keypoints: 17,
            d_model: 8,
            n_heads: 2,
            n_gnn_layers: 1,
        };
        let transformer = CsiToPoseTransformer::new(tf_config);
        let config = TrainerConfig {
            epochs: 3,
            batch_size: 4,
            lr: 0.0001,
            warmup_epochs: 0,
            early_stop_patience: 100,
            ..Default::default()
        };
        let mut t = Trainer::with_transformer(config, transformer);

        let samples: Vec<TrainingSample> = (0..4)
            .map(|i| TrainingSample {
                csi_features: vec![vec![(i as f32 * 0.2).sin(); 8]; 4],
                target_keypoints: (0..17)
                    .map(|k| (k as f32 * 0.5, k as f32 * 0.3, 1.0))
                    .collect(),
                target_body_parts: vec![],
                target_uv: (vec![], vec![]),
            })
            .collect();

        let result = t.run_training(&samples, &[]);
        assert!(
            result.history.iter().all(|s| s.train_loss.is_finite()),
            "all losses should be finite"
        );

        // Sync weights back and verify transformer still works
        t.sync_transformer_weights();
        if let Some(tf) = t.transformer() {
            let out = tf.forward(&vec![vec![1.0; 8]; 4]);
            assert_eq!(out.keypoints.len(), 17);
            for (i, &(x, y, z)) in out.keypoints.iter().enumerate() {
                assert!(
                    x.is_finite() && y.is_finite() && z.is_finite(),
                    "kp {i} not finite after training"
                );
            }
        }
    }

    #[test]
    fn test_pretrain_epoch_loss_decreases() {
        use crate::embedding::{CsiAugmenter, EmbeddingConfig, ProjectionHead};
        use crate::graph_transformer::{CsiToPoseTransformer, TransformerConfig};

        let tf_config = TransformerConfig {
            n_subcarriers: 8,
            n_keypoints: 17,
            d_model: 8,
            n_heads: 2,
            n_gnn_layers: 1,
        };
        let transformer = CsiToPoseTransformer::new(tf_config);
        let config = TrainerConfig {
            epochs: 10,
            batch_size: 4,
            lr: 0.001,
            warmup_epochs: 0,
            early_stop_patience: 100,
            pretrain_temperature: 0.5,
            ..Default::default()
        };
        let mut trainer = Trainer::with_transformer(config, transformer);

        let e_config = EmbeddingConfig {
            d_model: 8,
            d_proj: 16,
            temperature: 0.5,
            normalize: true,
        };
        let mut projection = ProjectionHead::new(e_config);
        let augmenter = CsiAugmenter::new();

        // Synthetic CSI windows (8 windows, each 4 frames of 8 subcarriers)
        let csi_windows: Vec<Vec<Vec<f32>>> = (0..8)
            .map(|i| {
                (0..4)
                    .map(|a| {
                        (0..8)
                            .map(|s| ((i * 7 + a * 3 + s) as f32 * 0.41).sin() * 0.5)
                            .collect()
                    })
                    .collect()
            })
            .collect();

        let loss_0 = trainer.pretrain_epoch(&csi_windows, &augmenter, &mut projection, 0.5, 0);
        let loss_1 = trainer.pretrain_epoch(&csi_windows, &augmenter, &mut projection, 0.5, 1);
        let loss_2 = trainer.pretrain_epoch(&csi_windows, &augmenter, &mut projection, 0.5, 2);

        assert!(
            loss_0.is_finite(),
            "epoch 0 loss should be finite: {loss_0}"
        );
        assert!(
            loss_1.is_finite(),
            "epoch 1 loss should be finite: {loss_1}"
        );
        assert!(
            loss_2.is_finite(),
            "epoch 2 loss should be finite: {loss_2}"
        );
        // Loss should generally decrease (or at least the final loss should be less than initial)
        assert!(
            loss_2 <= loss_0 + 0.5,
            "loss should not increase drastically: epoch0={loss_0}, epoch2={loss_2}"
        );
    }

    #[test]
    fn test_contrastive_loss_weight_in_composite() {
        let c = LossComponents {
            keypoint: 0.0,
            body_part: 0.0,
            uv: 0.0,
            temporal: 0.0,
            edge: 0.0,
            symmetry: 0.0,
            contrastive: 1.0,
        };
        let w = LossWeights {
            keypoint: 0.0,
            body_part: 0.0,
            uv: 0.0,
            temporal: 0.0,
            edge: 0.0,
            symmetry: 0.0,
            contrastive: 0.5,
        };
        assert!((composite_loss(&c, &w) - 0.5).abs() < 1e-6);
    }

    // ── Phase 7: EWC++ in Trainer tests ───────────────────────────────

    #[test]
    fn test_ewc_consolidation_reduces_forgetting() {
        // Setup: create trainer, set params, consolidate, then train.
        // EWC penalty should resist large param changes.
        let config = TrainerConfig {
            epochs: 5,
            batch_size: 4,
            lr: 0.01,
            warmup_epochs: 0,
            early_stop_patience: 100,
            ..Default::default()
        };
        let mut trainer = Trainer::new(config);
        let _pretrained_params = trainer.params().to_vec();

        // Consolidate pretrained state
        trainer.consolidate_pretrained();
        assert!(
            trainer.embedding_ewc.is_some(),
            "EWC should be set after consolidation"
        );

        // Train a few epochs (params will change)
        let samples = vec![sample()];
        for _ in 0..3 {
            trainer.train_epoch(&samples);
        }

        // With EWC penalty active, params should still be somewhat close
        // to pretrained values (EWC resists change)
        let penalty = trainer.ewc_penalty();
        assert!(
            penalty > 0.0,
            "EWC penalty should be > 0 after params changed"
        );

        // The penalty gradient should push params back toward pretrained values
        let grad = trainer.ewc_penalty_gradient();
        let any_nonzero = grad.iter().any(|&g| g.abs() > 1e-10);
        assert!(any_nonzero, "EWC gradient should have non-zero components");
    }

    #[test]
    fn test_ewc_penalty_nonzero_after_consolidation() {
        let config = TrainerConfig::default();
        let mut trainer = Trainer::new(config);

        // Before consolidation, penalty should be 0
        assert!(
            (trainer.ewc_penalty()).abs() < 1e-10,
            "no EWC => zero penalty"
        );

        // Consolidate
        trainer.consolidate_pretrained();

        // At the reference point, penalty = 0
        assert!(
            trainer.ewc_penalty().abs() < 1e-6,
            "penalty should be ~0 at reference point"
        );

        // Perturb params away from reference
        for p in trainer.params.iter_mut() {
            *p += 0.1;
        }

        let penalty = trainer.ewc_penalty();
        assert!(
            penalty > 0.0,
            "penalty should be > 0 after deviating from reference, got {penalty}"
        );
    }
}