//! The dense SVD (`quantum_svd`) against the scalar one-sided Jacobi rotations
//! the tensor networks used before it: time, accuracy, and the bits of the
//! result at each thread count.
//!
//! `cargo run --release --example svd_bench --features quantum_engine,quantum_receipt,quantum_bptn -- [sizes=64,128,256,512,1024,2048] [threads=1,4,16] [jacobi_max=512] [shapes=square,tall,wide] [dims=4096x256,16384x512]`
//!
//! Each matrix is complex, its entries from a fixed integer sequence, its
//! columns scaled by `0.995^j` (formed by repeated multiplication, so the input
//! is the same bits everywhere) so its spectrum decays as a truncated bond's
//! does. For each size and shape: the new SVD at every thread count given (the
//! result hashed, bit for bit; every count must give the same hash), the
//! Jacobi rotations where the size allows, the reconstruction error
//! `‖A − U·S·V†‖/‖A‖`, the departures of `U†U` and `V†V` from the identity, and
//! the largest difference between the two methods' singular values. The hash
//! is the same on every machine: compare a run on arm64 with one on x86_64.
//!
//! Writes a signed receipt to `RECEIPTS_DIR` (default `receipts`).

#[path = "tracker/common.rs"]
mod common;
#[path = "svd/inputs.rs"]
mod inputs;

use std::time::Instant;

use serde_json::{Value, json};
use inputs::{hash, input};
use wai_quantum::quantum_svd::{C, Keep, Svd, jacobi, matmul, phases, svd};

fn arg<T: std::str::FromStr>(args: &[String], key: &str) -> Option<Vec<T>> {
    args.iter().find_map(|a| a.strip_prefix(&format!("{key}=")).map(|v| v.split(',').map(|x| x.parse().unwrap_or_else(|_| panic!("{key}={v}"))).collect()))
}

fn frob(a: &[C]) -> f64 {
    a.iter().map(|x| x.re * x.re + x.im * x.im).sum::<f64>().sqrt()
}

/// `‖A − U·S·V†‖/‖A‖` and the departures of `U†U`, `V†V` from `I` (max entry).
fn accuracy(a: &[C], m: usize, n: usize, f: &Svd, threads: usize) -> (f64, f64, f64) {
    let k = f.k;
    let us: Vec<C> = (0..m * k).map(|i| C { re: f.u[i].re * f.s[i % k], im: f.u[i].im * f.s[i % k] }).collect();
    let back = matmul(&us, &f.vh, m, k, n, threads);
    let diff: Vec<C> = back.iter().zip(a).map(|(x, y)| C { re: x.re - y.re, im: x.im - y.im }).collect();
    let uh: Vec<C> = (0..k * m).map(|t| f.u[(t % m) * k + t / m]).map(|z| C { re: z.re, im: -z.im }).collect();
    let uu = matmul(&uh, &f.u, k, m, k, threads);
    let vv_h: Vec<C> = (0..n * k).map(|t| f.vh[(t % k) * n + t / k]).map(|z| C { re: z.re, im: -z.im }).collect();
    let vv = matmul(&f.vh, &vv_h, k, n, k, threads);
    let off = |g: &[C]| (0..k * k).map(|t| (g[t].re - if t / k == t % k { 1.0 } else { 0.0 }).abs().max(g[t].im.abs())).fold(0.0, f64::max);
    (frob(&diff) / frob(a), off(&uu), off(&vv))
}

