//! End-to-end JPEG AI decode in Rust: real bitstream → RGB pixels, no Python.
//!
//! Every stage was validated in isolation first (see jpegai_mans_golden.rs
//! and jpegai-vm-assets/README.md); this test chains them: codestream parse →
//! the entropy decodes and the integer σ-net through the integer operator set
//! (`jpegai_ops`, integer-operators §13.6.1) → dequant → MCM value path (ORT stage nets) →
//! LSBS → synthesis (ORT, 16×16-traced) → EFElinear → eICCI (ORT wrappers) →
//! EFEnonlinear → LEF → BT.709, compared against the reference decoder's PNG.
//!
//! The signalled parameters now come from a real PIH/TON header decode:
//! `parse_headers` reads the bitstream's own headers and supplies the beta
//! displacement, the EFElinear coefficients + means, the eICCI choices, the
//! EFEnonlinear weights + luma masks, the LEF model, and the per-region model
//! flags — all driven into the decode below, and independently checked against
//! the reference header trace by `jpegai_header_parser_matches_trace`. What
//! remains fixtured is sink-supplied model state (the ORT nets, the integer
//! σ-net weights, the normative tables, the per-checkpoint constants) plus a
//! few intermediate reference dumps kept as scaffolding for stage cross-checks
//! (the per-position dequant gain map, the z factorized-CDF offsets, the LSBS
//! thresholds); the header-signalled parameters are no longer among them.
//!
//! Needs both `neural` (the ORT stage/synthesis nets + the jpegai modules) and
//! `codecs` (the `image` crate, to load the reference PNG it compares against).
#![cfg(all(feature = "neural", feature = "codecs"))]

use std::path::{Path, PathBuf};

use ort::session::{builder::GraphOptimizationLevel, Session};
use ort::value::Value;

use wai::neural::jpegai_bitstream::{parse_ae_substream, parse_codestream};
use wai::neural::jpegai_header::parse_headers;
use wai::neural::jpegai_filters::{efe_linear_plane, efe_nonlinear_plane, lef_apply, EfeLinearPlane};

fn assets() -> PathBuf {
    PathBuf::from("../jpegai-vm-assets")
}

fn load_json(p: &Path) -> serde_json::Value {
    serde_json::from_slice(&std::fs::read(p).unwrap_or_else(|e| panic!("{}: {e}", p.display())))
        .unwrap()
}

fn i64s(v: &serde_json::Value) -> Vec<i64> {
    v.as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect()
}

fn f32s(v: &serde_json::Value) -> Vec<f32> {
    v.as_array().unwrap().iter().map(|n| n.as_f64().unwrap() as f32).collect()
}

fn sess(p: &Path) -> Session {
    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(p)
        .unwrap_or_else(|e| panic!("load {}: {e}", p.display()))
}

/// Run a single-input ONNX model on a [1,C,H,W] f32 tensor; returns (data, shape).
fn run1(s: &mut Session, x: Vec<f32>, c: usize, h: usize, w: usize) -> (Vec<f32>, Vec<usize>) {
    let name = s.inputs().first().map(|i| i.name().to_string()).unwrap();
    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();
    let shape = view.shape().to_vec();
    (view.iter().copied().collect(), shape)
}

/// Run a two-input MCM stage net: (p, prev) — both [1,C,8,8].
fn run_mcm(s: &mut Session, p: Vec<f32>, pc: usize, prev: Vec<f32>, vc: usize) -> Vec<f32> {
    let names: Vec<String> = s.inputs().iter().map(|i| i.name().to_string()).collect();
    let pa = ndarray::Array4::from_shape_vec((1, pc, 8, 8), p).unwrap();
    let mut feeds: Vec<(std::borrow::Cow<'static, str>, ort::session::SessionInputValue<'static>)> =
        vec![(names[0].clone().into(), Value::from_array(pa).unwrap().into())];
    if names.len() > 1 {
        let va = ndarray::Array4::from_shape_vec((1, vc, 8, 8), prev).unwrap();
        feeds.push((names[1].clone().into(), Value::from_array(va).unwrap().into()));
    }
    let out = s.run(feeds).unwrap();
    out[0].try_extract_array::<f32>().unwrap().iter().copied().collect()
}

/// Split a [C,16,16] plane into 4 stage slices [C,8,8] in decode order
/// (TL, BR, TR, BL — the VM's down_shuffle part order).
fn down_shuffle(x: &[f32], c: usize) -> [Vec<f32>; 4] {
    let phase = |py: usize, px: usize| -> Vec<f32> {
        let mut out = Vec::with_capacity(c * 64);
        for ch in 0..c {
            for y in 0..8 {
                for xx in 0..8 {
                    out.push(x[ch * 256 + (2 * y + py) * 16 + 2 * xx + px]);
                }
            }
        }
        out
    };
    [phase(0, 0), phase(1, 1), phase(0, 1), phase(1, 0)]
}

