//! V-JEPA: Joint Embedding Predictive Architecture for Video
//!
//! Extends I-JEPA (jepa.rs) to video:
//! - 3D patch embeddings (temporal × spatial)
//! - Tube masking: random space-time tubes covering 90% of volume
//! - Temporal context masking strategy
//! - Neuro-JEPA extension: apply same paradigm to EEG/MEG signal segments

use super::jepa::{
    BlockMaskConfig, EmaTargetEncoder, FeatureBank, JepaEvaluator, JepaLossResult, JepaPredictor,
    jepa_loss,
};

// ============================================================================
// Video Patch Configuration
// ============================================================================

/// Configuration for 3D (video) patch embeddings
#[derive(Debug, Clone)]
pub struct VideoPatchConfig {
    /// Spatial patch size in pixels (e.g. 16)
    pub spatial_patch_size: usize,
    /// Temporal patch size in frames (e.g. 2)
    pub temporal_patch_size: usize,
    /// Spatial resolution (height = width assumed square)
    pub image_size: usize,
    /// Number of input frames
    pub num_frames: usize,
    /// Embedding dimension
    pub embed_dim: usize,
}

impl Default for VideoPatchConfig {
    fn default() -> Self {
        Self {
            spatial_patch_size: 16,
            temporal_patch_size: 2,
            image_size: 224,
            num_frames: 16,
            embed_dim: 1024,
        }
    }
}

impl VideoPatchConfig {
    /// Spatial grid size (patches per dimension)
    pub fn spatial_grid(&self) -> usize {
        self.image_size / self.spatial_patch_size
    }

    /// Number of temporal segments
    pub fn temporal_segments(&self) -> usize {
        self.num_frames / self.temporal_patch_size
    }

    /// Total 3D patches: T × H × W
    pub fn total_patches(&self) -> usize {
        self.temporal_segments() * self.spatial_grid() * self.spatial_grid()
    }

    /// Convert flat patch index to (t, h, w) tuple
    pub fn patch_idx_to_thw(&self, idx: usize) -> (usize, usize, usize) {
        let sg = self.spatial_grid();
        let spatial = sg * sg;
        let t = idx / spatial;
        let h = (idx % spatial) / sg;
        let w = idx % sg;
        (t, h, w)
    }

    /// Convert (t, h, w) to flat patch index
    pub fn thw_to_patch_idx(&self, t: usize, h: usize, w: usize) -> usize {
        let sg = self.spatial_grid();
        t * sg * sg + h * sg + w
    }
}

// ============================================================================
// 3D Patch Embedder (CPU reference)
// ============================================================================

/// 3D patch embedder: video tensor → sequence of patch embeddings
/// Input: [T, H, W, C] video frames (frame-major)
/// Output: [T/Tp, H/Sp, W/Sp, embed_dim] patch sequence
#[derive(Debug, Clone)]
pub struct PatchEmbed3D {
    pub config: VideoPatchConfig,
    /// Projection weights [C * Tp * Sp * Sp, embed_dim]
    proj_w: Vec<f32>,
    proj_b: Vec<f32>,
    /// Position embeddings [total_patches, embed_dim]
    pos_embed: Vec<f32>,
}

impl PatchEmbed3D {
    pub fn new(config: VideoPatchConfig, channels: usize) -> Self {
        let tp = config.temporal_patch_size;
        let sp = config.spatial_patch_size;
        let d = config.embed_dim;
        let patch_vol = channels * tp * sp * sp;
        let n = config.total_patches();

        let init = |size: usize, s: u64| -> Vec<f32> {
            let mut v = vec![0.0f32; size];
            let mut lcg = s.wrapping_mul(2654435761).wrapping_add(1);
            let scale = (2.0f32 / size as f32).sqrt();
            for x in v.iter_mut() {
                lcg = lcg
                    .wrapping_mul(6364136223846793005)
                    .wrapping_add(1442695040888963407);
                let u = ((lcg >> 11) as f32) / (1u64 << 53) as f32;
                *x = (u * 2.0 - 1.0) * scale;
            }
            v
        };

        Self {
            config,
            proj_w: init(patch_vol * d, 77),
            proj_b: vec![0.0; d],
            pos_embed: init(n * d, 78),
        }
    }

