//! OTOC(2) at the edge of ergodicity on the dense state vector: the published
//! 18- to 31-qubit circuits' off-diagonal part `ΔC⁽⁴⁾`, against the values
//! published with them, with a signed receipt.
//!
//! ```text
//! cargo run --release --example otoc_sv --features quantum_sv,quantum_spd_receipt -- SET.txt [options]
//! ```
//!
//! `SET.txt` is a circuit set in the text form `quantum_otoc::EchoSet::parse`
//! reads (`wai/tools/quantum-otoc/` fetches the published circuits and converts
//! them). Options:
//!
//! - `--circuits 0-9` or `--circuits 0,4,7`: the published indices to run (all by default);
//! - `--instances N`: insertion instances for the diagonal part (20, as published);
//!   `--plain` runs the plain correlator only;
//! - `--seed S`: the insertion strings' seed (1);
//! - `--insert D`: the moment after which the strings are inserted. A set
//!   carries its script's `depth // 2 + shift` (9 for the 18-qubit set); the
//!   values published for the 18-qubit set are reproduced at moment 15
//!   (`tests/otoc_reference.rs`), and the receipt binds the moment used;
//! - `--width W`: the fusion width (4);
//! - `--receipt PATH`: where the signed receipt goes;
//! - the signing key is 32 bytes of hex in the file `OTOC_SIGNING_KEY` names, or
//!   a fixed demonstration key.
//!
//! Prints each circuit's plain value, the instances' mean and spread, the
//! estimate, the published value and the difference in units of the two
//! estimates' combined Monte Carlo error (the published one used 20 instances);
//! then the set's χ² per circuit, the work, the time and the energy. The error
//! comes from the circuit's own instances' spread, itself an estimate: with few
//! instances `z` is Student's t with one degree of freedom fewer than the
//! instances, not a normal variable (at three, `|z| > 3` one time in ten), and
//! the χ² band printed holds only when the instances are many
//! (`tests/otoc_reference.rs` reads the 31-qubit set so).

use std::time::Instant;
use wai_quantum::ed25519_dalek::SigningKey;
use wai_quantum::quantum_otoc::{EchoSet, Kind, Method, OtocJob, Quantity, published_for, seal_metered};

fn main() {
    let args: Vec<String> = std::env::args().skip(1).collect();
    let path = args.first().expect("usage: otoc_sv SET.txt [--circuits A-B] [--instances N | --plain] [--seed S] [--width W] [--receipt PATH]");
    let opt = |k: &str| args.iter().position(|a| a == k).map(|i| args[i + 1].clone());
    let mut set = EchoSet::parse(&std::fs::read_to_string(path).expect("the set reads")).expect("a valid set");
    if let Some(sel) = opt("--circuits") {
        set = set.subset(&indices(&sel));
    }
    if let Some(d) = opt("--insert") {
        // Another insertion moment than the set's (part of the patch's hash).
        set.insert = d.parse().expect("a moment");
    }
    if opt("--kind").as_deref() == Some("otoc1") {
        // Exploration: the OTOC(1) of the same blocks, which nothing published.
        set.kind = Kind::Otoc1;
        set.name += "-as-otoc1";
    }
    let width: u32 = opt("--width").map_or(4, |w| w.parse().expect("a width"));
    if let Some(mode) = opt("--explore") {
        explore(&set, &mode, opt("--instances").map_or(100, |n| n.parse().expect("a count")), width);
        return;
    }
    let quantity = if args.iter().any(|a| a == "--plain") || set.kind == Kind::Otoc1 {
        Quantity::Plain
    } else {
        Quantity::OffDiagonal { instances: opt("--instances").map_or(20, |n| n.parse().expect("a count")), seed: opt("--seed").map_or(1, |s| s.parse().expect("a seed")) }
    };
    let (key, id) = signer();
    let published = published_for(&set.name);
    println!("set {} · {} qubits · {} cycles · {} circuits · {quantity:?} · fusion width {width}", set.name, set.n, set.cycles, set.circuits.len());
    if let Some((_, what)) = published {
        println!("published: {what}");
    }
    let kind = set.kind;
    let job = OtocJob { set, quantity, method: Method::Sv { width } };
    let mut last = Instant::now();
    let mut z2 = Vec::new();
    let (receipt, json) = seal_metered(&key, &id, &job, |o| {
        let secs = last.elapsed().as_secs_f64();
        last = Instant::now();
        let ours = o.value();
        match published.and_then(|(p, _)| p.get(o.index as usize)) {
            Some(&p) if !o.instances.is_empty() => {
                // Each estimate's Monte Carlo error from this run's spread.
                let k = o.instances.len() as f64;
                let se = (o.sd() * o.sd() * (1.0 / 20.0 + 1.0 / k)).sqrt();
                let z = (ours - p) / se;
                z2.push(z * z);
                println!(
                    "{:4}  plain {:+.6}  mean {:+.6}  sd {:.6}  ΔC4 {:+.6}  published {:+.6}  diff {:+.6}  z {:+.2}  {secs:.1}s",
                    o.index, o.plain, o.mean(), o.sd(), ours, p, ours - p, z
                );
            }
            Some(&p) if kind == Kind::Otoc2 => {
                println!("{:4}  plain {:+.9}  published ΔC4 {:+.9} (the plain value less the diagonal part)  {secs:.1}s", o.index, ours, p)
            }
            Some(&p) => println!("{:4}  value {:+.9}  published {:+.9}  diff {:+.2e}  {secs:.1}s", o.index, ours, p, ours - p),
            None => println!("{:4}  value {:+.9}  {secs:.1}s", o.index, ours),
        }
    })
    .expect("the job runs");
    if !z2.is_empty() {
        let k = z2.len() as f64;
        println!("chi2 per circuit {:.3} over {} circuits (1 expected; 99% band for 50: 0.59-1.52)", z2.iter().sum::<f64>() / k, z2.len());
    }
    let work: u64 = receipt.outcomes.iter().map(|o| o.work).sum();
    println!("work {work} amplitude-gate updates · {:.1} s · {} µJ (this process's CPU energy)", receipt.elapsed_ms as f64 / 1000.0, receipt.joules_micro);
    if let Some(out) = opt("--receipt") {
        std::fs::write(&out, json + "\n").expect("the receipt writes");
        println!("receipt {out} (receipt hash {})", hex(&receipt.receipt_hash()));
    }
}

