//! Measure an encoder: its speed and energy over an odd number of rounds of one encode,
//! and the quality of what it made against what it was handed (extensions/platform-encode.md
//! §10). Informative: nothing here is a conformance check, and no figure is compared with
//! one taken under another counter.
//!
//! ```text
//! wai_encode_measure run --api <a> --codec avc|hevc --source <in.y4m> --kbps <n> --out <dir> --name <n>
//!     [--size WxH] [--idr-every <frames>] [--rounds <odd n>] [--counter <c>] [--loops <n>]
//!     [--library <p>] [--device <d>] [--hardware yes|no] [--key <seed>]
//! wai_encode_measure quality --codec avc|hevc --stream <es> --source <in.y4m> [--size WxH]   (macOS)
//! ```
//!
//! `run` writes `<dir>/<name>.h264|.h265`, `<dir>/<name>.record.json` (signed with `--key`,
//! or a key of 32 bytes of 0x45) and prints one JSON line of results. `quality` decodes the
//! stream with the platform's decoder and prints PSNR (luma, each chroma plane, and the
//! 6:1:1 weighted mean of their squared errors), SSIM of luma over 8×8 windows at a step of
//! 4 samples, and the SHA-256 of every decoded picture, one after another.

#![allow(clippy::float_arithmetic)]

use serde_json::json;
use std::path::Path;
use wai::platform_encode::{self as pe, Api, Codec, Options, record};

fn opt<'a>(args: &'a [String], flag: &str) -> Option<&'a str> {
    args.iter().position(|a| a == flag).and_then(|i| args.get(i + 1).map(String::as_str))
}

fn need<'a>(args: &'a [String], flag: &str) -> &'a str {
    opt(args, flag).unwrap_or_else(|| {
        eprintln!("{flag} is required");
        std::process::exit(2)
    })
}

fn codec(s: &str) -> Codec {
    match s {
        "avc" => Codec::Avc,
        "hevc" => Codec::Hevc,
        _ => {
            eprintln!("--codec avc|hevc");
            std::process::exit(2)
        }
    }
}

fn size(s: Option<&str>) -> Option<(u32, u32)> {
    s.map(|s| {
        let (w, h) = s.split_once('x').expect("WxH");
        (w.parse().expect("W"), h.parse().expect("H"))
    })
}

