//! One state's amplitudes by the cut-rank table on the host, in double and in
//! single precision, and on the platform's GPU, timed and compared.
//!
//! ```text
//! cargo run --release --example cutrank_gpu --features quantum_cutrank_metal -- QUBITS SEED [AMPLITUDES] [SHOTS]
//! ```
//!
//! The state is a random graph state on `QUBITS` qubits (a Hadamard on each,
//! then a CZ on each pair with probability ½), whose cut-rank is about half
//! the qubits on each side of the middle for any order, then a closing layer
//! of random single-qubit gates, summed along an order annealed for less work. Prints the widest cut, the seconds for each
//! amplitude on each, the largest difference of the GPU's amplitudes from the
//! double-precision ones (relative to each amplitude), and, with SHOTS, the
//! seconds for that many samples on each.

use std::time::Instant;
use wai_quantum::gates::complex;
use wai_quantum::quantum_cutrank::{Precision, ProductMeasured, Unitary};
use wai_quantum::quantum_cutrank_metal::GpuCutRank;
use wai_quantum::quantum_tableau::{Clifford, Tableau};

fn next(s: &mut u64) -> u64 {
    *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)
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args[1] == "passes" {
        // cutrank_gpu passes LOG_LEN REPS: each kernel's pass alone.
        let l: u32 = args[2].parse().expect("LOG_LEN");
        let reps: usize = args.get(3).map_or(10, |a| a.parse().expect("REPS"));
        let mut gpu = GpuCutRank::new(l as usize).expect("a GPU table");
        for (name, t) in gpu.time_passes(l, reps).expect("the passes run") {
            println!("{name:14} {:8.3} ms  {:6.2} G entries/s", t * 1e3, (1u64 << l) as f64 / t / 1e9);
        }
        return;
    }
    let n: usize = args[1].parse().expect("QUBITS");
    let seed: u64 = args[2].parse().expect("SEED");
    let amps: usize = args.get(3).map_or(3, |a| a.parse().expect("AMPLITUDES"));
    let shots: usize = args.get(4).map_or(0, |a| a.parse().expect("SHOTS"));
    let mut s = seed;
    let mut ops: Vec<Clifford> = (0..n as u32).map(Clifford::H).collect();
    for a in 0..n as u32 {
        for b in a + 1..n as u32 {
            if next(&mut s) & 1 == 1 {
                ops.push(Clifford::CZ(a, b));
            }
        }
    }
    let last: Vec<Unitary> = (0..n)
        .map(|_| {
            let mut ang = || (next(&mut s) >> 11) as f64 / (1u64 << 53) as f64 * std::f64::consts::TAU;
            let (t, p, l) = (ang(), ang(), ang());
            let (c, sn) = ((t / 2.0).cos(), (t / 2.0).sin());
            let e = |x: f64| complex(x.cos(), x.sin());
            let sc = |z: wai_quantum::gates::C, k: f64| complex(z.re * k, z.im * k);
            [complex(c, 0.0), sc(e(l), -sn), sc(e(p), sn), sc(e(p + l), c)]
        })
        .collect();
    let pm = ProductMeasured::new(&Tableau::run(n as u32, &ops).unwrap(), &last).unwrap();
    let order = pm.anneal_order(20_000, seed);
    let pm = pm.with_order(&order).unwrap();
    let cut = pm.max_cut();
    println!("{n} qubits, widest cut {cut} bits, work 2^{:.1}", pm.work().log2());
    let mut gpu = GpuCutRank::new(cut).expect("a GPU table");
    println!("device: {}", gpu.device());
    let mut worst: f64 = 0.0;
    let (mut td, mut ts, mut tg) = (0.0, 0.0, 0.0);
    for _ in 0..amps {
        let x: Vec<bool> = (0..n).map(|_| next(&mut s) & 1 == 1).collect();
        let t0 = Instant::now();
        let a = pm.amplitude(&x, 34).unwrap();
        td += t0.elapsed().as_secs_f64();
        let t0 = Instant::now();
        let _ = pm.amplitude_in(&x, 34, Precision::Single).unwrap();
        ts += t0.elapsed().as_secs_f64();
        let t0 = Instant::now();
        let b = gpu.amplitude(&pm, &x).unwrap();
        tg += t0.elapsed().as_secs_f64();
        worst = worst.max((a.re - b.re).hypot(a.im - b.im) / a.re.hypot(a.im).max(1e-300));
    }
    let k = amps.max(1) as f64;
    println!("seconds an amplitude: host double {:.3}, host single {:.3}, GPU {:.3}", td / k, ts / k, tg / k);
    println!("largest relative difference of the GPU's amplitudes: {worst:.2e}");
    if shots > 0 {
        let t0 = Instant::now();
        let _ = pm.sample_in(shots, seed, 34, Precision::Single).unwrap();
        let th = t0.elapsed().as_secs_f64();
        let t0 = Instant::now();
        let _ = gpu.sample(&pm, shots, seed).unwrap();
        let tg = t0.elapsed().as_secs_f64();
        let mut batched = GpuCutRank::with_batch(cut, 16).expect("a batch of GPU tables");
        let t0 = Instant::now();
        let _ = batched.sample(&pm, shots, seed).unwrap();
        let tb = t0.elapsed().as_secs_f64();
        println!("seconds for {shots} samples: host single {th:.3}, GPU {tg:.3}, GPU 16 at a time {tb:.3}");
    }
}