/// Inverse: stage slices (decode order TL,BR,TR,BL) → [C,16,16].
fn up_shuffle(st: &[Vec<f32>; 4], c: usize) -> Vec<f32> {
    let mut out = vec![0f32; c * 256];
    for (si, (py, px)) in [(0usize, (0usize, 0usize)), (1, (1, 1)), (2, (0, 1)), (3, (1, 0))] {
        for ch in 0..c {
            for y in 0..8 {
                for xx in 0..8 {
                    out[ch * 256 + (2 * y + py) * 16 + 2 * xx + px] = st[si][ch * 64 + y * 8 + xx];
                }
            }
        }
    }
    out
}

/// LSBS scaling: σ_skip-driven table lookup (validated in round 6b).
fn lsbs(yhat: &[f32], resi_dq: &[f32], sigma_skip: &[i64], c: usize, t0: &[i64], t1: &[i64]) -> Vec<f32> {
    // likely: per-channel 8×8 block sum of the SIGNED skip surface, (Σ+32)>>6
    let mut out = vec![0f32; c * 256];
    for ch in 0..c {
        let mut blocks = [[0i64; 2]; 2];
        for y in 0..16 {
            for x in 0..16 {
                blocks[y / 8][x / 8] += sigma_skip[ch * 256 + y * 16 + x];
            }
        }
        for y in 0..16 {
            for x in 0..16 {
                let lk = ((blocks[y / 8][x / 8] + 32) >> 6).clamp(0, 4351) as usize;
                let i = ch * 256 + y * 16 + x;
                let mean_full = yhat[i] - resi_dq[i];
                let additive = (t0[lk] as f32 * mean_full + t1[lk] as f32 * resi_dq[i]) / 8192.0;
                out[i] = yhat[i] + additive;
            }
        }
    }
    out
}

/// Diagnostic: max abs diff vs a checkpoint dumped from the live reference.
fn ckdiff(cp: &serde_json::Value, key: &str, got: &[f32]) {
    if cp.get(key).is_none() {
        return;
    }
    let want = f32s(&cp[key]);
    let n = want.len().min(got.len());
    let mut worst = 0f32;
    for i in 0..n {
        worst = worst.max((got[i] - want[i]).abs());
    }
    eprintln!("    [ck] {key}: max|err| {worst:.4e} (got {} want {})", got.len(), want.len());
}

