//! Time the dense state vector: single gates by position and kind, then whole
//! circuits (brickwork, QFT, QAOA) at each fusion width.
//!
//! ```text
//! cargo run --release --example sv_bench --features quantum_sv -- N [gates|circuits|all]
//! ```
//!
//! A gate pass reads and writes every amplitude once, 32 bytes each, so the
//! `GB/s` column is that traffic over the time.
use std::time::Instant;
use wai_quantum::gates::Gate;
use wai_quantum::quantum_sv::{Schedule, StateVector};

fn rng(seed: u64) -> impl FnMut() -> f64 {
    let mut s = seed;
    move || {
        s = s.wrapping_add(0x9e37_79b9_7f4a_7c15);
        let mut z = s;
        z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
        (((z ^ (z >> 31)) >> 11) as f64 + 0.5) / 9_007_199_254_740_992.0
    }
}

fn brickwork(n: u32, depth: usize, seed: u64) -> Vec<Gate> {
    let mut u = rng(seed);
    let mut g = Vec::new();
    for layer in 0..depth {
        for q in 0..n {
            g.push(Gate::rx(q, u() * 6.0));
            g.push(Gate::rz(q, u() * 6.0));
        }
        let mut q = (layer % 2) as u32;
        while q + 1 < n {
            g.push(Gate::fsim(q, q + 1, u() * 3.0, u() * 3.0));
            q += 2;
        }
    }
    g
}

fn qft(n: u32) -> Vec<Gate> {
    let mut g = Vec::new();
    for i in 0..n {
        g.push(Gate::h(i));
        for j in i + 1..n {
            g.push(Gate::cphase(j, i, core::f64::consts::PI / f64::from(1u32 << (j - i).min(30))));
        }
    }
    g
}

fn qaoa(n: u32, p: usize, seed: u64) -> Vec<Gate> {
    let mut u = rng(seed);
    let mut g: Vec<Gate> = (0..n).map(Gate::h).collect();
    for _ in 0..p {
        let gamma = u();
        for q in 0..n {
            g.push(Gate::rzz(q, (q + 1) % n, gamma));
            g.push(Gate::rzz(q, (q + 3) % n, gamma));
        }
        let beta = u();
        for q in 0..n {
            g.push(Gate::rx(q, beta));
        }
    }
    g
}

