//! Patched cross-entropy benchmarking from a circuit, its patches and the
//! measured shots: each patch simulated exactly on the dense state vector.
//!
//! ```text
//! cargo run --release --example patch_xeb --features quantum_sv -- CIRCUIT.qasm SHOTS.u64 PATCHES
//! ```
//!
//! `CIRCUIT.qasm`: OpenQASM 3 with `rz`, `rx`, `ry`, `h`, `x`, `sx`, `cz`,
//! `cx` (barriers and measurements are skipped). Its two-qubit gates must stay
//! inside the patches: a patched circuit has its boundary couplers removed.
//! `SHOTS.u64`: one little-endian `u64` a shot, qubit q at bit q.
//! `PATCHES`: the patches' qubits, `0,1,2;3,4,5`.
//!
//! Prints, per patch: its qubits, the mean of `2^n·p(x)` over the shots, its
//! ideal value `2^n·Σp² − 1` (the XEB of the patch's own exact samples), and
//! the normalised fidelity `(2^n·mean p − 1) / (2^n·Σp² − 1)`; then the
//! product of the fidelities and its standard error (the patches' errors
//! added in quadrature, relative).

use wai_quantum::gates::Gate;
use wai_quantum::quantum_sv::StateVector;

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let src = std::fs::read_to_string(&args[1]).expect("the circuit reads");
    let raw = std::fs::read(&args[2]).expect("the shots read");
    let shots: Vec<u64> = raw.chunks_exact(8).map(|b| u64::from_le_bytes(b.try_into().unwrap())).collect();
    let patches: Vec<Vec<u32>> = args[3].split(';').map(|p| p.split(',').map(|q| q.trim().parse().unwrap()).collect()).collect();
    // The gates, with their qubits.
    let mut gates: Vec<(String, f64, Vec<u32>)> = Vec::new();
    for line in src.lines().map(str::trim) {
        let line = line.trim_end_matches(';');
        if line.is_empty() || ["OPENQASM", "include", "bit", "qubit", "barrier", "c["].iter().any(|p| line.starts_with(p)) || line.contains("measure") {
            continue;
        }
        let (head, args) = line.split_once(' ').expect("a gate and its qubits");
        let (name, theta) = match head.split_once('(') {
            Some((n, p)) => (n.to_string(), p.trim_end_matches(')').parse::<f64>().expect("one angle")),
            None => (head.to_string(), 0.0),
        };
        let qs: Vec<u32> = args.split(',').map(|a| a.trim().trim_start_matches("q[").trim_end_matches(']').parse().unwrap()).collect();
        gates.push((name, theta, qs));
    }
    let mut product = 1.0;
    let mut rel_var = 0.0;
    for patch in &patches {
        let n = patch.len() as u32;
        let local = |q: u32| patch.iter().position(|&p| p == q).map(|i| i as u32);
        let mut sv = StateVector::zero(n).unwrap();
        for (name, theta, qs) in &gates {
            let l: Vec<Option<u32>> = qs.iter().map(|&q| local(q)).collect();
            if l.iter().all(Option::is_none) {
                continue;
            }
            let l: Vec<u32> = l.into_iter().map(|x| x.expect("a gate crossing the patch boundary")).collect();
            let g = match (name.as_str(), l.as_slice()) {
                ("rz", [a]) => Gate::rz(*a, *theta),
                ("rx", [a]) => Gate::rx(*a, *theta),
                ("ry", [a]) => Gate::ry(*a, *theta),
                ("h", [a]) => Gate::h(*a),
                ("x", [a]) => Gate::x(*a),
                ("sx", [a]) => Gate::sqrt_x(*a),
                ("cz", [a, b]) => Gate::cz(*a, *b),
                ("cx", [a, b]) => Gate::cx(*a, *b),
                _ => panic!("gate {name} on {qs:?}"),
            };
            sv.apply(&g).unwrap();
        }
        let dim = (n as f64).exp2();
        let prob = |i: usize| sv.re[i] * sv.re[i] + sv.im[i] * sv.im[i];
        let ideal = dim * (0..sv.re.len()).map(|i| prob(i) * prob(i)).sum::<f64>() - 1.0;
        // Each shot's patch bits, as the state vector's index (local qubit 0 most significant).
        let vals: Vec<f64> = shots
            .iter()
            .map(|&x| {
                let idx = patch.iter().enumerate().fold(0usize, |i, (l, &q)| i | (((x >> q) & 1) as usize) << (n as usize - 1 - l));
                dim * prob(idx)
            })
            .collect();
        let mean = vals.iter().sum::<f64>() / vals.len() as f64;
        let var = vals.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / (vals.len() as f64 - 1.0);
        let f = (mean - 1.0) / ideal;
        let se = (var / vals.len() as f64).sqrt() / ideal;
        product *= f;
        rel_var += (se / f).powi(2);
        println!("patch {n} qubits: mean 2^n·p {mean:.6}, ideal xeb {ideal:.6}, fidelity {f:.6} ± {se:.6}");
    }
    println!("fidelity {product:.6e} ± {:.6e} ({} shots)", product.abs() * rel_var.sqrt(), shots.len());
}