/// Other readings of the published estimator, for diagnosis (no receipt):
/// `same` puts one string in both echoes, `xyz` draws strings without the
/// identity.
fn explore(set: &EchoSet, mode: &str, instances: u32, width: u32) {
    use wai_quantum::quantum_otoc::{SvRunner, insertion_strings};
    let mut sv = SvRunner::new(width);
    let published = published_for(&set.name).map(|(p, _)| p);
    let mut z2 = Vec::new();
    for c in &set.circuits {
        let plain = sv.otoc(set, c, None).expect("runs");
        let mut vals = Vec::new();
        for k in 0..instances {
            let (mut a, mut b) = insertion_strings(set.n, 1, c.index, k);
            match mode {
                "same" => b = a.clone(),
                "xyz" => {
                    // Each identity redrawn from X, Y, Z (16 values mod 3: within 1/16 of uniform).
                    let (fa, fb) = insertion_strings(set.n, 2, c.index, k);
                    let (ga, gb) = insertion_strings(set.n, 3, c.index, k);
                    for (s, v) in a.iter_mut().zip(fa.iter().zip(&ga)).chain(b.iter_mut().zip(fb.iter().zip(&gb))) {
                        if *s == 0 {
                            *s = 1 + (4 * v.0 + v.1) % 3;
                        }
                    }
                }
                _ => panic!("unknown mode {mode}"),
            }
            vals.push(sv.otoc(set, c, Some((&a, &b))).expect("runs"));
        }
        let k = vals.len() as f64;
        let mean = vals.iter().sum::<f64>() / k;
        let sd = (vals.iter().map(|v| (v - mean) * (v - mean)).sum::<f64>() / (k - 1.0)).sqrt();
        let ours = plain - mean;
        let p = published.and_then(|p| p.get(c.index as usize)).copied().unwrap_or(f64::NAN);
        let z = (ours - p) / (sd * sd * (1.0 / 20.0 + 1.0 / k)).sqrt();
        z2.push(z * z);
        println!("{mode} {:4}  plain {plain:+.6}  mean {mean:+.6}  sd {sd:.6}  ΔC4 {ours:+.6}  published {p:+.6}  z {z:+.2}", c.index);
    }
    println!("{mode} chi2 per circuit {:.3}", z2.iter().sum::<f64>() / z2.len() as f64);
}

fn indices(s: &str) -> Vec<u32> {
    s.split(',')
        .flat_map(|part| match part.split_once('-') {
            Some((a, b)) => (a.parse::<u32>().expect("an index")..=b.parse::<u32>().expect("an index")).collect::<Vec<_>>(),
            None => vec![part.parse().expect("an index")],
        })
        .collect()
}

fn signer() -> (SigningKey, String) {
    match std::env::var("OTOC_SIGNING_KEY") {
        Ok(p) => {
            let h = std::fs::read_to_string(&p).expect("the key file reads");
            let h = h.trim();
            let b: Vec<u8> = (0..32).map(|i| u8::from_str_radix(&h[2 * i..2 * i + 2], 16).expect("hex")).collect();
            (SigningKey::from_bytes(&b.try_into().expect("32 bytes")), "wai-quantum otoc reproduction".into())
        }
        Err(_) => (SigningKey::from_bytes(&[0x0e; 32]), "wai-quantum otoc example (demonstration key)".into()),
    }
}

fn hex(b: &[u8]) -> String {
    b.iter().map(|x| format!("{x:02x}")).collect()
}