//! The rate and distortion of a residual against the exact upscale (extensions/exact-upscale.md
//! §4, §5), on the pictures `tools/wai_int_sr_train.py` writes.
//!
//!   wai_upscale_measure <dir> [q ...]
//!
//! `<dir>` is the training tool's output: `pictures/manifest.json` and its pictures, and
//! `priors/generic.wsr` and `priors/<title>.wsr`. For each step `q` (default 1 2 4 8 16) and
//! each picture, it codes the full-resolution picture four ways, every one with the repository's
//! own integer coders and every one decoded back and checked:
//!
//! - `direct`: the picture as a one-frame `WIV1` keyframe-0 payload (integer-payloads §3) at `q`;
//! - `up2`: the half-resolution picture (each 2 × 2 block's rounded mean) as `WIV1` at `q`, then a
//!   `WIR1` layer with the up2 predictor (staged-delivery Appendix A) at `q`;
//! - `sr`: the same half-resolution picture, `wai-upscale:sr2` of its decode under the generic
//!   prior, then a `WIR1` identity residual (exact-upscale §4.1) at `q`;
//! - `sr-title`: as `sr` under the title's own prior, whose bytes are counted once for the title.
//!
//! Every way quantises the last step's residual with the same `q` and the same rounding, so each
//! sample of every output lies within ⌊q/2⌋ of the picture's: that is the matched quality. Bytes
//! are payload bytes (staged-delivery §11): envelopes, receipts and records are not counted, and
//! the generic prior, which a sink holds once for every title, is not counted per title. The
//! output is one JSON object: per group (a title, or the held-out pictures), per `q`, each way's
//! bytes, PSNR and largest sample difference, and the PSNR of each upscale with no residual.

use std::collections::BTreeMap;
use std::path::Path;

use serde_json::{json, Value};
use wai::exact_upscale::encode_residual;
use wai::int_sr;
use wai::int_video_enc::{encode, EncParams};
use wai::staged::wir1::{self, Below, Next, Predictor, RefineParams, UnitKind};

fn box_down(px: &[u8], w: usize, h: usize) -> Vec<u8> {
    let mut out = Vec::with_capacity(w / 2 * h / 2 * 3);
    for y in (0..h).step_by(2) {
        for x in (0..w).step_by(2) {
            for c in 0..3 {
                let at = |yy: usize, xx: usize| px[(yy * w + xx) * 3 + c] as u32;
                out.push(((at(y, x) + at(y, x + 1) + at(y + 1, x) + at(y + 1, x + 1) + 2) >> 2) as u8);
            }
        }
    }
    out
}

fn sse(a: &[u8], b: &[u8]) -> u64 {
    a.iter().zip(b).map(|(x, y)| (*x as i64 - *y as i64).pow(2) as u64).sum()
}

fn max_diff(a: &[u8], b: &[u8]) -> u8 {
    a.iter().zip(b).map(|(x, y)| x.abs_diff(*y)).max().unwrap_or(0)
}

fn psnr(sse: u64, n: usize) -> f64 {
    if sse == 0 { 99.0 } else { 10.0 * ((255.0f64 * 255.0 * n as f64) / sse as f64).log10() }
}

fn wiv1(px: &[u8], w: usize, h: usize, q: u32) -> (usize, Vec<u8>) {
    let e = encode(&[px.to_vec()], h as u32, w as u32, &EncParams { block: 1, search: 0, mv_precision: 1, q }).expect("WIV1 encodes");
    (e.wiv1.len(), e.recon.into_iter().next().expect("one frame"))
}

fn up2_layer(lr: &[u8], hr: &[u8], w: usize, h: usize, q: u32) -> (usize, Vec<u8>) {
    let p = RefineParams { q, predictor: Predictor::Up2, rate: (0, 0), table_per_unit: false, max_symbols: 4095 };
    let (payload, next) = wir1::encode(UnitKind::Rgb8, &[Below::U8(lr)], ((h / 2) as u32, (w / 2) as u32, 3), &[Below::U8(hr)], (h as u32, w as u32, 3), &p).expect("WIR1 encodes");
    let Some(Next::U8(out)) = next.into_iter().next() else { panic!("one unit") };
    (payload.len(), out)
}

/// One way's figures over a group, summed over its pictures.
#[derive(Default)]
struct Sum {
    bytes: usize,
    sse: u64,
    n: usize,
    max: u8,
}

