/// Layer-by-layer parity test against PyTorch reference.
///
/// Requires: data/layer_outputs.json (from scripts/dump_layer_outputs.py)
///           data/model_base.safetensors, data/config.json
///
/// Run: cargo test --test parity_test -- --nocapture

use burn::backend::NdArray as B;
use burn::prelude::*;
use burn::tensor::TensorData;
use std::path::Path;

fn load_json(path: &str) -> serde_json::Value {
    let s = std::fs::read_to_string(path).expect("Failed to read JSON");
    serde_json::from_str(&s).expect("Failed to parse JSON")
}

fn json_to_f32_vec(val: &serde_json::Value) -> Vec<f32> {
    val.as_array().unwrap().iter().map(|v| v.as_f64().unwrap() as f32).collect()
}

fn json_to_usize_vec(val: &serde_json::Value) -> Vec<usize> {
    val.as_array().unwrap().iter().map(|v| v.as_u64().unwrap() as usize).collect()
}

fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 {
    a.iter().zip(b.iter()).map(|(x, y)| (x - y).abs()).fold(0.0f32, f32::max)
}

fn mean_abs_diff(a: &[f32], b: &[f32]) -> f32 {
    let sum: f32 = a.iter().zip(b.iter()).map(|(x, y)| (x - y).abs()).sum();
    sum / a.len() as f32
}

fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
    dot / (norm_a * norm_b + 1e-12)
}

