//! End-to-end integration check: run the synthesis decoders on the REAL reference
//! latent from `test256.bin` (checkpoint `yhat_y`/`yhat_uv`) and compare the pixels
//! to the reference synthesis output (`ypix`/`uvpix`). Unlike the per-net calibration
//! (synthetic latent vs the ONNX), this validates the synthesis stage on real
//! pipeline data — the first integration milestone. PSNR ≈ the 16-bit PTQ (~58 dB).
//!
//! env: WAI_JPEGAI_E2E_CKPT   = e2e_synth_checkpoint.json {yhat_y,yhat_uv,ypix,uvpix,y_lsbs}
//!      WAI_JPEGAI_Y_WEIGHTS  = synth16_y_weights.json
//!      WAI_JPEGAI_UV_WEIGHTS = synth16_uv_weights.json
#![cfg(all(feature = "neural_int", feature = "legacy_int"))]

use wai::jpegai_synthesis::{decode_y, decode_uv, Weights};

fn load(k: &str, v: &serde_json::Value) -> Vec<f64> {
    v[k].as_array().unwrap().iter().map(|x| x.as_f64().unwrap()).collect()
}
fn q16(f: &[f64]) -> Vec<i64> { f.iter().map(|&x| (x * 65536.0).round() as i64).collect() }
fn psnr_u8(got: &[u8], refpix: &[f64]) -> (f64, f64) {
    let (mut maxd, mut mse) = (0f64, 0f64);
    for i in 0..got.len() {
        let r = refpix[i].clamp(0.0, 255.0);
        let d = (got[i] as f64 - r).abs();
        maxd = maxd.max(d); mse += d * d;
    }
    mse /= got.len() as f64;
    (maxd, if mse > 0.0 { 20.0 * 255f64.log10() - 10.0 * mse.log10() } else { f64::INFINITY })
}

#[test]
fn synthesis_on_real_reference_latent() {
    let ckpt = match std::env::var("WAI_JPEGAI_E2E_CKPT") {
        Ok(p) if std::path::Path::new(&p).exists() => p,
        _ => { eprintln!("skip: set WAI_JPEGAI_E2E_CKPT"); return; }
    };
    let (yw, uw) = (std::env::var("WAI_JPEGAI_Y_WEIGHTS"), std::env::var("WAI_JPEGAI_UV_WEIGHTS"));
    let (yw, uw) = match (yw, uw) {
        (Ok(a), Ok(b)) if std::path::Path::new(&a).exists() && std::path::Path::new(&b).exists() => (a, b),
        _ => { eprintln!("skip: set WAI_JPEGAI_Y_WEIGHTS + WAI_JPEGAI_UV_WEIGHTS"); return; }
    };
    let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&ckpt).unwrap()).unwrap();
    let wy = Weights::from_json(&serde_json::from_slice(&std::fs::read(&yw).unwrap()).unwrap()).unwrap();
    let wu = Weights::from_json(&serde_json::from_slice(&std::fs::read(&uw).unwrap()).unwrap()).unwrap();

    // Y: decode_y(yhat_y) vs ypix. Try both the pre-lsbs (yhat_y) and post-lsbs
    // (y_lsbs) latent — whichever the synthesis actually consumes will hit ~58 dB.
    let ypix = load("ypix", &v);
    for cand in ["yhat_y", "y_lsbs"] {
        let lat = q16(&load(cand, &v));
        let y = decode_y(&lat, &wy);
        let (md, ps) = psnr_u8(&y, &ypix);
        eprintln!("Y synthesis on real latent [{cand}]: max|Δ|={md:.1}  PSNR={ps:.2} dB");
    }

    // Assert the better Y candidate reproduces the reference synthesis within PTQ.
    let best_y = ["yhat_y", "y_lsbs"].iter()
        .map(|c| psnr_u8(&decode_y(&q16(&load(c, &v)), &wy), &ypix).1)
        .fold(f64::MIN, f64::max);
    eprintln!("best Y PSNR = {best_y:.2} dB");
    assert!(best_y > 45.0, "Y synthesis on real latent regressed (best {best_y:.2} dB)");

    // UV: model_uv's input is the 256-ch (y‖uv) latent = concat(y_lsbs[160], uv_lsbs[96]).
    let uvpix = load("uvpix", &v);
    let mut uv_in: Vec<i64> = q16(&load("y_lsbs", &v)); // 160×16×16
    uv_in.extend(q16(&load("uv_lsbs", &v)));            // + 96×16×16 = 256×16×16
    let uv = decode_uv(&uv_in, &wu);
    let (mdu, psu) = psnr_u8(&uv, &uvpix);
    eprintln!("UV synthesis on real latent [concat(y_lsbs,uv_lsbs)]: max|Δ|={mdu:.1}  PSNR={psu:.2} dB");
    assert!(psu > 45.0, "UV synthesis on real latent regressed ({psu:.2} dB)");
}