fn time<F: FnMut()>(reps: usize, mut f: F) -> f64 {
    f();
    let t0 = Instant::now();
    for _ in 0..reps {
        f();
    }
    t0.elapsed().as_secs_f64() / reps as f64
}

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let n: u32 = a.get(1).map_or(24, |s| s.parse().unwrap());
    let what = a.get(2).map_or("all", |s| s.as_str());
    let bytes = 32.0 * (1u64 << n) as f64;
    let mut sv = StateVector::zero(n).unwrap();
    sv.run(&(0..n).map(Gate::h).collect::<Vec<_>>(), 0).unwrap();
    let threads = std::thread::available_parallelism().map_or(1, |t| t.get());
    println!("{n} qubits, {threads} threads, {:.2} GiB of state", 16.0 * (1u64 << n) as f64 / (1u64 << 30) as f64);
    if what == "gates" || what == "all" {
        let reps = if n >= 28 { 3 } else { 10 };
        let cases: Vec<(&str, Gate)> = vec![
            ("rx, top qubit", Gate::rx(0, 0.3)),
            ("rx, middle qubit", Gate::rx(n / 2, 0.3)),
            ("rx, bottom qubit", Gate::rx(n - 1, 0.3)),
            ("rz (diagonal), middle", Gate::rz(n / 2, 0.3)),
            ("fsim, top pair", Gate::fsim(0, 1, 0.4, 0.2)),
            ("fsim, middle pair", Gate::fsim(n / 2, n / 2 + 1, 0.4, 0.2)),
            ("fsim, bottom pair", Gate::fsim(n - 2, n - 1, 0.4, 0.2)),
            ("fsim, far apart", Gate::fsim(1, n - 2, 0.4, 0.2)),
            ("rzz (diagonal), middle", Gate::rzz(n / 2, n / 2 + 1, 0.3)),
        ];
        for (name, g) in cases {
            let t = time(reps, || sv.apply(&g).unwrap());
            println!("  {name:26} {:9.2} ms  {:7.1} GB/s", t * 1e3, bytes / t / 1e9);
        }
        for k in 3..=5u32 {
            // A dense k-qubit gate: a fused brickwork block on the bottom k qubits.
            let block = brickwork(k, 3, 7).into_iter().map(|mut g| {
                g.qubits.iter_mut().for_each(|q| *q += n - k);
                g
            });
            let fused = wai_quantum::quantum_sv::fuse(&block.collect::<Vec<_>>(), k as usize);
            if fused.len() == 1 {
                let g = &fused[0];
                let t = time(reps, || sv.apply(g).unwrap());
                println!("  dense {k}-qubit, bottom       {:9.2} ms  {:7.1} GB/s", t * 1e3, bytes / t / 1e9);
            }
        }
    }
    if what == "cache" {
        // Below 2^14 amplitudes every kernel runs on one thread: the in-cache
        // rate of one core, in amplitudes per second.
        let reps = (1usize << 26) >> n;
        let cases: Vec<(&str, Gate)> = vec![
            ("rx, top qubit", Gate::rx(0, 0.3)),
            ("rx, bottom qubit", Gate::rx(n - 1, 0.3)),
            ("rz (diagonal)", Gate::rz(n / 2, 0.3)),
            ("fsim, top pair", Gate::fsim(0, 1, 0.4, 0.2)),
            ("fsim, bottom pair", Gate::fsim(n - 2, n - 1, 0.4, 0.2)),
            ("fsim, far apart", Gate::fsim(1, n - 2, 0.4, 0.2)),
        ];
        for (name, g) in cases {
            let t = time(reps, || sv.apply(&g).unwrap());
            println!("  {name:26} {:9.1} M amplitudes/s per core", (1u64 << n) as f64 / t / 1e6);
        }
        for k in 3..=5u32 {
            let block = brickwork(k, 3, 7).into_iter().map(|mut g| {
                g.qubits.iter_mut().for_each(|q| *q += n - k);
                g
            });
            let fused = wai_quantum::quantum_sv::fuse(&block.collect::<Vec<_>>(), k as usize);
            if fused.len() == 1 {
                let g = &fused[0];
                let t = time(reps, || sv.apply(g).unwrap());
                println!("  dense {k}-qubit                {:9.1} M amplitudes/s per core", (1u64 << n) as f64 / t / 1e6);
            }
        }
    }
    if what == "circuits" || what == "all" {
        let circuits: Vec<(&str, Vec<Gate>)> = vec![("brickwork depth 20", brickwork(n, 20, 1)), ("qft", qft(n)), ("qaoa p=2", qaoa(n, 2, 3))];
        let tile: u32 = std::env::var("TILE").ok().map_or(15, |s| s.parse().unwrap());
        let low: u32 = std::env::var("RUN").ok().map_or(5, |s| s.parse().unwrap());
        let tiled = Schedule { tile_bits: tile, run_bits: low };
        for (name, gates) in &circuits {
            for width in [2usize, 3, 4, 5] {
                let passes = wai_quantum::quantum_sv::fuse(gates, width).len();
                let mut times = Vec::new();
                let mut states = Vec::new();
                for schedule in [Schedule::UNTILED, tiled] {
                    let mut s = StateVector::zero(n).unwrap();
                    let t0 = Instant::now();
                    s.run_scheduled(gates, width, schedule).unwrap();
                    times.push(t0.elapsed().as_secs_f64());
                    states.push(s);
                }
                let same = states[0].re.iter().chain(&states[0].im).zip(states[1].re.iter().chain(&states[1].im)).all(|(a, b)| a.to_bits() == b.to_bits());
                println!(
                    "  {name:20} {} gates, width {width}: {passes:5} gates fused; untiled {:8.3} s, tiled {:8.3} s ({:4.1}x), same bits: {same}",
                    gates.len(),
                    times[0],
                    times[1],
                    times[0] / times[1]
                );
            }
        }
    }
}