fn main() {
    let args: Vec<String> = std::env::args().skip(1).collect();
    let cores = std::thread::available_parallelism().map_or(1, |p| p.get());
    let sizes: Vec<usize> = arg(&args, "sizes").unwrap_or_else(|| vec![64, 128, 256, 512, 1024, 2048]);
    let threads: Vec<usize> = arg(&args, "threads").unwrap_or_else(|| vec![1, 4, cores]);
    let jacobi_max: usize = arg(&args, "jacobi_max").and_then(|v| v.first().copied()).unwrap_or(512);
    let shapes: Vec<String> = arg(&args, "shapes").unwrap_or_else(|| vec!["square".into(), "tall".into()]);
    // Explicit shapes, `dims=4096x256,16384x512`, run after the sized ones.
    let dims: Vec<(usize, usize)> = arg::<String>(&args, "dims")
        .unwrap_or_default()
        .iter()
        .map(|d| {
            let (m, n) = d.split_once('x').unwrap_or_else(|| panic!("dims: MxN, not {d}"));
            (m.parse().unwrap_or_else(|_| panic!("dims: {d}")), n.parse().unwrap_or_else(|_| panic!("dims: {d}")))
        })
        .collect();
    println!("host            {} {}, {cores} logical cores", std::env::consts::OS, std::env::consts::ARCH);
    println!("{:<7} {:>11} {:>8} {:>10} {:>10} {:>9} {:>9} {:>9} {:>9}  hash", "shape", "m x n", "threads", "new s", "jacobi s", "speedup", "resid", "U'U-I", "V'V-I");
    let (rows, secs, energy) = common::metered(|| {
        let mut rows: Vec<Value> = Vec::new();
        let mut cases: Vec<(String, usize, usize)> = Vec::new();
        for shape in &shapes {
            for &n0 in &sizes {
                cases.push(match shape.as_str() {
                    "square" => (shape.clone(), n0, n0),
                    "tall" => (shape.clone(), 4 * n0, n0),
                    "wide" => (shape.clone(), n0 / 4, n0),
                    other => panic!("square, tall or wide, not {other}"),
                });
            }
        }
        cases.extend(dims.iter().map(|&(m, n)| ("dims".to_string(), m, n)));
        for (shape, m, n) in cases {
            {
                let a = input(m, n, (m * 1_000_003 + n) as u64);
                let mut runs = Vec::new();
                let mut first: Option<(Svd, String)> = None;
                for &t in &threads {
                    let (before, r0) = phases();
                    let started = Instant::now();
                    let f = svd(&a, m, n, Keep::All, t);
                    let s = started.elapsed().as_secs_f64();
                    let (after, r1) = phases();
                    let split: Vec<Value> = before.iter().zip(&after).filter(|(x, y)| y.1 > x.1).map(|(x, y)| json!([x.0, y.1 - x.1])).collect();
                    let h = hash(&f);
                    runs.push(json!({"threads": t, "seconds": s, "hash": h, "converged": f.converged, "phases": split, "rotations": r1 - r0}));
                    if first.is_none() {
                        first = Some((f, h));
                    }
                }
                let (f, h0) = first.expect("a thread count");
                let same = runs.iter().all(|r| r["hash"] == h0.as_str());
                let (resid, ou, ov) = accuracy(&a, m, n, &f, cores);
                let (jac_secs, sdiff) = if m.max(n) <= jacobi_max {
                    let started = Instant::now();
                    let (_, sj, _) = jacobi(&a, m, n);
                    let s = started.elapsed().as_secs_f64();
                    let d = f.s.iter().zip(&sj).map(|(x, y)| (x - y).abs()).fold(0.0, f64::max) / f.s[0];
                    (Some(s), Some(d))
                } else {
                    (None, None)
                };
                let best = runs.iter().map(|r| r["seconds"].as_f64().unwrap()).fold(f64::INFINITY, f64::min);
                for r in &runs {
                    let t = r["threads"].as_u64().unwrap();
                    let s = r["seconds"].as_f64().unwrap();
                    let js = jac_secs.map_or("-".to_string(), |j| format!("{j:.3}"));
                    let sp = jac_secs.map_or("-".to_string(), |j| format!("{:.1}x", j / s));
                    println!(
                        "{shape:<7} {:>11} {t:>8} {s:>10.4} {js:>10} {sp:>9} {resid:>9.1e} {ou:>9.1e} {ov:>9.1e}  {}",
                        format!("{m}x{n}"),
                        &r["hash"].as_str().unwrap()[..16]
                    );
                    let ph: Vec<String> = r["phases"].as_array().unwrap().iter().map(|p| format!("{} {:.3}", p[0].as_str().unwrap(), p[1].as_f64().unwrap())).collect();
                    println!("                {}; {} rotations", ph.join(", "), r["rotations"]);
                }
                assert!(same, "{m}x{n}: the thread counts gave different bits");
                rows.push(json!({
                    "shape": shape, "m": m, "n": n, "runs": runs, "same_bits_at_every_thread_count": same,
                    "best_seconds": best, "jacobi_seconds": jac_secs, "jacobi_max_singular_value_difference": sdiff,
                    "residual": resid, "u_orthonormality": ou, "v_orthonormality": ov,
                }));
            }
        }
        rows
    });
    let request = json!({
        "benchmark": "dense complex SVD (quantum_svd::svd, Keep::All) against the scalar one-sided Jacobi rotations (quantum_svd::jacobi)",
        "input": "entries from a splitmix64 sequence, uniform in [-1/2, 1/2) for each part, column j scaled by 0.995^j (by repeated multiplication)",
        "sizes": sizes, "shapes": shapes, "dims": dims.iter().map(|(m, n)| format!("{m}x{n}")).collect::<Vec<_>>(), "threads": threads, "jacobi_max": jacobi_max,
        "host": { "os": std::env::consts::OS, "arch": std::env::consts::ARCH, "logical_cores": cores },
    });
    let result = json!({ "rows": rows, "wall_seconds": secs, "peak_rss_bytes": common::peak_rss_bytes(), "energy_figure": energy });
    let out = common::receipt_path(&format!("svd_bench_{}_{}.json", std::env::consts::OS, std::env::consts::ARCH));
    assert!(common::verifies(&common::seal(request, result, Some(&out))));
    println!("receipt         {}", out.display());
}