//! Calibration: integer MCM stage decoders (`mcm_stage0` / `mcm_stage_n`) vs the
//! float reference — run the SAME inputs through each `MCM/stage{i}.onnx` (ort) and
//! through the Rust stage, compare the fixed-point context features `y[i]` (SNR,
//! since the output is params not pixels). Each stage validated independently.
//! Needs `neural` (ort).
//!
//! env: WAI_JPEGAI_MCM_DIR         = dir with stage0..3.onnx,
//!      WAI_JPEGAI_MCM_WEIGHTS_DIR = dir with mcm_stage0..3_y.json
#![cfg(all(feature = "neural", feature = "legacy_int"))]

use ort::session::{builder::GraphOptimizationLevel, Session};
use ort::value::Value;
use wai::jpegai_synthesis::{mcm_stage0, mcm_stage_n, Weights};

fn arr(c: usize, h: usize, w: usize, seed: i64) -> (Vec<f32>, Vec<i64>) {
    let n = c * h * w;
    let f: Vec<f32> = (0..n).map(|i| ((((i as i64 + seed) * 31) % 197 - 98) as f32) / 64.0).collect();
    let q: Vec<i64> = f.iter().map(|&x| (x * 65536.0).round() as i64).collect();
    (f, q)
}

fn run_stage(path: &str, named: Vec<(&str, Vec<f32>, usize)>, h: usize, w: usize) -> Vec<f32> {
    let mut s = Session::builder(wai::ort_environment().unwrap()).unwrap()
        .with_intra_threads(1).unwrap().with_inter_threads(1).unwrap()
        .with_optimization_level(GraphOptimizationLevel::Level3).unwrap()
        .commit_from_file(path).unwrap();
    let inputs: Vec<(String, Value)> = named.into_iter().map(|(name, data, c)| {
        let a = ndarray::Array4::from_shape_vec((1, c, h, w), data).unwrap();
        (name.to_string(), Value::from_array(a).unwrap().into_dyn())
    }).collect();
    let out = s.run(inputs).unwrap();
    out[0].try_extract_array::<f32>().unwrap().iter().copied().collect()
}

fn snr(int_q16: &[i64], flt: &[f32]) -> (f64, f64) {
    let (mut maxd, mut se, mut sig) = (0f64, 0f64, 0f64);
    for i in 0..int_q16.len() {
        let q = int_q16[i] as f64 / 65536.0;
        let f = flt[i] as f64;
        let d = (q - f).abs();
        maxd = maxd.max(d); se += d * d; sig += f * f;
    }
    (maxd, if se > 0.0 { 10.0 * (sig / se).log10() } else { f64::INFINITY })
}

#[test]
fn mcm_stages_int_vs_float_snr() {
    let dir = match std::env::var("WAI_JPEGAI_MCM_DIR") {
        Ok(p) if std::path::Path::new(&p).exists() => p,
        _ => { eprintln!("skip: set WAI_JPEGAI_MCM_DIR to the MCM onnx dir"); return; }
    };
    let wdir = match std::env::var("WAI_JPEGAI_MCM_WEIGHTS_DIR") {
        Ok(p) if std::path::Path::new(&p).exists() => p,
        _ => { eprintln!("skip: set WAI_JPEGAI_MCM_WEIGHTS_DIR to the weights dir"); return; }
    };
    let (h, w) = (16usize, 16usize);
    let load = |i: usize| -> Weights {
        let v: serde_json::Value = serde_json::from_slice(&std::fs::read(format!("{wdir}/mcm_stage{i}_y.json")).unwrap()).unwrap();
        Weights::from_json(&v).unwrap()
    };
    let mut worst = f64::INFINITY;

    // stage 0: single input p0
    {
        let (pf, pq) = arr(160, h, w, 1);
        let rf = run_stage(&format!("{dir}/stage0.onnx"), vec![("p0", pf, 160)], h, w);
        let mine = mcm_stage0(&pq, h, w, &load(0));
        let (md, sn) = snr(&mine, &rf);
        eprintln!("MCM stage0:  max|Δ| = {md:.4}  SNR = {sn:.2} dB  (len {})", mine.len());
        worst = worst.min(sn);
    }
    // stages 1..3: inputs (p_i, context of width 160/320/480)
    for (i, (pn, cn, cc)) in [("p1", "y[0]", 160usize), ("p2", "y[1-0]", 320), ("p3", "y[2-0]", 480)].into_iter().enumerate() {
        let stage = i + 1;
        let (pf, pq) = arr(160, h, w, 7 * stage as i64 + 1);
        let (cf, cq) = arr(cc, h, w, 13 * stage as i64 + 2);
        let rf = run_stage(&format!("{dir}/stage{stage}.onnx"), vec![(pn, pf, 160), (cn, cf, cc)], h, w);
        let mine = mcm_stage_n(&pq, &cq, cc, h, w, &load(stage));
        let (md, sn) = snr(&mine, &rf);
        eprintln!("MCM stage{stage}:  max|Δ| = {md:.4}  SNR = {sn:.2} dB  (len {})", mine.len());
        worst = worst.min(sn);
    }
    eprintln!("MCM worst-stage SNR = {worst:.2} dB");
    assert!(worst > 45.0, "an MCM stage regressed (worst SNR {worst:.2} dB)");
}