//! Union-find against exact matching on rotated surface-code memory
//! experiments: logical failures on the same shots, and time per shot.
//!
//! ```text
//! cargo run --release --example uf_bench --features quantum_uf -- P SHOTS D [D …] [--uf-only]
//! ```
use wai_quantum::quantum_frame::{error_model, parse, sample, surface_code_memory};
use wai_quantum::quantum_match::MatchingGraph;
use wai_quantum::quantum_uf::UnionFind;

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let p: f64 = a[1].parse().unwrap();
    let shots: usize = a[2].parse().unwrap();
    let uf_only = a.iter().any(|x| x == "--uf-only");
    for d in a[3..].iter().filter(|x| !x.starts_with("--")).map(|x| x.parse::<u32>().unwrap()) {
        let c = parse(&surface_code_memory(d, d, p)).unwrap();
        let g = MatchingGraph::from_model(&error_model(&c).unwrap()).unwrap();
        let det = sample(&c, shots, 7);
        let uf = UnionFind::new(&g);
        let mut ws = uf.workspace();
        let shots_fired: Vec<Vec<u32>> = (0..shots).map(|s| det.fired(s)).collect();
        let t0 = std::time::Instant::now();
        let mut uf_fail = 0;
        for (s, f) in shots_fired.iter().enumerate() {
            uf_fail += u64::from(uf.decode(f, &mut ws).unwrap() != det.flips(s));
        }
        let uf_us = t0.elapsed().as_secs_f64() * 1e6 / shots as f64;
        let mean_fired = shots_fired.iter().map(Vec::len).sum::<usize>() as f64 / shots as f64;
        if uf_only {
            println!("d={d:>2} p={p} shots={shots} fired/shot {mean_fired:.1}  uf {uf_fail} ({uf_us:.2} us/shot)");
            continue;
        }
        let t1 = std::time::Instant::now();
        let mut mw_fail = 0;
        for (s, f) in shots_fired.iter().enumerate() {
            mw_fail += u64::from(g.decode(f).unwrap() != det.flips(s));
        }
        let mw_us = t1.elapsed().as_secs_f64() * 1e6 / shots as f64;
        println!(
            "d={d:>2} p={p} shots={shots} fired/shot {mean_fired:.1}  uf {uf_fail} ({uf_us:.2} us/shot)  exact {mw_fail} ({mw_us:.1} us/shot)  ratio {:.2}  speedup {:.0}x",
            uf_fail as f64 / mw_fail.max(1) as f64,
            mw_us / uf_us
        );
    }
}