//! Out-of-time-order correlators of the published circuits by tensor networks,
//! as a measured exploration of the 65- and 95-qubit circuits against the
//! processor's values, and held to the exact state vector where that exists.
//! Two kinds, each sealed in a receipt per run:
//!
//! - `--chi X`: a state on the circuits' coupling graph, evolved by simple
//!   update with bonds of at most X, read through belief-propagation
//!   environments;
//! - `--bond X`: a matrix-product state along the qubits' index order, bonds of
//!   at most X, gates between non-neighbours applied through swaps.
//!
//! ```text
//! cargo run --release --example otoc_tn --features quantum_sv,quantum_spd_receipt,quantum_bptn,quantum_mpsim -- SET.txt --chi X|--bond X [options]
//! ```
//!
//! Other options as `otoc_spd` takes them. Each circuit prints its value, the
//! weight its truncations discarded (a diagnostic, not a bound), the widest
//! bond, the published value and the difference.

use std::time::Instant;
use wai_quantum::ed25519_dalek::SigningKey;
use wai_quantum::quantum_otoc::{EchoSet, 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_tn SET.txt --chi X|--bond X [--circuits A-B] [--instances N] [--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 opt("--kind").as_deref() == Some("otoc1") {
        // Exploration: the OTOC(1) of the same blocks, which nothing published.
        set.kind = wai_quantum::quantum_otoc::Kind::Otoc1;
        set.name += "-as-otoc1";
    }
    let method = match (opt("--chi"), opt("--bond")) {
        (Some(x), None) => Method::Bptn { chi: x.parse().expect("a bond") },
        (None, Some(x)) => Method::Mps { max_bond: x.parse().expect("a bond") },
        _ => panic!("give one of --chi X and --bond X"),
    };
    let quantity = match opt("--instances") {
        Some(n) => Quantity::OffDiagonal { instances: n.parse().expect("a count"), seed: opt("--seed").map_or(1, |s| s.parse().expect("a seed")) },
        None => Quantity::Plain,
    };
    let (key, id) = signer();
    let published = published_for(&set.name);
    println!("set {} · {} qubits · {} cycles · {} circuits · {quantity:?} · {method:?}", set.name, set.n, set.cycles, set.circuits.len());
    if let Some((_, what)) = published {
        println!("published: {what}");
    }
    let job = OtocJob { set, quantity, method };
    let mut last = Instant::now();
    let mut pairs = 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();
        let p = published.and_then(|(p, _)| p.get(o.index as usize)).copied();
        let tail = p.map_or(String::new(), |p| {
            pairs.push((ours, p));
            format!("  published {p:+.6}  diff {:+.6}", ours - p)
        });
        println!("{:4}  value {ours:+.6}  plain {:+.6}  discarded {:.3e}  bond {:>4}{tail}  {secs:.1}s", o.index, o.plain, o.plain_bound, o.peak);
    })
    .expect("the job runs");
    if pairs.len() > 1 {
        let k = pairs.len() as f64;
        let mad = pairs.iter().map(|(a, b)| (a - b).abs()).sum::<f64>() / k;
        let (ma, mb) = (pairs.iter().map(|p| p.0).sum::<f64>() / k, pairs.iter().map(|p| p.1).sum::<f64>() / k);
        let cov: f64 = pairs.iter().map(|(a, b)| (a - ma) * (b - mb)).sum();
        let (va, vb): (f64, f64) = (pairs.iter().map(|(a, _)| (a - ma) * (a - ma)).sum(), pairs.iter().map(|(_, b)| (b - mb) * (b - mb)).sum());
        let rms_pub = (pairs.iter().map(|(_, b)| b * b).sum::<f64>() / k).sqrt();
        println!("mean |diff| {mad:.3e} · rms published {rms_pub:.3e} · correlation {:.3} · means {ma:+.3e} / {mb:+.3e}", cov / (va * vb).sqrt());
    }
    let work: u64 = receipt.outcomes.iter().map(|o| o.work).sum();
    let peak = receipt.outcomes.iter().map(|o| o.peak).max().unwrap_or(0);
    println!("work {work} (rotations or chain work) · widest bond {peak} · {:.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()));
    }
}

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()
}