/// End-to-end pipeline test: compare Rust pipeline output against
/// the full_pipeline.safetensors reference generated by gen_vectors.py.
mod common;
use common::{load_vectors, max_abs_diff, array_std};
use exg::{preprocess, PipelineConfig};
use ndarray::Array2;

#[test]
fn full_pipeline_matches_mne_reference() {
    let vecs = load_vectors("full_pipeline");

    let x_arr  = vecs.get("input_raw").expect("missing 'input_raw'");
    let y_ref   = vecs.get("output").expect("missing 'output'");    // [E, C, T]
    let src_hz  = vecs.get("src_sfreq").unwrap()[[0]] as f32;

    let shape = x_arr.shape();
    let data = Array2::from_shape_vec(
        (shape[0], shape[1]),
        x_arr.iter().cloned().collect(),
    ).unwrap();

    // Channel positions: dummy zeros (not needed for DSP steps).
    let chan_pos = Array2::zeros((shape[0], 3));

    let cfg = PipelineConfig::default();
    let epochs = preprocess(data, chan_pos, src_hz, &cfg).unwrap();

    // Expected: [E, C, T] in y_ref.
    let ref_shape = y_ref.shape();
    // let (n_e, n_c, n_t) = (ref_shape[0], ref_shape[1], ref_shape[2]);
    let (n_e, _, _) = (ref_shape[0], ref_shape[1], ref_shape[2]);

    assert_eq!(epochs.len(), n_e,
        "epoch count: Rust={} MNE={n_e}", epochs.len());

    // Flatten for comparison.
    let got_flat: Vec<f32> = epochs.iter()
        .flat_map(|(ep, _)| ep.iter().cloned())
        .collect();
    let ref_flat: Vec<f32> = y_ref.iter().cloned().collect();

    assert_eq!(got_flat.len(), ref_flat.len());

    let got_arr = ndarray::Array::from_vec(got_flat).into_dyn();
    let ref_arr = ndarray::Array::from_vec(ref_flat).into_dyn();

    let max_err = max_abs_diff(&got_arr, &ref_arr);
    let sigma   = array_std(&ref_arr);

    // Cumulative tolerance: 0.5 % of signal std.
    assert!(max_err < 5e-3,
        "max abs error {max_err:.2e} >= 5e-3");
    assert!(max_err / sigma * 100.0 < 0.5,
        "relative error {:.4}% >= 0.5%  (sigma={sigma:.4})",
        max_err / sigma * 100.0);
}

/// Verify each intermediate step individually by comparing to the
/// snapshots saved in full_pipeline.safetensors.
#[test]
fn pipeline_step_resample() {
    use exg::resample::resample;
    let vecs = load_vectors("full_pipeline");
    let x_arr = vecs.get("input_raw").unwrap();
    let y_ref  = vecs.get("after_resample").unwrap();
    let src    = vecs.get("src_sfreq").unwrap()[[0]] as f32;

    let shape = x_arr.shape();
    let x = Array2::from_shape_vec(
        (shape[0], shape[1]),
        x_arr.iter().cloned().collect(),
    ).unwrap();
    let got = resample(&x, src, 256.0).unwrap();
    let got_dyn = got.into_dyn();
    let max_err = max_abs_diff(&got_dyn, y_ref);
    let sigma = array_std(y_ref);
    assert!(max_err / sigma * 100.0 < 0.2,
        "resample step: rel error {:.4}%", max_err / sigma * 100.0);
}

#[test]
fn pipeline_step_highpass() {
    use exg::filter::{design_highpass, apply_fir_zero_phase};
    let vecs = load_vectors("full_pipeline");
    let x_arr = vecs.get("after_resample").unwrap();
    let y_ref  = vecs.get("after_hp").unwrap();

    let shape = x_arr.shape();
    let mut data = Array2::from_shape_vec(
        (shape[0], shape[1]),
        x_arr.iter().cloned().collect(),
    ).unwrap();
    let h = design_highpass(0.5, 256.0);
    apply_fir_zero_phase(&mut data, &h).unwrap();
    let got_dyn = data.into_dyn();
    let max_err = max_abs_diff(&got_dyn, y_ref);
    let sigma = array_std(y_ref);
    assert!(max_err / sigma * 100.0 < 0.01,
        "hp filter step: rel error {:.4}%", max_err / sigma * 100.0);
}

#[test]
fn pipeline_step_reference() {
    use exg::reference::average_reference_inplace;
    let vecs = load_vectors("full_pipeline");
    let x_arr = vecs.get("after_hp").unwrap();
    let y_ref  = vecs.get("after_ref").unwrap();

    let shape = x_arr.shape();
    let mut data = Array2::from_shape_vec(
        (shape[0], shape[1]),
        x_arr.iter().cloned().collect(),
    ).unwrap();
    average_reference_inplace(&mut data);
    let got_dyn = data.into_dyn();
    let max_err = max_abs_diff(&got_dyn, y_ref);
    assert!(max_err < 1e-4, "reference step: max error {max_err:.2e}");
}

#[test]
fn pipeline_step_zscore() {
    use exg::normalize::zscore_global_inplace;
    let vecs = load_vectors("full_pipeline");
    let x_arr = vecs.get("after_ref").unwrap();
    let y_ref  = vecs.get("after_zscore").unwrap();

    let shape = x_arr.shape();
    let mut data = Array2::from_shape_vec(
        (shape[0], shape[1]),
        x_arr.iter().cloned().collect(),
    ).unwrap();
    zscore_global_inplace(&mut data);
    let got_dyn = data.into_dyn();
    let max_err = max_abs_diff(&got_dyn, y_ref);
    assert!(max_err < 1e-4, "zscore step: max error {max_err:.2e}");
}