    /// Embed a video frame sequence.
    ///
    /// `video`: flat [num_frames * H * W * C] in frame-major order
    /// Returns: [total_patches, embed_dim]
    pub fn embed(&self, video: &[f32], channels: usize) -> Vec<f32> {
        let cfg = &self.config;
        let ts = cfg.temporal_segments();
        let sg = cfg.spatial_grid();
        let tp = cfg.temporal_patch_size;
        let sp = cfg.spatial_patch_size;
        let h = cfg.image_size;
        let w = cfg.image_size;
        let d = cfg.embed_dim;
        let patch_vol = channels * tp * sp * sp;
        let n = cfg.total_patches();
        let mut out = vec![0.0f32; n * d];

        for ti in 0..ts {
            for hi in 0..sg {
                for wi in 0..sg {
                    let patch_idx = ti * sg * sg + hi * sg + wi;
                    // Extract patch volume: tp frames × sp × sp × C
                    let mut patch = vec![0.0f32; patch_vol];
                    let mut pi = 0;
                    for tf in 0..tp {
                        let frame = ti * tp + tf;
                        for ph in 0..sp {
                            let row = hi * sp + ph;
                            for pw in 0..sp {
                                let col = wi * sp + pw;
                                for c in 0..channels {
                                    let flat_idx = frame * h * w * channels
                                        + row * w * channels
                                        + col * channels
                                        + c;
                                    if flat_idx < video.len() {
                                        patch[pi] = video[flat_idx];
                                    }
                                    pi += 1;
                                }
                            }
                        }
                    }
                    // Linear projection
                    for j in 0..d {
                        let mut val = self.proj_b[j];
                        for k in 0..patch_vol {
                            val += patch[k] * self.proj_w[k * d + j];
                        }
                        // Add position embedding
                        val += self.pos_embed[patch_idx * d + j];
                        out[patch_idx * d + j] = val;
                    }
                }
            }
        }
        out
    }
}

// ============================================================================
// Tube Masking (V-JEPA)
// ============================================================================

/// V-JEPA tube masking: mask spatially-consistent tubes across time.
///
/// For each selected spatial position, ALL temporal frames at that
/// spatial location are masked. This prevents the model from inferring
/// masked patches from adjacent frames at the same spatial location.
#[derive(Debug, Clone)]
pub struct TubeMaskConfig {
    /// Fraction of spatial positions to mask (V-JEPA uses ~0.90)
    pub mask_ratio: f64,
    /// Number of target tubes to sample (analogous to target blocks)
    pub num_target_tubes: usize,
    /// Spatial scale range for each tube block
    pub spatial_scale_range: (f64, f64),
}

impl Default for TubeMaskConfig {
    fn default() -> Self {
        Self {
            mask_ratio: 0.90,
            num_target_tubes: 8,
            spatial_scale_range: (0.15, 0.20),
        }
    }
}

/// Result of tube masking for a video
#[derive(Debug, Clone)]
pub struct TubeMaskResult {
    /// Flat patch indices visible to context encoder [t, h, w] → flat
    pub context_indices: Vec<usize>,
    /// Flat patch indices for each target tube
    pub target_tubes: Vec<Vec<usize>>,
    /// Union of all target indices
    pub all_target_indices: Vec<usize>,
    /// Config reference
    pub video_config: VideoPatchConfig,
}

impl TubeMaskResult {
    pub fn num_context(&self) -> usize {
        self.context_indices.len()
    }
    pub fn num_target_patches(&self) -> usize {
        self.all_target_indices.len()
    }
    pub fn num_target_tubes(&self) -> usize {
        self.target_tubes.len()
    }
}

/// Video tube masking strategy
#[derive(Debug, Clone)]
pub struct TubeMaskStrategy {
    config: TubeMaskConfig,
    lcg: u64,
}

impl TubeMaskStrategy {
    pub fn new(config: TubeMaskConfig) -> Self {
        Self {
            config,
            lcg: 9876543210987654321,
        }
    }

    fn lcg_next(&mut self) -> u64 {
        self.lcg = self
            .lcg
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        self.lcg
    }
    fn rand_f64(&mut self) -> f64 {
        (self.lcg_next() >> 11) as f64 / (1u64 << 53) as f64
    }
    fn rand_range(&mut self, lo: f64, hi: f64) -> f64 {
        lo + self.rand_f64() * (hi - lo)
    }
    fn rand_usize(&mut self, n: usize) -> usize {
        (self.lcg_next() % n as u64) as usize
    }

