//! A Clifford circuit with one closing layer of single-qubit gates, read
//! from OpenQASM 2, summed exactly by `quantum_cutrank`.
//!
//! ```text
//! cargo run --release --example cutrank_qasm --features quantum_cutrank -- FILE cuts
//! cargo run --release --example cutrank_qasm --features quantum_cutrank -- FILE amplitudes N [SEED]
//! cargo run --release --example cutrank_qasm --features quantum_cutrank -- FILE sample N [SEED]
//! cargo run --release --example cutrank_qasm --features quantum_cutrank -- FILE probs STRINGS [qiskit|q0]
//! ```
//!
//! `probs` reads one bit string per line from STRINGS (as a device's samples)
//! and prints each one's `2^n · p`, then their mean and the linear cross-entropy
//! `2^n · mean(p) − 1` with its standard error. `qiskit` (the default) reads a
//! string with qubit n−1 first, as Qiskit writes it; `q0` with qubit 0 first.
//!
//! The bulk takes `h s sdg sx x y z cx cz swap`; a qubit's closing layer
//! starts at its first non-Clifford gate (`t tdg rz rx ry p u u3`) and may hold
//! any single-qubit gates. `2^n · p` is printed per bit string: for exact
//! samples of a state whose probabilities follow Porter–Thomas its mean is
//! about 2 (a linear cross-entropy of about 1).

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

fn mul(a: &Unitary, b: &Unitary) -> Unitary {
    let m = |x: C, y: C| complex(x.re * y.re - x.im * y.im, x.re * y.im + x.im * y.re);
    let add = |x: C, y: C| complex(x.re + y.re, x.im + y.im);
    [add(m(a[0], b[0]), m(a[1], b[2])), add(m(a[0], b[1]), m(a[1], b[3])), add(m(a[2], b[0]), m(a[3], b[2])), add(m(a[2], b[1]), m(a[3], b[3]))]
}

