//! Calibration: my integer `jpegai_synthesis::decode_uv` vs the float reference —
//! run the SAME 256-channel latent through `synthesis16.onnx` (model_uv, ort) and
//! through `decode_uv`, compare the UV planes (max abs diff + PSNR). This is the
//! numerical-faithfulness pass for the chroma decoder. Needs `neural` (ort).
//!
//! env: WAI_JPEGAI_SYNTH16_UV = path to model_uv/synthesis16.onnx,
//!      WAI_JPEGAI_UV_WEIGHTS  = path to the exported 16-bit uv weights JSON.
#![cfg(all(feature = "neural", feature = "legacy_int"))]

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

fn ort_synth(path: &str, x: Vec<f32>) -> (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, 256, 16, 16), 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 synthesis_uv_int_vs_float_psnr() {
    let onnx = match std::env::var("WAI_JPEGAI_SYNTH16_UV") {
        Ok(p) if std::path::Path::new(&p).exists() => p,
        _ => { eprintln!("skip: set WAI_JPEGAI_SYNTH16_UV to model_uv/synthesis16.onnx"); return; }
    };
    let wjson = match std::env::var("WAI_JPEGAI_UV_WEIGHTS") {
        Ok(p) if std::path::Path::new(&p).exists() => p,
        _ => { eprintln!("skip: set WAI_JPEGAI_UV_WEIGHTS to the weights JSON"); return; }
    };
    // deterministic small latent [256,16,16] (typical post-LSBS magnitude)
    let n = 256 * 16 * 16;
    let lf: Vec<f32> = (0..n).map(|i| (((i % 97) as f32) - 48.0) / 96.0).collect();

    let (uvf, shape) = ort_synth(&onnx, lf.clone());
    let (fmn, fmx) = uvf.iter().fold((f32::MAX, f32::MIN), |(a, b), &v| (a.min(v), b.max(v)));
    eprintln!("float uv-synth out shape = {:?}, len {}, range [{:.3}, {:.3}]", shape, uvf.len(), fmn, fmx);
    eprintln!("float sample [0..6] = {:?}", &uvf[..6]);

    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 uvi = decode_uv(&li, &wt);

    assert_eq!(uvi.len(), uvf.len(), "output size mismatch (2×256×256 expected)");
    let (mut maxd, mut mse) = (0f64, 0f64);
    for i in 0..uvi.len() {
        let f = uvf[i].clamp(0.0, 255.0) as f64;
        let d = (uvi[i] as f64 - f).abs();
        maxd = maxd.max(d);
        mse += d * d;
    }
    mse /= uvi.len() as f64;
    let psnr = if mse > 0.0 { 20.0 * (255.0f64).log10() - 10.0 * mse.log10() } else { f64::INFINITY };
    eprintln!("INT vs FLOAT uv-synthesis:  max|Δ| = {:.2}   MSE = {:.2}   PSNR = {:.2} dB", maxd, mse, psnr);
    assert!(psnr > 45.0, "uv-synthesis PSNR {psnr:.2} dB too low (expected ~55-58)");
}