//! Time noisy shots two ways: independent trajectories, and branching
//! trajectories (shots that share a history share a state), at three error
//! rates.
//!
//! ```text
//! cargo run --release --example noisy_bench --features quantum_noisy -- N SHOTS DEPTH
//! ```
use std::time::Instant;
use wai_quantum::gates::Gate;
use wai_quantum::quantum_noisy::{branching_trajectories, trajectories, NoiseModel};

fn brickwork(n: u32, depth: usize) -> Vec<Gate> {
    let mut g = Vec::new();
    let mut a = 0.1f64;
    for layer in 0..depth {
        for q in 0..n {
            a += 0.7311;
            g.push(if q % 2 == 0 { Gate::rx(q, a) } else { Gate::ry(q, a) });
        }
        let mut q = (layer % 2) as u32;
        while q + 1 < n {
            g.push(Gate::cz(q, q + 1));
            q += 2;
        }
    }
    g
}

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let n: u32 = a.get(1).map_or(14, |s| s.parse().unwrap());
    let shots: u64 = a.get(2).map_or(1000, |s| s.parse().unwrap());
    let depth: usize = a.get(3).map_or(10, |s| s.parse().unwrap());
    let threads = std::thread::available_parallelism().map_or(1, |t| t.get());
    let g = brickwork(n, depth);
    let obs = vec![vec![(1.0, vec![(0, 'Z'), (1, 'Z')])]];
    println!("{n} qubits, {} gates, {shots} shots, {threads} threads", g.len());
    for (name, noise) in [
        ("rare: 1e-4 / 1e-3", NoiseModel { depol1: 1e-4, depol2: 1e-3, readout: 0.01, ..NoiseModel::default() }),
        ("moderate: 1e-3 / 5e-3", NoiseModel { depol1: 1e-3, depol2: 5e-3, readout: 0.01, ..NoiseModel::default() }),
        ("heavy: 5e-3 / 2e-2, damping 1e-3", NoiseModel { depol1: 5e-3, depol2: 2e-2, damping: 1e-3, readout: 0.01, ..NoiseModel::default() }),
    ] {
        let t0 = Instant::now();
        let b = branching_trajectories(&g, n, &noise, &obs, shots, 1, 16, threads).unwrap();
        let tb = t0.elapsed().as_secs_f64();
        let t0 = Instant::now();
        let t = trajectories(&g, n, &noise, &obs, shots as usize, 1, threads).unwrap();
        let tt = t0.elapsed().as_secs_f64();
        println!(
            "  {name:34} independent {tt:7.2} s; branching {tb:7.2} s ({:5.1}x): {} histories, {} gate applications of {}; <ZZ> {:.4} ± {:.4} vs {:.4} ± {:.4}",
            tt / tb,
            b.histories,
            b.gate_applications,
            shots * g.len() as u64,
            b.expectations[0].mean,
            b.expectations[0].stderr,
            t.expectations[0].mean,
            t.expectations[0].stderr
        );
    }
}