    /// Sample a spatial block (set of spatial positions) and extend across all T
    fn sample_tube(&mut self, video_cfg: &VideoPatchConfig) -> Vec<usize> {
        let sg = video_cfg.spatial_grid();
        let ts = video_cfg.temporal_segments();
        let total_spatial = sg * sg;

        let scale = self.rand_range(
            self.config.spatial_scale_range.0,
            self.config.spatial_scale_range.1,
        );
        let area = (scale * total_spatial as f64).ceil() as usize;
        let block_h = (area as f64).sqrt().ceil() as usize;
        let block_h = block_h.clamp(1, sg);
        let block_w = (area as f64 / block_h as f64).ceil() as usize;
        let block_w = block_w.clamp(1, sg);

        let top = self.rand_usize((sg - block_h + 1).max(1));
        let left = self.rand_usize((sg - block_w + 1).max(1));

        let mut indices = Vec::with_capacity(ts * block_h * block_w);
        for t in 0..ts {
            for h in top..top + block_h {
                for w in left..left + block_w {
                    indices.push(video_cfg.thw_to_patch_idx(t, h.min(sg - 1), w.min(sg - 1)));
                }
            }
        }
        indices.sort_unstable();
        indices.dedup();
        indices
    }

    /// Generate a tube mask for one video
    pub fn generate(&mut self, video_cfg: &VideoPatchConfig, seed_offset: u64) -> TubeMaskResult {
        self.lcg = self.lcg.wrapping_add(seed_offset.wrapping_mul(1234567891));
        let total_patches = video_cfg.total_patches();
        let cfg = &self.config.clone();

        let mut target_tubes: Vec<Vec<usize>> = Vec::with_capacity(cfg.num_target_tubes);
        let mut all_target_set: std::collections::BTreeSet<usize> = Default::default();

        for _ in 0..cfg.num_target_tubes {
            let tube = self.sample_tube(video_cfg);
            for &idx in &tube {
                all_target_set.insert(idx);
            }
            target_tubes.push(tube);
        }
        let all_target_indices: Vec<usize> = all_target_set.into_iter().collect();

        // Context = non-target, subsampled to (1 - mask_ratio) fraction
        let mut context_cands: Vec<usize> = (0..total_patches)
            .filter(|p| !all_target_indices.contains(p))
            .collect();
        let keep = ((1.0 - cfg.mask_ratio) * context_cands.len() as f64).ceil() as usize;
        let keep = keep.max(1).min(context_cands.len());
        for i in 0..keep {
            let j = i + self.rand_usize(context_cands.len() - i);
            context_cands.swap(i, j);
        }
        let mut context_indices = context_cands[..keep].to_vec();
        context_indices.sort_unstable();

        TubeMaskResult {
            context_indices,
            target_tubes,
            all_target_indices,
            video_config: video_cfg.clone(),
        }
    }

    /// Batch generation
    pub fn generate_batch(
        &mut self,
        batch_size: usize,
        video_cfg: &VideoPatchConfig,
    ) -> Vec<TubeMaskResult> {
        (0..batch_size)
            .map(|i| self.generate(video_cfg, i as u64))
            .collect()
    }
}

// ============================================================================
// V-JEPA Configuration and Trainer
// ============================================================================

/// Full V-JEPA configuration
#[derive(Debug, Clone)]
pub struct VJepaConfig {
    /// Video patch configuration
    pub video: VideoPatchConfig,
    /// Predictor depth (narrow transformer, same structure as I-JEPA)
    pub predictor_depth: usize,
    /// EMA tau schedule
    pub ema_tau_start: f64,
    pub ema_tau_end: f64,
    /// Masking strategy
    pub mask: TubeMaskConfig,
}

impl Default for VJepaConfig {
    fn default() -> Self {
        Self {
            video: VideoPatchConfig::default(),
            predictor_depth: 6,
            ema_tau_start: 0.996,
            ema_tau_end: 1.0,
            mask: TubeMaskConfig::default(),
        }
    }
}

impl VJepaConfig {
    /// Small config for testing
    pub fn small_test() -> Self {
        Self {
            video: VideoPatchConfig {
                spatial_patch_size: 8,
                temporal_patch_size: 2,
                image_size: 32,
                num_frames: 4,
                embed_dim: 64,
            },
            predictor_depth: 2,
            ema_tau_start: 0.996,
            ema_tau_end: 1.0,
            mask: TubeMaskConfig {
                mask_ratio: 0.75,
                num_target_tubes: 2,
                spatial_scale_range: (0.15, 0.30),
            },
        }
    }
}