fn run(args: &[String]) -> Result<(), String> {
    let api = Api::from_name(need(args, "--api")).ok_or("--api platform|va|gpu-driver|software")?;
    let codec = codec(need(args, "--codec"));
    let kbps: u32 = need(args, "--kbps").parse().map_err(|_| "--kbps")?;
    let src = std::fs::read(need(args, "--source")).map_err(|e| e.to_string())?;
    let mut p = pe::prepare(&src, codec, size(opt(args, "--size")), kbps).map_err(|e| e.to_string())?;
    drop(src);
    let loops: usize = opt(args, "--loops").map_or(1, |l| l.parse().expect("--loops"));
    if loops > 1 {
        let one = p.frames.clone();
        for _ in 1..loops {
            p.frames.extend(one.iter().cloned());
        }
        p.source.frames *= loops as u64;
    }
    let frames = p.frames.len() as u64;
    let every: u64 = opt(args, "--idr-every").map_or(2 * u64::from(p.cfg.fps_num) / u64::from(p.cfg.fps_den), |e| e.parse().expect("--idr-every"));
    let idr = pe::idr_every(frames, every);
    let hardware = match opt(args, "--hardware") {
        Some("yes") => Some(true),
        Some("no") => Some(false),
        _ => None,
    };
    let opts = Options { library: opt(args, "--library").map(Into::into), device: opt(args, "--device").map(String::from), hardware };
    let rounds: usize = opt(args, "--rounds").map_or(1, |r| r.parse().expect("--rounds"));
    let mut walls = Vec::new();
    let mut digests = Vec::new();
    let encode_once = |walls: &mut Vec<u64>, digests: &mut Vec<String>| -> Result<pe::Encoded, String> {
        let e = pe::encode(api, &p.cfg, &opts, p.frames.iter().map(Vec::as_slice), &idr).map_err(|e| e.to_string())?;
        walls.push(e.wall_us);
        digests.push(pe::sha256_hex(&e.stream));
        Ok(e)
    };
    let (energy, round_info, enc) = match opt(args, "--counter") {
        Some(c) => {
            let mut counter = pe::energy::open(c)?;
            let (e, r, enc) = pe::energy::measure(counter.as_mut(), rounds, std::time::Duration::from_millis(2000), || encode_once(&mut walls, &mut digests))?;
            (e, r, enc)
        }
        None => {
            let mut last = None;
            for _ in 0..rounds {
                last = Some(encode_once(&mut walls, &mut digests)?);
            }
            (None, Vec::new(), last.ok_or("no round")?)
        }
    };
    let out = Path::new(need(args, "--out"));
    std::fs::create_dir_all(out).map_err(|e| e.to_string())?;
    let name = need(args, "--name");
    let es = out.join(format!("{name}.{}", codec.extension()));
    std::fs::write(&es, &enc.stream).map_err(|e| e.to_string())?;
    let seed = match opt(args, "--key") {
        Some(k) => {
            let t = std::fs::read_to_string(k).map_err(|e| e.to_string())?;
            let t = t.trim();
            let mut s = [0u8; 32];
            for (i, b) in s.iter_mut().enumerate() {
                *b = u8::from_str_radix(&t[2 * i..2 * i + 2], 16).map_err(|e| e.to_string())?;
            }
            s
        }
        None => [0x45; 32],
    };
    let key = ed25519_dalek::SigningKey::from_bytes(&seed);
    let body = record::body(&enc, &p.cfg, &p.source, &idr, energy.as_ref(), None, &record::now_rfc3339());
    let signed = record::sign(body, &key);
    let text = record::to_text(&signed);
    record::verify(&text, &record::Held { stream: Some(&enc.stream), source: None }).map_err(|r| format!("the record does not verify: {r}"))?;
    std::fs::write(out.join(format!("{name}.record.json")), &text).map_err(|e| e.to_string())?;
    let mut sorted = walls.clone();
    sorted.sort();
    let median_us = sorted[sorted.len() / 2];
    let secs = frames as f64 * f64::from(p.cfg.fps_den) / f64::from(p.cfg.fps_num);
    let same = digests.windows(2).all(|w| w[0] == w[1]);
    let line = json!({
        "name": name, "api": api.name(), "codec": codec.name(), "width": p.cfg.width, "height": p.cfg.height,
        "frames": frames, "kbps_target": kbps, "bytes": enc.stream.len(),
        "kbps_made": (enc.stream.len() as f64 * 8.0 / secs / 1000.0 * 10.0).round() / 10.0,
        "rounds": rounds, "wall_us": walls, "median_wall_us": median_us,
        "fps_median": (frames as f64 * 1e6 / median_us as f64 * 10.0).round() / 10.0,
        "rounds_same_bytes": same, "stream_sha256": pe::sha256_hex(&enc.stream),
        "identity": {"encoder": enc.identity.encoder, "version": enc.identity.version, "hardware": enc.identity.hardware, "device": enc.identity.device, "driver": enc.identity.driver, "library_sha256": enc.identity.library_sha256},
        "energy": energy.as_ref().map(record::Energy::to_json),
        "energy_rounds": round_info.iter().map(|r| json!({"uj": r.uj.as_ref().ok(), "why": r.uj.as_ref().err(), "busy_us": r.busy_us, "idle_us": r.idle_us})).collect::<Vec<_>>(),
        "uj_per_frame": energy.as_ref().map(|e| e.joules_micro / frames),
        "record_digest": record::digest(&signed),
    });
    println!("{line}");
    Ok(())
}

#[cfg_attr(not(target_vendor = "apple"), allow(dead_code))]
/// Mean squared error of two planes, as a sum and a count.
fn sse(a: &[u8], b: &[u8]) -> (u64, u64) {
    (a.iter().zip(b).map(|(&x, &y)| (i64::from(x) - i64::from(y)).pow(2) as u64).sum(), a.len() as u64)
}

#[cfg_attr(not(target_vendor = "apple"), allow(dead_code))]
fn psnr(sum: u64, n: u64) -> f64 {
    if sum == 0 { 100.0 } else { 10.0 * (255.0f64 * 255.0 / (sum as f64 / n as f64)).log10() }
}

#[cfg_attr(not(target_vendor = "apple"), allow(dead_code))]
/// SSIM of one plane over 8×8 windows at a step of 4 (Wang et al. 2004, with
/// C1 = (0.01·255)², C2 = (0.03·255)²), the mean over windows.
fn ssim(a: &[u8], b: &[u8], w: usize, h: usize) -> (f64, u64) {
    let (c1, c2) = ((0.01f64 * 255.0).powi(2), (0.03f64 * 255.0).powi(2));
    let (mut total, mut n) = (0.0f64, 0u64);
    let mut y = 0;
    while y + 8 <= h {
        let mut x = 0;
        while x + 8 <= w {
            let (mut sa, mut sb, mut saa, mut sbb, mut sab) = (0u64, 0u64, 0u64, 0u64, 0u64);
            for r in y..y + 8 {
                for c in x..x + 8 {
                    let (p, q) = (u64::from(a[r * w + c]), u64::from(b[r * w + c]));
                    sa += p;
                    sb += q;
                    saa += p * p;
                    sbb += q * q;
                    sab += p * q;
                }
            }
            let k = 64.0f64;
            let (ma, mb) = (sa as f64 / k, sb as f64 / k);
            let va = saa as f64 / k - ma * ma;
            let vb = sbb as f64 / k - mb * mb;
            let cov = sab as f64 / k - ma * mb;
            total += ((2.0 * ma * mb + c1) * (2.0 * cov + c2)) / ((ma * ma + mb * mb + c1) * (va + vb + c2));
            n += 1;
            x += 4;
        }
        y += 4;
    }
    (total, n)
}

