//! Calibration: integer `jpegai_synthesis::decode_hyper` vs the float reference —
//! run the SAME hyper-latent `z` through `hyper_decoder.onnx` (ort) and through
//! `decode_hyper`, compare the 640-channel entropy params `p`. The output is not
//! pixels but fixed-point SGMM params, so we use a scale-invariant SNR
//! (20·log10(rms(signal)/rms(error))) plus max abs error. Needs `neural` (ort).
//!
//! env: WAI_JPEGAI_HYPERDEC = path to common_modules/hyper_decoder.onnx,
//!      WAI_JPEGAI_HYPERDEC_WEIGHTS = path to the exported 16-bit weights JSON.
#![cfg(all(feature = "neural", feature = "legacy_int"))]

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

fn ort_run(path: &str, x: Vec<f32>, c: usize, h: usize, w: usize) -> (Vec<f32>, Vec<usize>) {
    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 name = s.inputs().first().unwrap().name().to_string();
    let arr = ndarray::Array4::from_shape_vec((1, c, h, w), x).unwrap();
    let v = Value::from_array(arr).unwrap();
    let out = s.run(ort::inputs![name.as_str() => v]).unwrap();
    let view = out[0].try_extract_array::<f32>().unwrap();
    (view.iter().copied().collect(), view.shape().to_vec())
}

#[test]
fn hyper_decoder_int_vs_float_snr() {
    let onnx = match std::env::var("WAI_JPEGAI_HYPERDEC") {
        Ok(p) if std::path::Path::new(&p).exists() => p,
        _ => { eprintln!("skip: set WAI_JPEGAI_HYPERDEC to hyper_decoder.onnx"); return; }
    };
    let wjson = match std::env::var("WAI_JPEGAI_HYPERDEC_WEIGHTS") {
        Ok(p) if std::path::Path::new(&p).exists() => p,
        _ => { eprintln!("skip: set WAI_JPEGAI_HYPERDEC_WEIGHTS to the weights JSON"); return; }
    };
    let (c, h, w) = (160usize, 8usize, 8usize);
    let n = c * h * w;
    // integer-valued hyper-latent (entropy-decoded symbols), small magnitude
    let lf: Vec<f32> = (0..n).map(|i| ((i % 11) as f32) - 5.0).collect();

    let (pf, shape) = ort_run(&onnx, lf.clone(), c, h, w);
    let (fmn, fmx) = pf.iter().fold((f32::MAX, f32::MIN), |(a, b), &v| (a.min(v), b.max(v)));
    eprintln!("float hyper out shape = {:?}, len {}, range [{:.3}, {:.3}]", shape, pf.len(), fmn, fmx);

    let li: Vec<i64> = lf.iter().map(|&x| (x * 65536.0).round() as i64).collect();
    let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&wjson).unwrap()).unwrap();
    let wt = Weights::from_json(&v).unwrap();
    let (pi, ph, pw) = decode_hyper(&li, c, h, w, &wt);
    eprintln!("int hyper out = {} vals at [{},{}]", pi.len(), ph, pw);

    assert_eq!(pi.len(), pf.len(), "output size mismatch (640×2h×2w expected)");
    let (mut maxd, mut se, mut sig) = (0f64, 0f64, 0f64);
    for i in 0..pi.len() {
        let f = pf[i] as f64;
        let q = pi[i] as f64 / 65536.0; // dequantize Q16
        let d = (q - f).abs();
        maxd = maxd.max(d);
        se += d * d;
        sig += f * f;
    }
    let snr = 10.0 * (sig / se).log10();
    eprintln!("INT vs FLOAT hyper_decoder:  max|Δ| = {:.4}   SNR = {:.2} dB", maxd, snr);
    assert!(snr > 45.0, "hyper_decoder SNR {snr:.2} dB too low (expected high — pure conv chain)");
}