/// Minimal mock video encoder (CPU reference for testing)
#[derive(Debug, Clone)]
pub struct MockVideoEncoder {
    pub embed_dim: usize,
    pub total_patches: usize,
    weights: Vec<f32>,
}

impl MockVideoEncoder {
    pub fn new(embed_dim: usize, total_patches: usize) -> Self {
        let mut lcg = 54321u64;
        let weights: Vec<f32> = (0..total_patches * embed_dim)
            .map(|_| {
                lcg = lcg
                    .wrapping_mul(6364136223846793005)
                    .wrapping_add(1442695040888963407);
                ((lcg >> 11) as f32 / (1u64 << 53) as f32) * 2.0 - 1.0
            })
            .collect();
        Self {
            embed_dim,
            total_patches,
            weights,
        }
    }

    pub fn encode(&self, patch_indices: &[usize]) -> Vec<f32> {
        let n = patch_indices.len();
        let d = self.embed_dim;
        let mut out = vec![0.0f32; n * d];
        for (i, &pi) in patch_indices.iter().enumerate() {
            let pi = pi.min(self.total_patches - 1);
            for dd in 0..d {
                out[i * d + dd] = self.weights[pi * d + dd];
            }
        }
        out
    }

    pub fn l2_normalize(reps: &mut Vec<f32>, n: usize, d: usize) {
        for i in 0..n {
            let norm = (0..d)
                .map(|dd| reps[i * d + dd].powi(2))
                .sum::<f32>()
                .sqrt()
                .max(1e-8);
            for dd in 0..d {
                reps[i * d + dd] /= norm;
            }
        }
    }
}

/// V-JEPA training step metrics
#[derive(Debug, Clone)]
pub struct VJepaStepMetrics {
    pub loss: f32,
    pub tube_losses: Vec<f32>,
    pub ema_tau: f64,
    pub num_context_patches: usize,
    pub num_target_patches: usize,
    pub mask_ratio_actual: f64,
}

/// V-JEPA trainer (video analog of JepaTrainer)
#[derive(Debug, Clone)]
pub struct VJepaTrainer {
    pub config: VJepaConfig,
    pub context_encoder: MockVideoEncoder,
    pub target_ema: Vec<f32>, // shadow weights of context_encoder
    pub predictor: JepaPredictor,
    pub mask_strategy: TubeMaskStrategy,
    pub step: usize,
    pub total_steps: usize,
}

impl VJepaTrainer {
    pub fn new(config: VJepaConfig, total_steps: usize) -> Self {
        let total_patches = config.video.total_patches();
        let embed_dim = config.video.embed_dim;
        let pred_dim = (embed_dim / 4).max(64);
        let num_heads = (pred_dim / 32).max(1).min(8);

        let context_encoder = MockVideoEncoder::new(embed_dim, total_patches);
        let target_ema = context_encoder.weights.clone();
        let predictor = JepaPredictor::new(
            embed_dim,
            pred_dim,
            config.predictor_depth,
            num_heads,
            total_patches,
        );
        let mask_strategy = TubeMaskStrategy::new(config.mask.clone());

        Self {
            config,
            context_encoder,
            target_ema,
            predictor,
            mask_strategy,
            step: 0,
            total_steps,
        }
    }

    fn tau(&self) -> f64 {
        let p = self.step as f64 / self.total_steps.max(1) as f64;
        self.config.ema_tau_start + (self.config.ema_tau_end - self.config.ema_tau_start) * p
    }

    fn target_encode(&self, patch_indices: &[usize]) -> Vec<f32> {
        let n = patch_indices.len();
        let d = self.config.video.embed_dim;
        let np = self.config.video.total_patches();
        let mut out = vec![0.0f32; n * d];
        for (i, &pi) in patch_indices.iter().enumerate() {
            let pi = pi.min(np - 1);
            for dd in 0..d {
                out[i * d + dd] = self.target_ema[pi * d + dd];
            }
        }
        out
    }

    fn ema_update(&mut self, tau: f64) {
        for (sw, &ow) in self
            .target_ema
            .iter_mut()
            .zip(self.context_encoder.weights.iter())
        {
            *sw = (tau as f32) * *sw + (1.0 - tau as f32) * ow;
        }
    }