#[test]
fn test_layer_by_layer_parity() {
    let layer_path = "data/layer_outputs.json";
    let config_path = Path::new("data/config.json");
    let weights_path = Path::new("data/model_base.safetensors");

    if !Path::new(layer_path).exists() || !config_path.exists() || !weights_path.exists() {
        eprintln!("Skipping parity test: required files not found.");
        eprintln!("  Run: python3 scripts/dump_layer_outputs.py");
        return;
    }

    let ref_data = load_json(layer_path);
    let device = burn::backend::ndarray::NdArrayDevice::Cpu;

    // Load model
    let cfg_str = std::fs::read_to_string(config_path).unwrap();
    let model_cfg: sleepfm::ModelConfig = serde_json::from_str(&cfg_str).unwrap();
    let model = sleepfm::weights::load_model::<B>(&model_cfg, weights_path.to_str().unwrap(), &device)
        .expect("Failed to load model");

    // Load input
    let input_flat = json_to_f32_vec(&ref_data["input"]);
    let n_channels = ref_data["n_channels"].as_u64().unwrap() as usize;
    let n_samples = ref_data["n_samples"].as_u64().unwrap() as usize;
    let input = Tensor::<B, 3>::from_data(
        TensorData::new(input_flat.clone(), vec![1, n_channels, n_samples]),
        &device,
    );
    let mask_data: Vec<bool> = vec![false; n_channels];
    let mask = Tensor::<B, 2, Bool>::from_data(
        TensorData::new(mask_data, vec![1, n_channels]),
        &device,
    );

    eprintln!("\n=== Layer-by-layer parity test ===");
    eprintln!("Input: [{}, {}, {}]", 1, n_channels, n_samples);

    // ── Step 1: Tokenizer ────────────────────────────────────────────────────
    let tokens = model.patch_embedding.forward(input.clone());
    let ref_tokens = json_to_f32_vec(&ref_data["after_tokenizer"]);
    let ref_shape = json_to_usize_vec(&ref_data["after_tokenizer_shape"]);
    let rust_tokens: Vec<f32> = tokens.clone().into_data().to_vec::<f32>().unwrap();

    let tok_max = max_abs_diff(&rust_tokens, &ref_tokens);
    let tok_mean = mean_abs_diff(&rust_tokens, &ref_tokens);
    let tok_cos = cosine_similarity(&rust_tokens, &ref_tokens);
    eprintln!("Tokenizer:    max_diff={tok_max:.6}  mean_diff={tok_mean:.6}  cos_sim={tok_cos:.6}  shape={:?}", ref_shape);

    // ── Step 2+3: Spatial pooling ────────────────────────────────────────────
    let [b, c, s, e] = tokens.dims();
    // rearrange [B, C, S, E] -> [(B*S), C, E]
    let tokens_re = tokens.swap_dims(1, 2).reshape([b * s, c, e]);
    // mask: [B, C] -> expand to [B, S, C] -> reshape [(B*S), C]
    let mask_expanded = mask.clone()
        .unsqueeze_dim::<3>(1)
        .expand([b, s, n_channels])
        .reshape([b * s, n_channels]);
    let spatial = model.spatial_pooling.forward(tokens_re, Some(mask_expanded));
    let spatial = spatial.reshape([b, s, e]);

    let ref_spatial = json_to_f32_vec(&ref_data["after_spatial_pooling"]);
    let rust_spatial: Vec<f32> = spatial.clone().into_data().to_vec::<f32>().unwrap();
    let sp_max = max_abs_diff(&rust_spatial, &ref_spatial);
    let sp_mean = mean_abs_diff(&rust_spatial, &ref_spatial);
    let sp_cos = cosine_similarity(&rust_spatial, &ref_spatial);
    eprintln!("Spatial pool:  max_diff={sp_max:.6}  mean_diff={sp_mean:.6}  cos_sim={sp_cos:.6}");

    // ── Step 4: Positional encoding + LayerNorm ──────────────────────────────
    let spatial_pe = sleepfm::model::positional_encoding::add_positional_encoding(
        spatial, &model.pe_buffer,
    );
    let spatial_ln = model.layer_norm.forward(spatial_pe);

    let ref_pe_ln = json_to_f32_vec(&ref_data["after_pos_enc_ln"]);
    let rust_pe_ln: Vec<f32> = spatial_ln.clone().into_data().to_vec::<f32>().unwrap();
    let pe_max = max_abs_diff(&rust_pe_ln, &ref_pe_ln);
    let pe_mean = mean_abs_diff(&rust_pe_ln, &ref_pe_ln);
    let pe_cos = cosine_similarity(&rust_pe_ln, &ref_pe_ln);
    eprintln!("PosEnc+LN:     max_diff={pe_max:.6}  mean_diff={pe_mean:.6}  cos_sim={pe_cos:.6}");

    // ── Step 5: Transformer encoder ──────────────────────────────────────────
    let mut te_out = spatial_ln;
    for layer in &model.transformer_encoder {
        te_out = layer.forward(te_out, None);
    }
    let ref_te = json_to_f32_vec(&ref_data["after_transformer_encoder"]);
    let rust_te: Vec<f32> = te_out.clone().into_data().to_vec::<f32>().unwrap();
    let te_max = max_abs_diff(&rust_te, &ref_te);
    let te_mean = mean_abs_diff(&rust_te, &ref_te);
    let te_cos = cosine_similarity(&rust_te, &ref_te);
    eprintln!("Transformer:   max_diff={te_max:.6}  mean_diff={te_mean:.6}  cos_sim={te_cos:.6}");

    // ── Step 6: Temporal pooling ─────────────────────────────────────────────
    let pooled = model.temporal_pooling.forward(te_out, None);
    let ref_pooled = json_to_f32_vec(&ref_data["pooled"]);
    let rust_pooled: Vec<f32> = pooled.into_data().to_vec::<f32>().unwrap();
    let p_max = max_abs_diff(&rust_pooled, &ref_pooled);
    let p_mean = mean_abs_diff(&rust_pooled, &ref_pooled);
    let p_cos = cosine_similarity(&rust_pooled, &ref_pooled);
    eprintln!("Temporal pool: max_diff={p_max:.6}  mean_diff={p_mean:.6}  cos_sim={p_cos:.6}");

    // ── Also test full forward ───────────────────────────────────────────────
    let (pooled_full, _) = model.forward(input, mask);
    let rust_pooled_full: Vec<f32> = pooled_full.into_data().to_vec::<f32>().unwrap();
    let full_max = max_abs_diff(&rust_pooled_full, &ref_pooled);
    let full_cos = cosine_similarity(&rust_pooled_full, &ref_pooled);
    eprintln!("Full forward:  max_diff={full_max:.6}  cos_sim={full_cos:.6}");

    eprintln!("\n=== Summary ===");
    eprintln!("Rust first 5 pooled:   {:?}", &rust_pooled[..5]);
    eprintln!("PyTorch first 5 pooled: {:?}", &ref_pooled[..5]);

    // ── Assertions ───────────────────────────────────────────────────────────
    // Tokenizer: effectively exact (float rounding only)
    assert!(tok_max < 1e-4, "Tokenizer max_abs_diff too large: {tok_max}");
    assert!(tok_cos > 0.999999, "Tokenizer cosine similarity too low: {tok_cos}");

    // Spatial pooling: effectively exact
    assert!(sp_max < 1e-4, "Spatial pooling max_abs_diff too large: {sp_max}");
    assert!(sp_cos > 0.999999, "Spatial pooling cosine similarity too low: {sp_cos}");

    // PosEnc + LayerNorm: effectively exact
    assert!(pe_max < 1e-4, "PosEnc+LN max_abs_diff too large: {pe_max}");

    // Transformer encoder: effectively exact
    assert!(te_max < 1e-4, "Transformer max_abs_diff too large: {te_max}");
    assert!(te_cos > 0.999999, "Transformer cosine similarity too low: {te_cos}");

    // Final pooled output: effectively exact
    assert!(p_max < 1e-4, "Final pooled max_abs_diff too large: {p_max}");
    assert!(p_cos > 0.999999, "Final pooled cosine similarity too low: {p_cos}");

    eprintln!("\n✓ All parity checks passed!");
}

