//! The equivalence checker (`engine::miter`) on a published benchmark set:
//! each circuit against its compilations to three native gate sets, made by
//! the suite's own compiler and by ours, equal and broken one way each; the
//! time and peak memory of every check, and one sealed receipt.
//!
//! ```text
//! equiv_bench DIR [--receipt OUT.json] [--only NAME] [--min-n N] [--max-n N] [--heavy-hex CONFIG.json]
//!                 [--limit-secs S] [--no-ours] [budget flags]
//! equiv_bench --one SPEC.json        (one check, in its own process: the driver runs each so)
//! ```
//!
//! `DIR` holds `manifest.json` and the circuits it names: each benchmark at
//! each size, target-independent (the source), and mapped by the suite's own
//! compiler to three devices with each one's initial and final layout
//! (`wai/wai-quantum/tests/data/equivalence/README.md`). Our own compiles are
//! of the source, to `rz sx x cz` on a heavy-hex map (the coupling of the
//! configuration given, else a grid), to `prx cz` on a 6×9 grid and to
//! `gpi gpi2 zz` on 36 qubits all coupled.
//!
//! Broken one way each, as the published comparison does: one gate removed
//! (a seeded choice), and one rotation turned by an eighth of a turn. Every
//! broken compile must be refuted.
//!
//! The receipt seals the inputs' digests, the budget, every check's verdict
//! (with its routes), its time and its peak memory, and the published numbers
//! beside ours. Times and memory are this host's; the verdicts are the same
//! on every host.

use serde_json::{Value, json};
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};
use wai_quantum::ed25519_dalek::SigningKey;
use wai_quantum::engine::miter::{self, Budget};
use wai_quantum::quantum_toolchain::native::{self, Basis, Coupling};

#[path = "equiv_bench/common.rs"]
mod common;
use common::{break_native, break_text, brief};