    /// One V-JEPA training step
    pub fn train_step(&mut self, batch_size: usize) -> VJepaStepMetrics {
        let video_cfg = self.config.video.clone();
        let embed_dim = video_cfg.embed_dim;
        let masks = self.mask_strategy.generate_batch(batch_size, &video_cfg);
        let tau = self.tau();

        let mut total_loss = 0.0f32;
        let num_tubes = self.config.mask.num_target_tubes;
        let mut tube_loss_acc = vec![0.0f32; num_tubes];
        let mut total_ctx = 0usize;
        let mut total_tgt = 0usize;

        for (bi, mask) in masks.iter().enumerate() {
            let _ = bi;
            let ctx_reps = self.context_encoder.encode(&mask.context_indices);
            let all_target = &mask.all_target_indices;
            let predicted = self
                .predictor
                .forward(&ctx_reps, &mask.context_indices, all_target);
            let mut target_reps = self.target_encode(all_target);
            let n_tgt = all_target.len();
            MockVideoEncoder::l2_normalize(&mut target_reps, n_tgt, embed_dim);
            let mut pred_n = predicted.clone();
            MockVideoEncoder::l2_normalize(&mut pred_n, n_tgt, embed_dim);

            let mut offsets = vec![0usize; mask.target_tubes.len()];
            let mut cursor = 0usize;
            for (k, tube) in mask.target_tubes.iter().enumerate() {
                offsets[k] = cursor.min(n_tgt);
                cursor += tube.len();
            }

            let loss_result = jepa_loss(
                &pred_n,
                &target_reps,
                embed_dim,
                &mask.target_tubes,
                &offsets,
            );
            total_loss += loss_result.loss;
            for (k, &bl) in loss_result.block_losses.iter().enumerate() {
                if k < tube_loss_acc.len() {
                    tube_loss_acc[k] += bl;
                }
            }
            total_ctx += mask.num_context();
            total_tgt += mask.num_target_patches();
        }

        self.ema_update(tau);
        self.step += 1;

        let bs = batch_size.max(1) as f32;
        let avg_ctx = total_ctx / batch_size.max(1);
        let total_p = video_cfg.total_patches();
        VJepaStepMetrics {
            loss: total_loss / bs,
            tube_losses: tube_loss_acc.iter().map(|&v| v / bs).collect(),
            ema_tau: tau,
            num_context_patches: avg_ctx,
            num_target_patches: total_tgt / batch_size.max(1),
            mask_ratio_actual: 1.0 - (avg_ctx as f64 / total_p as f64),
        }
    }
}

// ============================================================================
// Neuro-JEPA: Signal-domain JEPA for EEG/MEG
// ============================================================================

/// Configuration for applying JEPA to multi-channel neural time series
#[derive(Debug, Clone)]
pub struct NeuroJepaConfig {
    /// Number of EEG/MEG channels
    pub num_channels: usize,
    /// Segment length in samples per patch
    pub segment_len: usize,
    /// Total number of time samples in a recording window
    pub total_samples: usize,
    /// Embedding dimension for each channel-segment patch
    pub embed_dim: usize,
    /// Fraction of segments to mask as target
    pub mask_ratio: f64,
    /// EMA tau
    pub ema_tau: f64,
}

impl Default for NeuroJepaConfig {
    fn default() -> Self {
        Self {
            num_channels: 64,    // typical EEG montage
            segment_len: 64,     // ~250ms at 256Hz
            total_samples: 1024, // ~4s window at 256Hz
            embed_dim: 256,
            mask_ratio: 0.75,
            ema_tau: 0.996,
        }
    }
}

impl NeuroJepaConfig {
    /// Number of time segments (patches in time dimension)
    pub fn num_time_segments(&self) -> usize {
        self.total_samples / self.segment_len
    }
    /// Total number of patches = channels × time_segments
    pub fn total_patches(&self) -> usize {
        self.num_channels * self.num_time_segments()
    }
}

/// Mask result for neural signal JEPA
#[derive(Debug, Clone)]
pub struct NeuroMaskResult {
    /// Patch indices (channel × time) visible to context encoder
    pub context_indices: Vec<usize>,
    /// Target patch indices
    pub target_indices: Vec<usize>,
}

/// Neuro-JEPA masking: independent channel-time tube masking
///
/// Masks entire time-segments of selected channels (tube in channel-time space)
#[derive(Debug, Clone)]
pub struct NeuroMaskStrategy {
    lcg: u64,
}

impl NeuroMaskStrategy {
    pub fn new() -> Self {
        Self { lcg: 1111111111 }
    }