/// Full-model parity with the multi-modality reference (random inputs from torch.manual_seed(42))
#[test]
fn test_full_parity_all_modalities() {
    let ref_path = "data/reference_outputs.json";
    let config_path = Path::new("data/config.json");
    let weights_path = Path::new("data/model_base.safetensors");

    if !Path::new(ref_path).exists() || !config_path.exists() || !weights_path.exists() {
        eprintln!("Skipping full parity test: required files not found.");
        return;
    }

    let ref_data = load_json(ref_path);
    let device = burn::backend::ndarray::NdArrayDevice::Cpu;

    let (encoder, _) = sleepfm::SleepFmEncoder::<B>::load(
        config_path, weights_path, device.clone(),
    ).unwrap();

    eprintln!("\n=== Full model parity test (all modalities) ===");

    for name in &["BAS", "RESP", "EKG", "EMG"] {
        let entry = &ref_data[*name];
        let n_channels = entry["n_channels"].as_u64().unwrap() as usize;
        let input_flat = json_to_f32_vec(&entry["input_flat"]);
        let ref_pooled = json_to_f32_vec(&entry["pooled"]);
        let t_samples = input_flat.len() / n_channels;

        let signal = Tensor::<B, 3>::from_data(
            TensorData::new(input_flat, vec![1, n_channels, t_samples]),
            &device,
        );
        let mask_data: Vec<bool> = vec![false; n_channels];
        let mask = Tensor::<B, 2, Bool>::from_data(
            TensorData::new(mask_data, vec![1, n_channels]),
            &device,
        );

        let result = encoder.run(signal, mask).unwrap();

        let cos = cosine_similarity(&result.pooled, &ref_pooled);
        let max_d = max_abs_diff(&result.pooled, &ref_pooled);
        let mean_d = mean_abs_diff(&result.pooled, &ref_pooled);
        eprintln!("{name:>4}: cos_sim={cos:.6}  max_diff={max_d:.6}  mean_diff={mean_d:.6}");
        eprintln!("      Rust:    {:?}", &result.pooled[..5]);
        eprintln!("      PyTorch: {:?}", &ref_pooled[..5]);

        assert!(cos > 0.9999, "{name} cosine similarity too low: {cos}");
        assert!(max_d < 0.01, "{name} max_abs_diff too large: {max_d}");
    }

    eprintln!("\n✓ All modality parity checks passed!");
}