/// The published times for compiled circuits of these benchmarks (Peham,
/// Burgholzer and Wille, IEEE JETCAS 12(3), 2022, Table III, "Compiled
/// circuits", equivalent instances; a 4.2 GHz desktop CPU, 32 GiB, 1 h limit):
/// name, qubits, |G|, |G'|, the rewriting checker's seconds and the decision
/// diagram checker's (`None` past the hour).
/// A published row: name, qubits, |G|, |G'|, and the two checkers' seconds.
type Published = (&'static str, u32, u32, u32, Option<f64>, Option<f64>);

const PUBLISHED: &[Published] = &[
    ("grover", 6, 1606, 2803, Some(0.39), Some(3.40)),
    ("grover", 7, 4732, 8476, Some(1.24), Some(0.30)),
    ("grover", 8, 12482, 22860, Some(12.15), Some(0.91)),
    ("qft", 23, 1311, 3741, Some(0.06), Some(2.00)),
    ("qft", 38, 3591, 10449, Some(0.32), None),
    ("qwalk", 7, 6523, 8955, Some(150.36), Some(0.24)),
    ("qwalk", 8, 14084, 19755, Some(1289.13), Some(0.57)),
    ("qwalk", 9, 29325, 41942, None, Some(1.31)),
    ("qpeexact", 22, 1217, 3006, Some(0.83), Some(0.10)),
    ("qpeexact", 39, 3823, 11552, Some(3.19), None),
    ("ghz", 65, 130, 493, Some(0.06), Some(0.005)),
    ("graphstate", 62, 403, 2041, Some(0.36), Some(0.17)),
];

fn fail(msg: &str) -> ! {
    eprintln!("equiv_bench: {msg}");
    std::process::exit(2);
}

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

fn flag(args: &[String], name: &str) -> bool {
    args.iter().any(|a| a == name)
}

fn num<T: std::str::FromStr>(args: &[String], name: &str, d: T) -> T {
    opt(args, name).map_or(d, |s| s.parse().unwrap_or_else(|_| fail(&format!("{name} {s}: not a number"))))
}

fn budget(args: &[String]) -> Budget {
    let d = Budget::default();
    Budget {
        stimuli: num(args, "--stimuli", d.stimuli),
        sv_qubits: num(args, "--sv-qubits", d.sv_qubits),
        dd_nodes: num(args, "--nodes", d.dd_nodes),
        dd_steps: num(args, "--steps", d.dd_steps),
        zx_steps: num(args, "--zx", d.zx_steps),
        exact_bits: num(args, "--exact-bits", d.exact_bits),
        seed: d.seed,
    }
}

fn hex(b: &[u8]) -> String {
    b.iter().map(|x| format!("{x:02x}")).collect()
}

/// The peak resident memory of this process, in bytes.
#[cfg(unix)]
fn max_rss() -> Option<u64> {
    // struct rusage on 64-bit macOS and Linux: two timevals (16 bytes each), then longs.
    #[repr(C)]
    struct Rusage([i64; 18]);
    unsafe extern "C" {
        fn getrusage(who: i32, usage: *mut Rusage) -> i32;
    }
    let mut r = Rusage([0; 18]);
    // SAFETY: getrusage writes one struct rusage, 144 bytes on these targets, into r.
    if unsafe { getrusage(0, &mut r) } != 0 {
        return None;
    }
    let v = r.0[4].max(0) as u64;
    Some(if cfg!(target_os = "macos") { v } else { v * 1024 })
}

#[cfg(not(unix))]
fn max_rss() -> Option<u64> {
    None
}

/// `f`'s result and the milliseconds each of the checker's routes took in it
/// (`engine::miter::observe_routes`: the checker reads no clock; this does).
fn timed<T>(f: impl FnOnce() -> T) -> (T, Value) {
    use std::{cell::RefCell, rc::Rc};
    let marks: Rc<RefCell<Vec<(&'static str, Instant)>>> = Rc::new(RefCell::new(Vec::new()));
    let m = marks.clone();
    let t0 = Instant::now();
    miter::observe_routes(Some(Box::new(move |route| {
        // On standard error too, so a check stopped at the limit still says where it was.
        eprintln!("route {route} at {} ms", t0.elapsed().as_millis());
        m.borrow_mut().push((route, Instant::now()));
    })));
    let out = f();
    miter::observe_routes(None);
    let marks = marks.borrow();
    let mut t = serde_json::Map::new();
    for w in marks.windows(2) {
        if w[0].0 != "end" {
            t.insert(w[0].0.to_string(), json!((w[1].1 - w[0].1).as_millis() as u64));
        }
    }
    (out, Value::Object(t))
}

/// One check, in this process: the spec says what.
fn one(spec: &Value) -> Value {
    let b = miter::budget_of(spec).unwrap_or_else(|| fail("the spec names no budget"));
    let read = |k: &str| std::fs::read_to_string(spec[k].as_str().unwrap_or("")).unwrap_or_else(|e| fail(&format!("{}: {e}", spec[k])));
    let pick = spec["pick"].as_u64().unwrap_or(0);
    let how = spec["break"].as_str();
    let t0 = Instant::now();
    let mut out = json!({});
    match spec["kind"].as_str() {
        Some("theirs") => {
            let (src, mut comp) = (read("source"), read("compiled"));
            if let Some(h) = how {
                comp = break_text(&comp, h, pick);
            }
            let list = |k: &str| -> Vec<u32> { spec[k].as_array().into_iter().flatten().filter_map(|x| x.as_u64().map(|x| x as u32)).collect() };
            let (layout, fin) = (list("layout"), list("final_layout"));
            // Where each logical qubit ends, read from the measurements: the source's
            // own final permutation is in its measurements (a closing SWAP elided).
            let (v, route_ms) = timed(|| miter::texts_by_readout(&src, &comp, &layout, Some(&fin), Some(&[]), &b));
            out["route_ms"] = route_ms;
            out["final_layout_read"] = v["final_layout"].clone();
            out["final_layout_reported"] = json!(fin);
            out["verdict"] = brief(&v);
            // A compiler may drop diagonal gates just before the measurements: where one
            // unitary is not proved, the statement up to phases before them is decided too.
            // The same input already shows other outcome distributions: refuted there too.
            if let Some(m) = common::measured_by_the_same_input(&v) {
                out["measured"] = m;
            } else if v["verdict"] != "proved" {
                let (w, route_ms) = timed(|| miter::texts_by_readout(&src, &comp, &layout, Some(&fin), None, &b));
                out["measured"] = brief(&w);
                out["measured_route_ms"] = route_ms;
            }
        }
        Some("ours") => {
            let text = read("source");
            let program = wai_quantum::qasm::parse(&text).unwrap_or_else(|e| fail(&e));
            let c = match wai_quantum::quantum_toolchain::from_qasm(&text) {
                Ok(c) => c,
                Err(e) => return json!({ "skipped": format!("our device compiler takes dyadic angles only: {e}") }),
            };
            let coupling = coupling_of(spec);
            let basis = Basis::ALL.into_iter().find(|x| format!("{x:?}") == spec["basis"].as_str().unwrap_or("")).unwrap_or_else(|| fail("no such basis"));
            let tc = Instant::now();
            let mut p = match native::compile(&c, basis, &coupling) {
                Ok(p) => p,
                Err(e) => return json!({ "skipped": format!("our compile refuses it: {e}") }),
            };
            out["compile_ms"] = json!(tc.elapsed().as_millis() as u64);
            out["native_gates"] = json!(p.gates.len());
            out["swaps"] = json!(p.swaps);
            if let Some(h) = how {
                p = break_native(&p, h, pick);
            }
            let tv = Instant::now();
            let (v, route_ms) = timed(|| miter::native(&program, &p, &b));
            out["check_ms"] = json!(tv.elapsed().as_millis() as u64);
            out["route_ms"] = route_ms;
            out["verdict"] = brief(&v);
        }
        _ => fail("the spec's kind is theirs or ours"),
    }
    out["ms"] = json!(t0.elapsed().as_millis() as u64);
    out["max_rss"] = json!(max_rss());
    out
}

fn coupling_of(spec: &Value) -> Coupling {
    match spec["coupling"].as_str().unwrap_or("") {
        "grid" => Coupling::grid(6, 9),
        "all-to-all" => {
            let pairs: Vec<(u16, u16)> = (0..36u16).flat_map(|a| (a + 1..36).map(move |b| (a, b))).collect();
            Coupling::new(36, &pairs).expect("all pairs")
        }
        _ => match spec["heavy_hex"].as_str() {
            Some(path) => {
                let cfg: Value = serde_json::from_str(&std::fs::read_to_string(path).unwrap_or_else(|e| fail(&format!("{path}: {e}")))).unwrap_or_else(|e| fail(&format!("{path}: {e}")));
                let n = cfg["n_qubits"].as_u64().unwrap_or(156) as u16;
                let pairs: Vec<(u16, u16)> = cfg["coupling_map"].as_array().into_iter().flatten().filter_map(|p| Some((p[0].as_u64()? as u16, p[1].as_u64()? as u16))).collect();
                Coupling::new(n, &pairs).unwrap_or_else(|e| fail(&e.to_string()))
            }
            None => Coupling::grid(12, 13),
        },
    }
}

/// Runs one check in a child process, within `limit`.
fn child(spec: &Value, limit: Duration) -> Value {
    let dir = std::env::temp_dir().join(format!("equiv-bench-{}", std::process::id()));
    std::fs::create_dir_all(&dir).unwrap_or_else(|e| fail(&format!("{}: {e}", dir.display())));
    let path = dir.join("spec.json");
    std::fs::write(&path, spec.to_string()).unwrap_or_else(|e| fail(&format!("{}: {e}", path.display())));
    let out_path = dir.join("out.json");
    let _ = std::fs::remove_file(&out_path);
    let exe = std::env::current_exe().unwrap_or_else(|e| fail(&e.to_string()));
    let mut c = std::process::Command::new(exe)
        .arg("--one")
        .arg(&path)
        .stdout(std::fs::File::create(&out_path).unwrap_or_else(|e| fail(&e.to_string())))
        .spawn()
        .unwrap_or_else(|e| fail(&e.to_string()));
    let t0 = Instant::now();
    loop {
        match c.try_wait() {
            Ok(Some(status)) => {
                let text = std::fs::read_to_string(&out_path).unwrap_or_default();
                return serde_json::from_str(&text).unwrap_or_else(|_| json!({ "failed": format!("the check exited {status} with no answer") }));
            }
            Ok(None) if t0.elapsed() > limit => {
                let _ = c.kill();
                let _ = c.wait();
                return json!({ "killed": format!("past the host's limit of {} s", limit.as_secs()) });
            }
            Ok(None) => std::thread::sleep(Duration::from_millis(20)),
            Err(e) => return json!({ "failed": e.to_string() }),
        }
    }
}

fn signing_key() -> (SigningKey, &'static str) {
    match std::env::var("EQUIV_SIGNING_KEY") {
        Ok(path) => {
            let text = std::fs::read_to_string(&path).unwrap_or_else(|e| fail(&format!("the key {path}: {e}")));
            let text = text.trim();
            if text.len() != 64 {
                fail("the key file holds 64 hex digits");
            }
            let mut b = [0u8; 32];
            for (i, x) in b.iter_mut().enumerate() {
                *x = u8::from_str_radix(&text[2 * i..2 * i + 2], 16).unwrap_or_else(|_| fail("hex digits"));
            }
            (SigningKey::from_bytes(&b), "the operator's key")
        }
        Err(_) => (SigningKey::from_bytes(&[42u8; 32]), "a demonstration key every reader knows (set EQUIV_SIGNING_KEY to sign with your own)"),
    }
}

fn host() -> Value {
    let sysctl = |k: &str| std::process::Command::new("sysctl").args(["-n", k]).output().ok().map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string()).filter(|s| !s.is_empty());
    let proc_line = |file: &str, key: &str| -> Option<String> { std::fs::read_to_string(file).ok()?.lines().find(|l| l.starts_with(key)).and_then(|l| l.split_once(':')).map(|(_, v)| v.trim().to_string()) };
    let cpu = sysctl("machdep.cpu.brand_string").or_else(|| proc_line("/proc/cpuinfo", "model name"));
    let memory = sysctl("hw.memsize").and_then(|s| s.parse::<u64>().ok()).or_else(|| proc_line("/proc/meminfo", "MemTotal").and_then(|v| v.trim_end_matches(" kB").trim().parse::<u64>().ok()).map(|k| k * 1024));
    let load = sysctl("vm.loadavg").or_else(|| std::fs::read_to_string("/proc/loadavg").ok()).map(|s| s.trim().trim_matches(|c| c == '{' || c == '}').trim().to_string());
    json!({ "os": std::env::consts::OS, "arch": std::env::consts::ARCH, "cpu": cpu, "cores": std::thread::available_parallelism().map(|n| n.get()).ok(), "memory_bytes": memory, "load_at_start": load })
}

fn main() {
    let args: Vec<String> = std::env::args().skip(1).collect();
    if let Some(p) = opt(&args, "--one") {
        let spec: Value = serde_json::from_str(&std::fs::read_to_string(&p).unwrap_or_else(|e| fail(&format!("{p}: {e}")))).unwrap_or_else(|e| fail(&format!("{p}: {e}")));
        println!("{}", one(&spec));
        return;
    }
    let Some(dir) = args.first().filter(|a| !a.starts_with("--")).map(PathBuf::from) else { fail("give the benchmark directory") };
    let manifest: Value = serde_json::from_str(&std::fs::read_to_string(dir.join("manifest.json")).unwrap_or_else(|e| fail(&format!("manifest.json: {e}")))).unwrap_or_else(|e| fail(&e.to_string()));
    let b = budget(&args);
    let limit = Duration::from_secs(num(&args, "--limit-secs", 300));
    let max_n: u64 = num(&args, "--max-n", u64::MAX);
    let min_n: u64 = num(&args, "--min-n", 0);
    let only = opt(&args, "--only");
    let heavy_hex = opt(&args, "--heavy-hex").map(|p| std::fs::canonicalize(&p).unwrap_or_else(|e| fail(&format!("{p}: {e}"))).display().to_string());
    let path = |f: &str| dir.join(f).display().to_string();
    let budget_json = json!({ "budget": b.to_json() });
    let host_at_start = host();
    let mut rows = Vec::new();
    let mut inputs = Vec::new();
    let mut tally = std::collections::BTreeMap::<String, u64>::new();
    for row in manifest["rows"].as_array().into_iter().flatten() {
        let (name, n) = (row["benchmark"].as_str().unwrap_or(""), row["n"].as_u64().unwrap_or(0));
        if n > max_n || n < min_n || only.as_deref().is_some_and(|o| o != name) || row.get("indep").is_none() {
            continue;
        }
        let src = row["indep"]["file"].as_str().unwrap_or("");
        inputs.push(json!({ "file": src, "sha256": row["indep"]["sha256"] }));
        let mut checks = Vec::new();
        let mut run = |what: Value, mut spec: Value, checks: &mut Vec<Value>| {
            spec["budget"] = budget_json["budget"].clone();
            let r = child(&spec, limit);
            let verdict = r["verdict"]["verdict"].as_str().or(r["skipped"].as_str().map(|_| "skipped")).or(r["killed"].as_str().map(|_| "killed")).unwrap_or("failed").to_string();
            let broken = what["broken"].is_string();
            let key = format!("{}{}", if broken { "broken " } else { "equal " }, verdict);
            *tally.entry(key).or_default() += 1;
            println!("{name:<12} {n:>3} {:<24} {:<6} {verdict:<9} {:>9} ms {:>6} MB  {}", what["by"].as_str().unwrap_or(""), what["broken"].as_str().unwrap_or("-"), r["ms"], r["max_rss"].as_u64().map_or(0, |x| x >> 20), r["route_ms"]);
            let mut o = what;
            o["result"] = r;
            checks.push(o);
        };
        for (dev, m) in row["mapped"].as_object().into_iter().flatten() {
            let Some(file) = m["file"].as_str() else { continue };
            inputs.push(json!({ "file": file, "sha256": m["sha256"] }));
            for (k, brk) in [None, Some("drop"), Some("angle")].into_iter().enumerate() {
                let spec = json!({ "kind": "theirs", "source": path(src), "compiled": path(file), "layout": m["initial_layout"], "final_layout": m["final_layout"], "break": brk, "pick": 7919 * n + k as u64 });
                run(json!({ "by": format!("theirs:{dev}"), "gates": m["gates"], "broken": brk }), spec, &mut checks);
            }
        }
        if !flag(&args, "--no-ours") {
            for (fam, basis) in [("heavy-hex", "RzSxXCz"), ("grid", "PrxCz"), ("all-to-all", "GpiZz")] {
                for (k, brk) in [None, Some("drop"), Some("angle")].into_iter().enumerate() {
                    let spec = json!({ "kind": "ours", "source": path(src), "basis": basis, "coupling": fam, "heavy_hex": heavy_hex, "break": brk, "pick": 7919 * n + k as u64 });
                    run(json!({ "by": format!("ours:{fam}:{basis}"), "broken": brk }), spec, &mut checks);
                }
            }
        }
        rows.push(json!({ "benchmark": name, "n": n, "source_gates": row["indep"]["gates"], "checks": checks }));
    }
    let published: Vec<Value> = PUBLISHED.iter().map(|&(name, n, g, g2, zx, dd)| json!({ "benchmark": name, "n": n, "G": g, "G_compiled": g2, "rewriting_s": zx, "decision_diagram_s": dd })).collect();
    let request = json!({
        "what": "the equivalence checker on a published benchmark set: each circuit against its compilations by the suite's compiler (three devices) and by ours (three native gate sets), equal and broken one way each (a gate removed, a rotation turned by π/4)",
        "suite": { "cite": "Quetschlich, Burgholzer and Wille, Quantum 7, 1062 (2023), MIT licence", "generator": manifest["generator"], "compiler": manifest["compiler"], "inputs": inputs },
        "budget": b.to_json(),
        "limit_secs": limit.as_secs(),
        "heavy_hex": heavy_hex.as_deref().map(|p| Path::new(p).file_name().map(|f| f.to_string_lossy().to_string())),
    });
    let mut result = json!({
        "rows": rows,
        "tally": tally,
        "published": { "cite": "Peham, Burgholzer and Wille, IEEE JETCAS 12(3), 2022, Table III (compiled circuits, equivalent instances; their compiler then, another device)", "rows": published },
        "host": host_at_start,
        "times_are": "wall-clock milliseconds of each check in its own process on this host; max_rss its peak resident memory; the verdicts are the same bytes on every host",
    });
    let (key, which) = signing_key();
    result["signed_with"] = json!(which);
    let receipt = wai_quantum::engine::receipt::seal(&key, &request, &result);
    let doc = json!({ "request": request, "result": result, "receipt": receipt });
    let ok = wai_quantum::engine::receipt::verify(&doc["request"], &doc["result"], &doc["receipt"]);
    if let Some(out) = opt(&args, "--receipt") {
        std::fs::write(&out, serde_json::to_string_pretty(&doc).unwrap_or_default() + "\n").unwrap_or_else(|e| fail(&format!("{out}: {e}")));
        println!("receipt {out} (seal verifies: {ok}), key {}", hex(key.verifying_key().as_bytes()));
    }
    println!("{}", serde_json::to_string(&doc["result"]["tally"]).unwrap_or_default());
}