    fn lcg_next(&mut self) -> u64 {
        self.lcg = self
            .lcg
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        self.lcg
    }
    fn rand_usize(&mut self, n: usize) -> usize {
        (self.lcg_next() % n as u64) as usize
    }

    /// Generate channel-tube masks: select `mask_fraction * num_channels` channels,
    /// mask ALL time segments for those channels (tube masking in channel dimension).
    pub fn generate(&mut self, cfg: &NeuroJepaConfig, seed_offset: u64) -> NeuroMaskResult {
        self.lcg = self.lcg.wrapping_add(seed_offset.wrapping_mul(987654321));
        let ts = cfg.num_time_segments();
        let nc = cfg.num_channels;
        let total = cfg.total_patches();

        // Select channels to mask (without replacement)
        let num_mask_channels = ((cfg.mask_ratio * nc as f64).ceil() as usize).min(nc);
        let mut channel_order: Vec<usize> = (0..nc).collect();
        for i in 0..num_mask_channels {
            let j = i + self.rand_usize(nc - i);
            channel_order.swap(i, j);
        }
        let masked_channels: std::collections::BTreeSet<usize> =
            channel_order[..num_mask_channels].iter().cloned().collect();

        let mut context_indices: Vec<usize> = Vec::new();
        let mut target_indices: Vec<usize> = Vec::new();
        for c in 0..nc {
            for t in 0..ts {
                let idx = c * ts + t;
                if idx < total {
                    if masked_channels.contains(&c) {
                        target_indices.push(idx);
                    } else {
                        context_indices.push(idx);
                    }
                }
            }
        }
        NeuroMaskResult {
            context_indices,
            target_indices,
        }
    }
}

impl Default for NeuroMaskStrategy {
    fn default() -> Self {
        Self::new()
    }
}