impl Sum {
    fn add(&mut self, bytes: usize, out: &[u8], hr: &[u8]) {
        self.bytes += bytes;
        self.sse += sse(out, hr);
        self.n += hr.len();
        self.max = self.max.max(max_diff(out, hr));
    }
    fn json(&self) -> Value {
        json!({"bytes": self.bytes, "psnr_db": (psnr(self.sse, self.n) * 1000.0).round() / 1000.0, "max_difference": self.max})
    }
}

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let Some(dir) = a.get(1).map(Path::new) else {
        eprintln!("usage: wai_upscale_measure <dir> [q ...]");
        std::process::exit(2);
    };
    let qs: Vec<u32> = if a.len() > 2 { a[2..].iter().map(|s| s.parse().expect("a q")).collect() } else { vec![1, 2, 4, 8, 16] };
    let manifest: Value = serde_json::from_slice(&std::fs::read(dir.join("pictures/manifest.json")).expect("manifest")).expect("JSON");
    let generic_bytes = std::fs::read(dir.join("priors/generic.wsr")).expect("the generic prior");
    let generic = int_sr::load(&generic_bytes).expect("the generic prior loads");
    let mut groups: BTreeMap<String, Vec<(String, usize, usize)>> = BTreeMap::new();
    for p in manifest["pictures"].as_array().expect("pictures") {
        let g = p["group"].as_str().expect("group").to_string();
        groups.entry(g).or_default().push((p["name"].as_str().unwrap().into(), p["width"].as_u64().unwrap() as usize, p["height"].as_u64().unwrap() as usize));
    }
    let mut report = serde_json::Map::new();
    report.insert("generic_prior".into(), json!({"bytes": generic_bytes.len(), "sha256": wai::pin::sha256_hex(&generic_bytes), "acc32": int_sr::meets_acc32(&generic_bytes)}));
    for (group, pics) in &groups {
        let title_bytes = std::fs::read(dir.join(format!("priors/{group}.wsr"))).ok();
        let title = title_bytes.as_ref().map(|b| int_sr::load(b).expect("a title prior loads"));
        let mut g = serde_json::Map::new();
        if let Some(b) = &title_bytes {
            g.insert("title_prior".into(), json!({"bytes": b.len(), "sha256": wai::pin::sha256_hex(b), "acc32": int_sr::meets_acc32(b)}));
        }
        g.insert("pictures".into(), json!(pics.len()));
        // The upscales with no residual, from the lossless half-resolution picture.
        let (mut s_up2, mut s_gen, mut s_title) = (Sum::default(), Sum::default(), Sum::default());
        let pictures: Vec<(Vec<u8>, usize, usize)> = pics
            .iter()
            .map(|(name, w, h)| (std::fs::read(dir.join(format!("pictures/{name}.rgb"))).expect("a picture"), *w, *h))
            .collect();
        for (hr, w, h) in &pictures {
            let lr = box_down(hr, *w, *h);
            s_up2.add(0, &int_sr::up2(&lr, w / 2, h / 2).unwrap(), hr);
            let gd = int_sr::sr2(&generic, &lr, w / 2, h / 2).unwrap();
            // The graph and the loop-by-loop arithmetic agree on every measured picture.
            assert_eq!(gd, int_sr::direct(&int_sr::parse(&generic_bytes).unwrap(), &lr, w / 2, h / 2).unwrap());
            s_gen.add(0, &gd, hr);
            if let Some(t) = &title {
                s_title.add(0, &int_sr::sr2(t, &lr, w / 2, h / 2).unwrap(), hr);
            }
        }
        g.insert(
            "upscale_only".into(),
            json!({"up2_psnr_db": s_up2.json()["psnr_db"], "sr_psnr_db": s_gen.json()["psnr_db"], "sr_title_psnr_db": if title.is_some() { s_title.json()["psnr_db"].clone() } else { Value::Null }}),
        );
        let mut by_q = serde_json::Map::new();
        for &q in &qs {
            let (mut direct, mut up2, mut sr, mut sr_title) = (Sum::default(), Sum::default(), Sum::default(), Sum::default());
            for (hr, w, h) in &pictures {
                let (b, out) = wiv1(hr, *w, *h, q);
                direct.add(b, &out, hr);
                let lr = box_down(hr, *w, *h);
                let (lb, lr_dec) = wiv1(&lr, w / 2, h / 2, q);
                let (b2, out2) = up2_layer(&lr_dec, hr, *w, *h, q);
                up2.add(lb + b2, &out2, hr);
                let up = int_sr::sr2(&generic, &lr_dec, w / 2, h / 2).unwrap();
                let (r, out3) = encode_residual(&up, hr, *w as u32, *h as u32, q).unwrap();
                sr.add(lb + r.len(), &out3, hr);
                if let Some(t) = &title {
                    let up = int_sr::sr2(t, &lr_dec, w / 2, h / 2).unwrap();
                    let (r, out4) = encode_residual(&up, hr, *w as u32, *h as u32, q).unwrap();
                    sr_title.add(lb + r.len(), &out4, hr);
                }
            }
            let mut row = serde_json::Map::new();
            row.insert("direct".into(), direct.json());
            row.insert("up2".into(), up2.json());
            row.insert("sr".into(), sr.json());
            if let Some(b) = &title_bytes {
                sr_title.bytes += b.len();
                row.insert("sr_title".into(), sr_title.json());
            }
            by_q.insert(format!("q{q}"), Value::Object(row));
        }
        g.insert("rate".into(), Value::Object(by_q));
        report.insert(group.clone(), Value::Object(g));
    }
    println!("{}", serde_json::to_string_pretty(&Value::Object(report)).unwrap());
}