#[test]
fn jpegai_end_to_end_decode() {
    let a = assets();
    for need in ["e2e/test256.bin", "e2e/route_fixture.json", "model16/CCS_SGMM"] {
        if !a.join(need).exists() {
            eprintln!("skipping: {need} not present");
            return;
        }
    }
    let cp = load_json(&a.join("e2e/checkpoints_fixture.json"));
    // ---- fixtures ----
    let ent = load_json(&a.join("e2e/jpegai_real_entropy_fixture.json"));
    let weights = load_json(&a.join("hyper_scale_decoder_int_weights.json"));
    let adds = load_json(&a.join("e2e/sigma_additions_fixture.json"));
    let gains = load_json(&a.join("e2e/dequant_gain_map_fixture.json"));
    let route = load_json(&a.join("e2e/route_fixture.json"));
    let norm = load_json(&a.join("jpegai_normative_tables.json"));

    // ---- 1. codestream + headers ----
    let bin = std::fs::read(a.join("e2e/test256.bin")).unwrap();
    let subs = parse_codestream(&bin).unwrap();
    let payload = |m: u16| subs.iter().find(|s| s.marker == m).unwrap().payload;

    // ---- headers: typed parse (PIH + TON, separate cursors) ----
    let hdr = parse_headers(payload(0xFF82), payload(0xFF83)).expect("headers");
    assert_eq!((hdr.img_width, hdr.img_height), (256, 256));
    let consts = load_json(&a.join("model_constants.json"));

    // ---- 2–3. SOZ → z → the scale graph → σ, mask → the residuals, through the integer
    // operator set as both decode paths run them (integer-operators §13.6.1), from this
    // bundle's members: the normative PMFs and bounds, the factorized CDFs, the gains and
    // the scale network. The derived quantities are checked against the dumped ones
    // while both exist. ----
    let fac = load_json(&a.join("jpegai_factorized_cdfs.json"));
    let cdfs = |p: &str| -> Vec<i64> { fac[p].as_array().unwrap().iter().flat_map(i64s).collect() };
    for plane in ["y", "uv"] {
        assert!(i64s(&route["z_offsets"][plane]).iter().all(|&o| o == 31), "the hyper-latent offsets are §13.6.1's");
    }
    let bounds: Vec<i64> = std::fs::read_to_string(a.join("me-tANS/bounds.csv")).unwrap()
        .lines().filter(|l| !l.trim().is_empty()).map(|l| l.trim().parse().unwrap()).collect();
    let gain_of = |p: &str| -> Vec<i64> { i64s(&consts["tools_0"][format!("gain_mlog_{p}")]).iter().map(|&g| g + hdr.beta_displacement_log).collect() };
    let net = wai_ops::lower::jpegai::JpegaiNet {
        norm: wai_ops::lower::jpegai::Norm {
            pmf: norm["pmf"].as_array().unwrap().iter().map(i64s).collect(),
            bounds,
            per_channel: fac["per_channel"].as_i64().unwrap(),
            cdfs_y: cdfs("y"),
            cdfs_uv: cdfs("uv"),
            gain_y: gain_of("y"),
            gain_uv: gain_of("uv"),
        },
        hs_y: wai::jpegai_ops::scale_net_of_blocks(&weights["y"]).unwrap(),
        hs_uv: wai::jpegai_ops::scale_net_of_blocks(&weights["uv"]).unwrap(),
    };
    let sp = wai::jpegai_ops::load_net(net).expect("the bundle's scale path loads under acc32");
    let (mem_z, ends_z) = parse_ae_substream(payload(0xFF88), 1).unwrap();
    let mut sigma_final: Vec<Vec<i64>> = Vec::new();   // per plane [C*256]
    let mut sigma_skip: Vec<Vec<i64>> = Vec::new();
    let mut z_net_planes: Vec<Vec<i64>> = Vec::new();
    let mut resi: Vec<Vec<i16>> = Vec::new();
    for (pi, (plane, marker, stream, p)) in [("y", 0xFF89u16, "SORP", wai::jpegai_ops::Plane::Luma), ("uv", 0xFF8A, "SORS", wai::jpegai_ops::Plane::Chroma)].into_iter().enumerate() {
        let (mem, ends) = parse_ae_substream(payload(marker), 1).unwrap();
        let grfs: Vec<bool> = hdr.models[pi].grfs_channel_flag.iter().map(|&f| f == 1).collect();
        let ints = wai::jpegai_ops::plane(&sp, p, &mem_z[..ends_z[0] as usize], &mem[..ends[0] as usize], &grfs).expect("the plane decodes");
        let fin: Vec<i64> = ints.v.iter().zip(&ints.rvs).map(|(v, r)| v + r).collect();
        let rvs_fix = i64s(&adds[plane]["rvs"]);
        for (i, &f) in fin.iter().enumerate() {
            assert_eq!(f, ints.v[i] + rvs_fix[i], "derived rvs mismatch at {i}");
        }
        let calls = ent[stream]["calls"].as_array().unwrap();
        for (c, call) in calls.iter().enumerate() {
            let fx: Vec<bool> = i64s(&call["masks"]).iter().map(|&v| v != 0).collect();
            assert_eq!(&ints.mask[c * 256..(c + 1) * 256], &fx[..], "derived skip mask != fixture at {stream} ch {c}");
            let want: Vec<i64> = i64s(&call["x"]).iter().map(|&v| v as i16 as i64).collect();
            assert_eq!(&ints.y_q[c * 256..(c + 1) * 256], &want[..], "{stream} ch {c}: the residuals");
        }
        resi.push(ints.y_q.iter().map(|&q| q as i16).collect());
        sigma_final.push(fin);
        z_net_planes.push(ints.z);
        sigma_skip.push(ints.v);
    }

    // ---- 4. dequant (per-position gain map) ----
    let mut resi_dq: Vec<Vec<f32>> = Vec::new();
    for (pi, plane) in ["y", "uv"].iter().enumerate() {
        let g = f32s(&gains[plane]["gain"]);
        resi_dq.push(resi[pi].iter().zip(&g).map(|(&r, &gv)| r as f32 * gv).collect());
    }

    // ---- 5. hyper decoder (ORT) → ψ ----
    let mut psi: Vec<Vec<f32>> = Vec::new();
    for (pi, (plane, c)) in [("y", 160usize), ("uv", 96)].iter().enumerate() {
        let mut s = sess(&a.join(format!("model/CCS_SGMM/tools_0/model_{plane}/common_modules/hyper_decoder.onnx")));
        let zf: Vec<f32> = z_net_planes[pi].iter().map(|&v| v as f32).collect();
        let (out, _) = run1(&mut s, zf, *c, 4, 4);
        psi.push(out);
    }

    // ---- 6. value path: Y via MCM stages, UV via ψ up-shuffle ----
    let resi_y_st = down_shuffle(&resi_dq[0], 160);
    let mut yhat_st: Vec<Vec<f32>> = Vec::new();
    for s_i in 0..4 {
        let mut st = sess(&a.join(format!("model/CCS_SGMM/tools_0/model_y/common_modules/MCM/stage{s_i}.onnx")));
        let p: Vec<f32> = psi[0][s_i * 160 * 64..(s_i + 1) * 160 * 64].to_vec();
        let prev: Vec<f32> = yhat_st.iter().flatten().copied().collect();
        let prev_c = 160 * yhat_st.len().max(1);
        let prev = if yhat_st.is_empty() { vec![0f32; 160 * 64] } else { prev };
        let mu = run_mcm(&mut st, p, 160, prev, prev_c);
        yhat_st.push(resi_y_st[s_i].iter().zip(&mu).map(|(a, b)| a + b).collect());
    }
    let yhat_y = up_shuffle(&[yhat_st[0].clone(), yhat_st[1].clone(), yhat_st[2].clone(), yhat_st[3].clone()], 160);
    ckdiff(&cp, "resi_dq_y", &resi_dq[0]);
    ckdiff(&cp, "psi_y", &psi[0]);
    ckdiff(&cp, "yhat_y", &yhat_y);
    // UV: μ = up_shuffle of ψ chunks (decode order = chunks 0,2,3,1 → TL,BR,TR,BL)
    let ch = |i: usize| psi[1][i * 96 * 64..(i + 1) * 96 * 64].to_vec();
    let mu_uv = up_shuffle(&[ch(0), ch(1), ch(2), ch(3)], 96);
    let yhat_uv: Vec<f32> = resi_dq[1].iter().zip(&mu_uv).map(|(a, b)| a + b).collect();
    ckdiff(&cp, "yhat_uv", &yhat_uv);

    // ---- 7. LSBS ----
    let t0 = i64s(&route["lsbs_t0"]);
    let t1 = i64s(&route["lsbs_t1"]);
    let y_ls = lsbs(&yhat_y, &resi_dq[0], &sigma_skip[0], 160, &t0, &t1);
    let uv_ls = lsbs(&yhat_uv, &resi_dq[1], &sigma_skip[1], 96, &t0, &t1);
    ckdiff(&cp, "y_lsbs", &y_ls);
    ckdiff(&cp, "uv_lsbs", &uv_ls);

    // ---- 8. synthesis (ORT, 16×16-traced) ----
    let mut sy = sess(&a.join("model16/CCS_SGMM/tools_0/model_y/synthesis16.onnx"));
    let (ypix, _) = run1(&mut sy, y_ls.clone(), 160, 16, 16);          // [1,1,256,256]
    let mut suv = sess(&a.join("model16/CCS_SGMM/tools_0/model_uv/synthesis16.onnx"));
    let mut uv_in = y_ls;
    uv_in.extend_from_slice(&uv_ls);
    let (uvpix, _) = run1(&mut suv, uv_in, 256, 16, 16);               // [1,2,256,256]
    let (h, w) = (256usize, 256usize);
    let mut yp = ypix;
    let mut up: Vec<f32> = uvpix[0..h * w].to_vec();
    let mut vp: Vec<f32> = uvpix[h * w..2 * h * w].to_vec();
    ckdiff(&cp, "ypix", &yp);
    ckdiff(&cp, "uvpix", &uvpix);

    // ---- 9. EFElinear (parameters straight from the tool header) ----
    let deint = |q: u64| (q as f32 - 32767.0) / 2048.0;
    assert!(hdr.efe_linear_enabled);
    for (ci, wq, w2q, mean_q) in [
        (1usize, hdr.efe_linear.weights_u[0], hdr.efe_linear.weights_u2[0], hdr.efe_linear.mean_u_q100),
        (2, hdr.efe_linear.weights_v[0], hdr.efe_linear.weights_v2[0], hdr.efe_linear.mean_v_q100),
    ] {
        let mean = mean_q as f32 / 100.0;
        let p = EfeLinearPlane {
            mean,
            w_uv: [deint(wq) + 1.0; 4],
            bias: [mean; 4],
            w_y: [deint(w2q); 4],
        };
        let src = if ci == 1 { &up } else { &vp };
        let out = efe_linear_plane(&yp, src, h, w, &p);
        if ci == 1 { up = out } else { vp = out }
    }

    {
        let mut all = yp.clone(); all.extend_from_slice(&up); all.extend_from_slice(&vp);
        ckdiff(&cp, "efl_out", &all);
    }
    // ---- 10. eICCI (ORT wrappers; [0,1] domain) ----
    let mut yuv01 = Vec::with_capacity(3 * h * w);
    for pl in [&yp, &up, &vp] {
        yuv01.extend(pl.iter().map(|&v| v / 255.0));
    }
    let crop = |data: &[f32], shape: &[usize], ch: usize| -> Vec<f32> {
        let (oh, ow) = (shape[2], shape[3]);
        let mut out = Vec::with_capacity(h * w);
        for y in 0..h {
            for x in 0..w {
                out.push(data[(ch * oh + y) * ow + x] * 255.0);
            }
        }
        out
    };
    let mut si_y = sess(&a.join("e2e/pfchain/icci_y_m6.onnx"));
    let (fy, fysh) = run1(&mut si_y, yuv01.clone(), 3, h, w);
    let mut si_uv = sess(&a.join("e2e/pfchain/icci_uv_m3.onnx"));
    let (fuv, fuvsh) = run1(&mut si_uv, yuv01, 3, h, w);
    yp = crop(&fy, &fysh, 0);
    up = crop(&fuv, &fuvsh, 0);
    vp = crop(&fuv, &fuvsh, 1);

    {
        let mut all = yp.clone(); all.extend_from_slice(&up); all.extend_from_slice(&vp);
        ckdiff(&cp, "icci_out", &all);
    }
    // ---- 11. EFEnonlinear (parameters from the tool header) ----
    assert!(hdr.efe_nonlinear_enabled && hdr.efn_v_enabled && !hdr.efn_u_enabled);
    let wq: Vec<i64> = hdr.efn_weights.iter().map(|&v| v as i64).collect();
    vp = efe_nonlinear_plane(&yp, &vp,
        &[wq[0], wq[1], wq[2], wq[3], wq[4], wq[5], wq[6], wq[7]],
        hdr.efn_min_luma_q100 as i64, hdr.efn_max_luma_q100 as i64);

    {
        let mut all = yp.clone(); all.extend_from_slice(&up); all.extend_from_slice(&vp);
        ckdiff(&cp, "efn_out", &all);
    }
    // ---- 12. LEF (op-point tables at header model_id; σ channel from header) ----
    assert!(hdr.lef_enabled);
    let mid = hdr.model_id as usize;
    let thr_t = i64s(&consts["lef_thr"][mid]);
    let mag_t = f32s(&consts["lef_mag"][mid]);
    let c0 = hdr.lef_ch_idx as usize * 256;
    yp = lef_apply(&yp, &sigma_final[0][c0..c0 + 256], h, w,
        [thr_t[0], thr_t[1], thr_t[2]], [mag_t[0], mag_t[1], mag_t[2]]);

    {
        let mut all = yp.clone(); all.extend_from_slice(&up); all.extend_from_slice(&vp);
        ckdiff(&cp, "lef_out", &all);
    }
    // ---- 13. BT.709 → RGB → compare with reference PNG ----
    let refimg = image::open(a.join("e2e/test256_dec.png")).unwrap().to_rgb8();
    let (kr, kg, kb): (f32, f32, f32) = (0.2126, 0.7152, 0.0722);
    let (kby, kry): (f32, f32) = (1.8556, 1.5748);
    let mut mismatched = 0usize;
    let mut maxdiff = 0i32;
    for i in 0..h * w {
        let y01 = yp[i] / 255.0;
        let u01 = up[i] / 255.0 - 0.5;
        let v01 = vp[i] / 255.0 - 0.5;
        let rgb = [
            y01 + kry * v01,
            y01 - (kb * kby / kg) * u01 - (kr * kry / kg) * v01,
            y01 + kby * u01,
        ];
        let px = refimg.get_pixel((i % w) as u32, (i / w) as u32);
        for (c, &val) in rgb.iter().enumerate() {
            let got = (val * 255.0).round_ties_even().clamp(0.0, 255.0) as i32;
            let d = (got - px[c] as i32).abs();
            if d > 0 { mismatched += 1; }
            maxdiff = maxdiff.max(d);
        }
    }
    let total = 3 * h * w;
    println!("  e2e decode: {mismatched}/{total} samples differ, max diff {maxdiff}");
    assert!(maxdiff <= 1, "pixel deviation beyond rounding: {maxdiff}");
    assert!(mismatched * 1000 <= total, "more than 0.1% mismatched: {mismatched}/{total}");
}