// ============================================================================
// Tests — Batch 23 (V-JEPA + Neuro-JEPA)
// ============================================================================

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

    // ── Video patch config ────────────────────────────────────────────────────

    #[test]
    fn test_video_patch_config_totals() {
        let cfg = VideoPatchConfig {
            spatial_patch_size: 16,
            temporal_patch_size: 2,
            image_size: 64,
            num_frames: 8,
            embed_dim: 128,
        };
        assert_eq!(cfg.spatial_grid(), 4);
        assert_eq!(cfg.temporal_segments(), 4);
        assert_eq!(cfg.total_patches(), 64); // 4 × 4 × 4
    }

    #[test]
    fn test_thw_roundtrip() {
        let cfg = VideoPatchConfig::default();
        let idx = cfg.thw_to_patch_idx(2, 3, 5);
        let (t, h, w) = cfg.patch_idx_to_thw(idx);
        assert_eq!((t, h, w), (2, 3, 5));
    }

    #[test]
    fn test_patch_embed_3d_output_shape() {
        let cfg = VideoPatchConfig {
            spatial_patch_size: 4,
            temporal_patch_size: 2,
            image_size: 8,
            num_frames: 4,
            embed_dim: 32,
        };
        let total = cfg.total_patches(); // (4/2)×(8/4)×(8/4) = 2×2×2 = 8
        let embed = PatchEmbed3D::new(cfg.clone(), 3);
        let video = vec![0.0f32; 4 * 8 * 8 * 3]; // [T, H, W, C]
        let out = embed.embed(&video, 3);
        assert_eq!(out.len(), total * 32);
    }

    #[test]
    fn test_patch_embed_nonzero() {
        let cfg = VideoPatchConfig {
            spatial_patch_size: 4,
            temporal_patch_size: 1,
            image_size: 8,
            num_frames: 2,
            embed_dim: 16,
        };
        let embed = PatchEmbed3D::new(cfg.clone(), 1);
        let video: Vec<f32> = (0..2 * 8 * 8).map(|i| i as f32 * 0.01).collect();
        let out = embed.embed(&video, 1);
        let any_nonzero = out.iter().any(|&v| v.abs() > 1e-8);
        assert!(
            any_nonzero,
            "non-zero input should produce non-zero embeddings"
        );
    }

    // ── Tube masking ──────────────────────────────────────────────────────────

    #[test]
    fn test_tube_mask_basic() {
        let video_cfg = VJepaConfig::small_test().video;
        let mask_cfg = TubeMaskConfig {
            mask_ratio: 0.75,
            num_target_tubes: 2,
            ..Default::default()
        };
        let mut strategy = TubeMaskStrategy::new(mask_cfg);
        let result = strategy.generate(&video_cfg, 0);
        assert!(!result.context_indices.is_empty());
        assert!(!result.all_target_indices.is_empty());
        assert_eq!(result.target_tubes.len(), 2);
    }

    #[test]
    fn test_tube_mask_no_overlap() {
        let video_cfg = VJepaConfig::small_test().video;
        let mut strategy = TubeMaskStrategy::new(TubeMaskConfig::default());
        let result = strategy.generate(&video_cfg, 1);
        for &ci in &result.context_indices {
            assert!(
                !result.all_target_indices.contains(&ci),
                "context patch {ci} appears in target"
            );
        }
    }

    #[test]
    fn test_tube_mask_indices_in_range() {
        let video_cfg = VJepaConfig::small_test().video;
        let total = video_cfg.total_patches();
        let mut strategy = TubeMaskStrategy::new(TubeMaskConfig::default());
        let result = strategy.generate(&video_cfg, 0);
        for &ci in &result.context_indices {
            assert!(ci < total);
        }
        for &ti in &result.all_target_indices {
            assert!(ti < total);
        }
    }

    #[test]
    fn test_tube_spans_all_time_segments() {
        let cfg = VideoPatchConfig {
            spatial_patch_size: 4,
            temporal_patch_size: 1,
            image_size: 8,
            num_frames: 4,
            embed_dim: 32,
        };
        let mut strategy = TubeMaskStrategy::new(TubeMaskConfig {
            mask_ratio: 0.9,
            num_target_tubes: 1,
            spatial_scale_range: (0.3, 0.4),
        });
        let result = strategy.generate(&cfg, 0);
        // Target tubes should span multiple time steps (temporal consistency)
        let ts = cfg.temporal_segments();
        let sg = cfg.spatial_grid();
        if !result.all_target_indices.is_empty() {
            // All time segments should appear somewhere in target
            let time_steps_covered: std::collections::BTreeSet<usize> = result
                .all_target_indices
                .iter()
                .map(|&idx| idx / (sg * sg))
                .collect();
            assert!(
                time_steps_covered.len() <= ts,
                "temporal coverage must be within bounds"
            );
        }
    }

    #[test]
    fn test_tube_mask_batch() {
        let video_cfg = VJepaConfig::small_test().video;
        let mut strategy = TubeMaskStrategy::new(TubeMaskConfig::default());
        let batch = strategy.generate_batch(4, &video_cfg);
        assert_eq!(batch.len(), 4);
        for mask in &batch {
            assert!(!mask.context_indices.is_empty());
        }
    }

    #[test]
    fn test_tube_mask_vary_across_batch() {
        // Use a larger grid so masks are unlikely to collide
        let video_cfg = VideoPatchConfig {
            spatial_patch_size: 4,
            temporal_patch_size: 2,
            image_size: 32,
            num_frames: 8,
            embed_dim: 64,
        };
        let mut strategy = TubeMaskStrategy::new(TubeMaskConfig::default());
        let batch = strategy.generate_batch(8, &video_cfg);
        // At least some pairs in the batch should differ
        let any_differ = batch
            .windows(2)
            .any(|w| w[0].all_target_indices != w[1].all_target_indices);
        assert!(
            any_differ,
            "different seeds should produce different masks in at least one pair"
        );
    }

    // ── V-JEPA trainer ────────────────────────────────────────────────────────

    #[test]
    fn test_vjepa_trainer_step_basic() {
        let cfg = VJepaConfig::small_test();
        let mut trainer = VJepaTrainer::new(cfg, 100);
        let metrics = trainer.train_step(2);
        assert!(metrics.loss.is_finite());
        assert!(metrics.ema_tau >= 0.996 && metrics.ema_tau <= 1.0);
        assert!(metrics.num_context_patches >= 1);
        assert!(metrics.num_target_patches >= 1);
    }

    #[test]
    fn test_vjepa_mask_ratio_actual() {
        let cfg = VJepaConfig::small_test();
        let mut trainer = VJepaTrainer::new(cfg, 100);
        let metrics = trainer.train_step(4);
        assert!(metrics.mask_ratio_actual >= 0.0 && metrics.mask_ratio_actual <= 1.0);
    }

    #[test]
    fn test_vjepa_step_increments() {
        let cfg = VJepaConfig::small_test();
        let mut trainer = VJepaTrainer::new(cfg, 100);
        trainer.train_step(2);
        assert_eq!(trainer.step, 1);
        trainer.train_step(2);
        assert_eq!(trainer.step, 2);
    }

    #[test]
    fn test_vjepa_ema_anneals() {
        let cfg = VJepaConfig::small_test();
        let mut trainer = VJepaTrainer::new(cfg, 5);
        let m0 = trainer.train_step(2);
        for _ in 0..4 {
            trainer.train_step(2);
        }
        let m4 = trainer.train_step(2);
        assert!(m4.ema_tau >= m0.ema_tau, "tau must anneal upward");
    }

    #[test]
    fn test_vjepa_tube_losses_count() {
        let mut cfg = VJepaConfig::small_test();
        cfg.mask.num_target_tubes = 3;
        let mut trainer = VJepaTrainer::new(cfg, 10);
        let metrics = trainer.train_step(2);
        assert_eq!(metrics.tube_losses.len(), 3);
    }

    // ── Neuro-JEPA ───────────────────────────────────────────────────────────

    #[test]
    fn test_neuro_config_patches() {
        let cfg = NeuroJepaConfig::default();
        assert_eq!(cfg.num_time_segments(), 16); // 1024 / 64
        assert_eq!(cfg.total_patches(), 64 * 16); // 64 channels × 16 segments
    }

    #[test]
    fn test_neuro_mask_basic() {
        let cfg = NeuroJepaConfig::default();
        let mut strategy = NeuroMaskStrategy::new();
        let result = strategy.generate(&cfg, 0);
        assert!(!result.context_indices.is_empty());
        assert!(!result.target_indices.is_empty());
    }

    #[test]
    fn test_neuro_mask_no_overlap() {
        let cfg = NeuroJepaConfig::default();
        let mut strategy = NeuroMaskStrategy::new();
        let result = strategy.generate(&cfg, 0);
        for &ci in &result.context_indices {
            assert!(!result.target_indices.contains(&ci));
        }
    }

    #[test]
    fn test_neuro_mask_covers_all_patches() {
        let cfg = NeuroJepaConfig {
            num_channels: 4,
            segment_len: 16,
            total_samples: 64,
            ..Default::default()
        };
        let mut strategy = NeuroMaskStrategy::new();
        let result = strategy.generate(&cfg, 0);
        let total = cfg.total_patches();
        let covered = result.context_indices.len() + result.target_indices.len();
        assert_eq!(
            covered, total,
            "all patches must appear in context or target"
        );
    }

    #[test]
    fn test_neuro_mask_ratio_respected() {
        let cfg = NeuroJepaConfig {
            num_channels: 8,
            mask_ratio: 0.5,
            ..Default::default()
        };
        let mut strategy = NeuroMaskStrategy::new();
        let result = strategy.generate(&cfg, 0);
        let tgt_channels = result.target_indices.len() / cfg.num_time_segments();
        let expected = ((0.5 * cfg.num_channels as f64).ceil() as usize).min(cfg.num_channels);
        assert_eq!(tgt_channels, expected, "masked channels should match ratio");
    }

    #[test]
    fn test_neuro_mask_tube_structure() {
        // Verify that masked channels are fully masked across all time segments
        let cfg = NeuroJepaConfig {
            num_channels: 4,
            segment_len: 4,
            total_samples: 16,
            mask_ratio: 0.5,
            ..Default::default()
        };
        let ts = cfg.num_time_segments();
        let mut strategy = NeuroMaskStrategy::new();
        let result = strategy.generate(&cfg, 0);

        // Each target channel should have ALL its time segments in the target
        let mut channel_times: std::collections::HashMap<usize, usize> = Default::default();
        for &ti in &result.target_indices {
            let channel = ti / ts;
            *channel_times.entry(channel).or_insert(0) += 1;
        }
        for (&_ch, &count) in &channel_times {
            assert_eq!(
                count, ts,
                "tube masking: a masked channel must have ALL time segments masked"
            );
        }
    }

    #[test]
    fn test_neuro_vary_across_seeds() {
        let cfg = NeuroJepaConfig::default();
        let mut strategy = NeuroMaskStrategy::new();
        let r0 = strategy.generate(&cfg, 0);
        let r1 = strategy.generate(&cfg, 999);
        assert_ne!(
            r0.target_indices, r1.target_indices,
            "different seeds should produce different masks"
        );
    }
}