#[cfg(target_vendor = "apple")]
fn quality(args: &[String]) -> Result<(), String> {
    let codec = codec(need(args, "--codec"));
    let stream = std::fs::read(need(args, "--stream")).map_err(|e| e.to_string())?;
    let src = std::fs::read(need(args, "--source")).map_err(|e| e.to_string())?;
    let p = pe::prepare(&src, codec, size(opt(args, "--size")), 1).map_err(|e| e.to_string())?;
    let n = p.frames.len() as u64;
    let checked = pe::annexb::check_stream(codec, &stream, n, &[0]).or_else(|_| {
        // The plan is not known here: read the IDRs from the stream itself.
        let nals = pe::annexb::nal_units(codec, &stream).map_err(|r| r.detail.clone())?;
        let aus = pe::annexb::access_units(codec, &stream, &nals).map_err(|r| r.detail.clone())?;
        let idr: Vec<u64> = aus.iter().enumerate().filter(|(_, au)| nals[(*au).clone()].iter().any(|x| match codec {
            Codec::Avc => x.kind == 5,
            Codec::Hevc => matches!(x.kind, 19 | 20),
        })).map(|(i, _)| i as u64).collect();
        pe::annexb::check_stream(codec, &stream, n, &idr).map_err(|r| format!("{}: {}", r.code, r.detail))
    })?;
    let (w, h) = (p.cfg.width as usize, p.cfg.height as usize);
    let (yl, cl) = (w * h, (w / 2) * (h / 2));
    let (mut sy, mut su, mut sv, mut ss, mut sn) = ((0u64, 0u64), (0u64, 0u64), (0u64, 0u64), 0.0f64, 0u64);
    let mut i = 0usize;
    let mut all = sha2::Sha256::default();
    use sha2::Digest as _;
    let count = pe::platform_decode(codec, &stream, &checked.units, |f| {
        if f.len() != yl + 2 * cl {
            return Err(format!("a decoded picture of {} bytes", f.len()));
        }
        let s = p.frames.get(i).ok_or("more pictures than the source")?;
        let a = sse(&s[..yl], &f[..yl]);
        let b = sse(&s[yl..yl + cl], &f[yl..yl + cl]);
        let c = sse(&s[yl + cl..], &f[yl + cl..]);
        sy = (sy.0 + a.0, sy.1 + a.1);
        su = (su.0 + b.0, su.1 + b.1);
        sv = (sv.0 + c.0, sv.1 + c.1);
        let (t, k) = ssim(&s[..yl], &f[..yl], w, h);
        ss += t;
        sn += k;
        all.update(f);
        i += 1;
        Ok(())
    })?;
    if count != n {
        return Err(format!("the decoder gave {count} pictures for {n} access units"));
    }
    // The 6:1:1 weighted mean of the planes' mean squared errors.
    let mse = |s: (u64, u64)| s.0 as f64 / s.1 as f64;
    let wmse = (6.0 * mse(sy) + mse(su) + mse(sv)) / 8.0;
    let line = json!({
        "pictures": count, "width": w, "height": h,
        "psnr_y": (psnr(sy.0, sy.1) * 1000.0).round() / 1000.0,
        "psnr_u": (psnr(su.0, su.1) * 1000.0).round() / 1000.0,
        "psnr_v": (psnr(sv.0, sv.1) * 1000.0).round() / 1000.0,
        "psnr_yuv_611": ((10.0 * (255.0f64 * 255.0 / wmse).log10()) * 1000.0).round() / 1000.0,
        "ssim_y": (ss / sn as f64 * 100000.0).round() / 100000.0,
        "decoded_sha256": pe::hex(&all.finalize()),
    });
    println!("{line}");
    Ok(())
}

#[cfg(not(target_vendor = "apple"))]
fn quality(_: &[String]) -> Result<(), String> {
    Err("quality decodes with the platform's decoder, which this build has not".into())
}

fn main() {
    let a: Vec<String> = std::env::args().skip(1).collect();
    let r = match a.first().map(String::as_str) {
        Some("run") => run(&a[1..]),
        Some("quality") => quality(&a[1..]),
        _ => Err("wai_encode_measure run|quality …".into()),
    };
    if let Err(e) = r {
        eprintln!("error: {e}");
        std::process::exit(1);
    }
}
