//! Gaussian boson sampling numbers, as JSON, for checking against any other
//! implementation.
//!
//! ```text
//! cargo run --release --example gbs_check --features quantum_gbs -- state MODES SEED ETA R1,R2,…
//! cargo run --release --example gbs_check --features quantum_gbs -- hafnian N SEED
//! ```
use wai_quantum::quantum_gbs::{hafnian, random_unitary, GaussianState, Mat, C};

fn c(z: C) -> String {
    format!("[{:e},{:e}]", z.re, z.im)
}

fn main() {
    let a: Vec<String> = std::env::args().collect();
    match a.get(1).map(String::as_str) {
        Some("state") => {
            let n: usize = a[2].parse().unwrap();
            let seed: u64 = a[3].parse().unwrap();
            let eta: f64 = a[4].parse().unwrap();
            let r: Vec<f64> = a[5].split(',').map(|x| x.parse().unwrap()).collect();
            let u = random_unitary(n, seed);
            let st = GaussianState::squeezed(&r, &u, eta).unwrap();
            let cov: Vec<String> = st.cov.a.iter().map(|z| c(*z)).collect();
            let clicks: Vec<String> = (0..1u32 << n)
                .map(|m| {
                    let p: Vec<bool> = (0..n).map(|i| m >> i & 1 == 1).collect();
                    format!("{:e}", st.click_probability(&p).unwrap())
                })
                .collect();
            // Photon patterns with up to 4 photons in total.
            let mut photons = Vec::new();
            let mut counts = vec![0usize; n];
            fn rec(i: usize, left: usize, counts: &mut Vec<usize>, st: &GaussianState, out: &mut Vec<String>) {
                if i == counts.len() {
                    out.push(format!("[{:?},{:e}]", counts, st.photon_probability(counts).unwrap()));
                    return;
                }
                for k in 0..=left {
                    counts[i] = k;
                    rec(i + 1, left - k, counts, st, out);
                }
                counts[i] = 0;
            }
            rec(0, 4, &mut counts, &st, &mut photons);
            println!("{{\"n\":{n},\"cov\":[{}],\"clicks\":[{}],\"photons\":[{}]}}", cov.join(","), clicks.join(","), photons.join(","));
        }
        Some("hafnian") => {
            let n: usize = a[2].parse().unwrap();
            let seed: u64 = a[3].parse().unwrap();
            let u = random_unitary(n, seed);
            let mut m = Mat::zeros(n);
            for i in 0..n {
                for j in 0..n {
                    let (x, y) = (u.at(i, j), u.at(j, i));
                    m.a[i * n + j] = C { re: x.re + y.re, im: x.im + y.im };
                }
            }
            let t0 = std::time::Instant::now();
            let h = hafnian(&m);
            let dt = t0.elapsed().as_secs_f64();
            let entries: Vec<String> = m.a.iter().map(|z| c(*z)).collect();
            println!("{{\"n\":{n},\"m\":[{}],\"haf\":{},\"seconds\":{dt}}}", entries.join(","), c(h));
        }
        _ => eprintln!("usage: gbs_check state MODES SEED ETA R1,R2,… | hafnian N SEED"),
    }
}