fn one_qubit(name: &str, p: &[f64]) -> Option<Unitary> {
    let c = complex;
    let r = std::f64::consts::FRAC_1_SQRT_2;
    let e = |x: f64| c(x.cos(), x.sin());
    let (o, one) = (c(0.0, 0.0), c(1.0, 0.0));
    Some(match name {
        "h" => [c(r, 0.0), c(r, 0.0), c(r, 0.0), c(-r, 0.0)],
        "x" => [o, one, one, o],
        "y" => [o, c(0.0, -1.0), c(0.0, 1.0), o],
        "z" => [one, o, o, c(-1.0, 0.0)],
        "s" => [one, o, o, c(0.0, 1.0)],
        "sdg" => [one, o, o, c(0.0, -1.0)],
        "t" => [one, o, o, e(std::f64::consts::FRAC_PI_4)],
        "tdg" => [one, o, o, e(-std::f64::consts::FRAC_PI_4)],
        "sx" => [c(0.5, 0.5), c(0.5, -0.5), c(0.5, -0.5), c(0.5, 0.5)],
        "rz" | "p" | "u1" => {
            let t = p[0];
            if name == "rz" { [e(-t / 2.0), o, o, e(t / 2.0)] } else { [one, o, o, e(t)] }
        }
        "rx" => {
            let (cs, sn) = ((p[0] / 2.0).cos(), (p[0] / 2.0).sin());
            [c(cs, 0.0), c(0.0, -sn), c(0.0, -sn), c(cs, 0.0)]
        }
        "ry" => {
            let (cs, sn) = ((p[0] / 2.0).cos(), (p[0] / 2.0).sin());
            [c(cs, 0.0), c(-sn, 0.0), c(sn, 0.0), c(cs, 0.0)]
        }
        _ => return None,
    })
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let (path, mode) = (&args[1], args[2].as_str());
    let count: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(4);
    let seed: u64 = args.get(4).and_then(|s| s.parse().ok()).unwrap_or(1);
    let src = std::fs::read_to_string(path).expect("the file reads");
    let mut n = 0u32;
    let mut bulk = Vec::new();
    let mut last: Vec<Unitary> = Vec::new();
    let mut closing: Vec<bool> = Vec::new();
    let id = [complex(1.0, 0.0), complex(0.0, 0.0), complex(0.0, 0.0), complex(1.0, 0.0)];
    for line in src.lines().map(str::trim).filter(|l| !l.is_empty()) {
        let line = line.trim_end_matches(';');
        if let Some(rest) = line.strip_prefix("qreg q[") {
            n = rest.trim_end_matches(']').parse().unwrap();
            last = vec![id; n as usize];
            closing = vec![false; n as usize];
            continue;
        }
        if line.starts_with("OPENQASM") || line.starts_with("include") || line.starts_with("creg") || line.starts_with("barrier") || line.starts_with("measure") {
            continue;
        }
        let (head, args) = line.split_once(' ').expect("a gate and its qubits");
        let (name, params): (&str, Vec<f64>) = match head.split_once('(') {
            Some((n, p)) => (n, p.trim_end_matches(')').split(',').map(|x| x.trim().parse().expect("a numeric angle")).collect()),
            None => (head, Vec::new()),
        };
        let qs: Vec<u32> = args.split(',').map(|a| a.trim().trim_start_matches("q[").trim_end_matches(']').parse().unwrap()).collect();
        let a = qs[0];
        let non_clifford = matches!(name, "t" | "tdg" | "rz" | "rx" | "ry" | "p" | "u1");
        if qs.len() == 1 && (closing[a as usize] || non_clifford) {
            closing[a as usize] = true;
            let u = one_qubit(name, &params).unwrap_or_else(|| panic!("gate {name}"));
            last[a as usize] = mul(&u, &last[a as usize]);
            continue;
        }
        assert!(qs.iter().all(|&q| !closing[q as usize]), "a gate on {qs:?} after its closing layer began: {line}");
        bulk.push(match (name, qs.as_slice()) {
            ("h", _) => Clifford::H(a),
            ("s", _) => Clifford::S(a),
            ("sdg", _) => Clifford::Sdg(a),
            ("sx", _) => Clifford::SX(a),
            ("x", _) => Clifford::X(a),
            ("y", _) => Clifford::Y(a),
            ("z", _) => Clifford::Z(a),
            ("cx", [_, b]) => Clifford::CX(a, *b),
            ("cz", [_, b]) => Clifford::CZ(a, *b),
            ("swap", [_, b]) => Clifford::Swap(a, *b),
            _ => panic!("not a bulk gate: {line}"),
        });
    }
    let t0 = Instant::now();
    let mut pm = ProductMeasured::new(&Tableau::run(n, &bulk).unwrap(), &last).unwrap();
    // CUTRANK_ORDER=FILE: an order (as `order` prints it) to sum in.
    if let Ok(path) = std::env::var("CUTRANK_ORDER") {
        let text = std::fs::read_to_string(&path).expect("the order file reads");
        let line = text.lines().find(|l| l.starts_with("order: ")).expect("an `order: [...]` line");
        let order: Vec<usize> = line["order: ".len()..].trim_matches(|c| c == '[' || c == ']').split(',').map(|x| x.trim().parse().unwrap()).collect();
        pm = pm.with_order(&order).expect("a permutation of the qubits");
    }
    println!("{n} qubits, {} bulk gates; graph form and cuts in {:?}", bulk.len(), t0.elapsed());
    println!("cut-ranks: {:?}", pm.cut_ranks());
    println!("max cut {} bits: a table of {} entries ({} GiB at f64)", pm.max_cut(), 1u64 << pm.max_cut(), (16u64 << pm.max_cut()) >> 30);
    let limit = std::env::var("CUTRANK_MAX").ok().and_then(|s| s.parse().ok()).unwrap_or(32);
    let precision = if std::env::var("CUTRANK_PRECISION").as_deref() == Ok("single") { Precision::Single } else { Precision::Double };
    println!("tables at {precision:?}: the largest {} GiB", pm.table_bytes(precision) >> 30);
    for (kind, passes, entries) in pm.passes() {
        println!("  {kind:8} {passes:3} passes, {:.2}e9 entries", entries as f64 / 1e9);
    }
    let scale = (n as f64).exp2();
    // `order N`: anneal the order for the least work, Σ 2^cut.
    let cost = |c: &[usize]| c.iter().map(|&k| (k as f64).exp2()).sum::<f64>();
    match mode {
        "cuts" => {}
        "order" => {
            let t0 = Instant::now();
            let best_order = pm.anneal_order(count, seed);
            let best = cost(pm.clone().with_order(&best_order).unwrap().cut_ranks());
            let b = pm.clone().with_order(&best_order).unwrap();
            println!("annealed {count} moves in {:?}: work Σ2^cut {best:.3e} (natural {:.3e}), max cut {}", t0.elapsed(), cost(pm.cut_ranks()), b.max_cut());
            println!("cuts: {:?}", b.cut_ranks());
            println!("order: {best_order:?}");
        }
        "amplitudes" => {
            let mut s = seed;
            for i in 0..count {
                let x: Vec<bool> = (0..n).map(|_| {
                    s = s.wrapping_add(0x9e37_79b9_7f4a_7c15);
                    let mut z = s;
                    z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
                    (z ^ (z >> 27)) & 1 == 1
                }).collect();
                let t = Instant::now();
                let a = pm.amplitude_in(&x, limit, precision).unwrap();
                println!("amplitude {i}: {a:?}, 2^n·p = {:.6}, in {:?}", scale * (a.re * a.re + a.im * a.im), t.elapsed());
            }
        }
        "sample" => {
            let t = Instant::now();
            let samples = pm.sample_in(count, seed, limit, precision).unwrap();
            let el = t.elapsed();
            let mean = samples.iter().map(|(_, p)| scale * p).sum::<f64>() / samples.len() as f64;
            // CUTRANK_OUT=FILE: every sample, `bits 2^n·p` a line (qubit 0 first).
            if let Ok(out) = std::env::var("CUTRANK_OUT") {
                let text: String = samples.iter().map(|(x, p)| format!("{} {:.9}\n", x.iter().map(|&b| if b { '1' } else { '0' }).collect::<String>(), scale * p)).collect();
                std::fs::write(&out, text).expect("the samples write");
            }
            for (x, p) in samples.iter().take(8) {
                let bits: String = x.iter().map(|&b| if b { '1' } else { '0' }).collect();
                println!("{bits} 2^n·p = {:.6}", scale * p);
            }
            println!("{count} samples in {el:?} ({:?} each); mean 2^n·p = {mean:.4}, linear XEB = {:.4}", el / count as u32, mean - 1.0);
        }
        // `probs FILE [qiskit|q0]`: the exact 2^n·p of each bit string in FILE (one
        // a line; `qiskit`, the default: the first character is qubit n−1, as
        // Qiskit prints counts; `q0`: the first is qubit 0), and their linear
        // XEB, 2^n·mean(p) − 1, with its standard error. `count` is unused.
        "probs" => {
            let file = args.get(3).expect("a file of bit strings");
            let qiskit = args.get(4).map(|s| s.as_str()) != Some("q0");
            let text = std::fs::read_to_string(file).expect("the bit strings read");
            let t = Instant::now();
            let mut v = Vec::new();
            for (i, line) in text.lines().map(str::trim).filter(|l| !l.is_empty()).enumerate() {
                assert_eq!(line.len(), n as usize, "line {i}: {} characters, not {n}", line.len());
                let c: Vec<bool> = line.bytes().map(|b| b == b'1').collect();
                let x: Vec<bool> = (0..n as usize).map(|q| if qiskit { c[n as usize - 1 - q] } else { c[q] }).collect();
                let a = pm.amplitude_in(&x, limit, precision).unwrap();
                let w = scale * (a.re * a.re + a.im * a.im);
                println!("{line} {w:.9}");
                v.push(w);
            }
            let m = v.len() as f64;
            let mean = v.iter().sum::<f64>() / m;
            let var = v.iter().map(|w| (w - mean) * (w - mean)).sum::<f64>() / (m - 1.0).max(1.0);
            println!("{} strings ({}) in {:?}: mean 2^n·p = {mean:.4}, linear XEB = {:.4} ± {:.4}", v.len(), if qiskit { "qubit n-1 first" } else { "qubit 0 first" }, t.elapsed(), mean - 1.0, (var / m).sqrt());
        }
        other => panic!("mode {other}"),
    }
}