//! The peaked circuits read through the structure they were built with
//! (`quantum_peaked`): the junction found by probing, the operator grown
//! outward from it, the peak read by exact contraction, and each string's
//! amplitude in the whole circuit at a sweep of truncations.

#![allow(dead_code)]

use serde_json::{Value, json};
use wai_quantum::quantum_peaked::{Layered, Truncation, Unswap, best_centre, cancel, cancel_unswapping, probe, scan, scan_unswapping};
use wai_quantum::quantum_spd::RotCircuit;

/// The peaks the tracker's classical submissions published, as printed there,
/// with where each was printed. The 49×5072 string is printed qubit 48 first;
/// the 56×1917 string as the submission's notebook gives it, compared here in
/// both orders.
pub const PUBLISHED: &[(&str, &str, &str)] = &[
    (
        "49x5072",
        "1001100011100000001010100001011111111011110010110",
        "tracker issue #105 (TNO + iterative cancellation), the string as submitted",
    ),
    (
        "P9_Hqap_56x1917",
        "01101110111001100000100000001010011100101101010111110111",
        "the notebook issue #106 links (MPO + unswapping), the string it checks its sample against",
    ),
];

/// The published string for `name`, if one was printed, and its source.
pub fn published(name: &str) -> Option<(&'static str, &'static str)> {
    PUBLISHED.iter().find(|(n, _, _)| *n == name).map(|(_, s, w)| (*s, *w))
}

fn bits(s: &str) -> Vec<u8> {
    s.bytes().map(|b| u8::from(b == b'1')).collect()
}

fn text(b: &[u8]) -> String {
    b.iter().map(|&x| if x == 1 { '1' } else { '0' }).collect()
}

/// The arguments after the mode, as `key=value`.
fn arg<T: std::str::FromStr>(args: &[String], key: &str, default: T) -> T {
    args.iter()
        .find_map(|a| a.strip_prefix(&format!("{key}=")).map(|v| v.parse().unwrap_or_else(|_| panic!("{key}={v} does not read"))))
        .unwrap_or(default)
}

/// Bits in the string for qubit order `q = 0, 1, …` from a string printed
/// either way, and which way matched (the order with the larger amplitude).
pub fn orient(printed: &str) -> [(Vec<u8>, &'static str); 2] {
    let fwd = bits(printed);
    let mut rev = fwd.clone();
    rev.reverse();
    [(fwd, "as printed, first character qubit 0"), (rev, "reversed, last character qubit 0")]
}

/// `structure`: scan every centre, grow the operator from the smallest, read
/// the peak; then the string's amplitude in that operator, and the published
/// string's when there is one. Writes a receipt.
pub fn structure(circuit: &RotCircuit, map: &[u32], name: &str, path: &str, args: &[String]) {
    let window: usize = arg(args, "window", 4);
    let chi: usize = arg(args, "chi", 16);
    let cutoff: f64 = arg(args, "cutoff", 1e-2);
    let probe_chi: usize = arg(args, "probe_chi", 8);
    let read_chi: usize = arg(args, "read_chi", 8);
    let read_cutoff: f64 = arg(args, "read_cutoff", 1e-2);
    let limit: usize = arg(args, "limit", 1 << 26);
    let forced: Option<usize> = args.iter().find_map(|a| a.strip_prefix("centre=").map(|v| v.parse().expect("a centre")));
    // unswap=T: unswapping wherever a bond reaches T (0: none), rounds=R passes a step.
    let threshold: usize = arg(args, "unswap", 0);
    let rounds: usize = arg(args, "rounds", 4);
    let unswap = (threshold > 0).then_some(Unswap { threshold, rounds });
    let trunc = Truncation { chi, cutoff };
    let ((layered, probes, centre, op, peak, amp, pub_amps), secs, energy) = super::common::metered(|| {
        let layered = Layered::new(circuit).expect("one- and two-qubit gates");
        let probe = Truncation { chi: probe_chi, cutoff };
        // scan=0 with a centre given: no scan.
        let centres: Vec<usize> = if forced.is_some() && arg(args, "scan", 1) == 0 { Vec::new() } else { (0..=2 * layered.depth()).collect() };
        let probes = match unswap {
            Some(u) => scan_unswapping(&layered, centres, window, probe, u),
            None => scan(&layered, centres, window, probe),
        }
        .expect("centres in range");
        let centre = forced.unwrap_or_else(|| best_centre(&layered, &probes).expect("a centre"));
        let op = match unswap {
            Some(u) => cancel_unswapping(&layered, centre, trunc, 1, u),
            None => cancel(&layered, centre, trunc, 1),
        }
        .expect("the centre");
        let peak = op.peak(Truncation { chi: read_chi, cutoff: read_cutoff }, 2, limit).expect("a readout within the limit");
        let amp = op.amplitude(&peak.string, limit).expect("the amplitude");
        let pub_amps: Vec<(String, f64, &'static str)> = published(name)
            .map(|(s, _)| orient(s).into_iter().map(|(b, how)| (text(&b), op.amplitude(&b, limit).expect("the amplitude").probability, how)).collect())
            .unwrap_or_default();
        (layered, probes, centre, op, peak, amp, pub_amps)
    });
    let mut ranked = probes.clone();
    ranked.sort_by(|a, b| a.step.log_size.total_cmp(&b.step.log_size).then(a.centre.cmp(&b.centre)));
    println!("layers          {} block layers, {} blocks, {} two-qubit gates", layered.depth(), layered.blocks(), layered.two_qubit_gates());
    println!("scan            {} centres, window {window}, χ {probe_chi}, cutoff {cutoff:e}", probes.len());
    for p in ranked.iter().take(5) {
        println!(
            "  centre {:>4} (layer {:>6.1})  log size {:>7.3}  max bond {:>3}  bonds {:>3}  discarded {:.3e}",
            p.centre,
            p.centre as f64 / 2.0,
            p.step.log_size,
            p.step.max_bond,
            p.step.bonds,
            p.step.discarded
        );
    }
    let median = {
        let mut v: Vec<f64> = probes.iter().map(|p| p.step.log_size).collect();
        v.sort_by(f64::total_cmp);
        v.get(v.len() / 2).copied().unwrap_or(0.0)
    };
    println!("  median log size over every centre {median:.3}");
    let peak_bond = op.steps().iter().map(|s| s.max_bond).max().unwrap_or(1);
    let peak_elems = op.steps().iter().map(|s| s.elements).max().unwrap_or(0);
    println!("centre          {centre} (layer {:.1}), χ {chi}, cutoff {cutoff:e}", centre as f64 / 2.0);
    let n = op.steps().len();
    for (i, s) in op.steps().iter().enumerate() {
        if i % 10 == 0 || i + 1 == n {
            println!(
                "  step {:>4}  rows {:>4} cols {:>4}  max bond {:>3}  bonds {:>3}  elements {:>7}  discarded {:.3e}",
                i, s.right, s.left, s.max_bond, s.bonds, s.elements, s.discarded
            );
        }
    }
    let weakest = peak.marginals.iter().map(|p| (p - 0.5).abs() * 2.0).fold(f64::INFINITY, f64::min);
    println!("peak            {}  (bit q = qubit q)", text(&peak.string));
    println!("                printed qubit {} first: {}", circuit.n - 1, text(&peak.string.iter().rev().copied().collect::<Vec<_>>()));
    println!("probability     {:.6} in the read state (bond {}), weakest bit margin {weakest:.3}", peak.probability, peak.max_bond);
    println!("amplitude       |⟨s|C|0⟩|² = {:.6e}, ‖C|0⟩‖² = {:.6e} in the operator, ratio {:.6}", amp.probability, amp.norm, amp.probability / amp.norm);
    for (s, p, how) in &pub_amps {
        let same = *s == text(&peak.string);
        println!("published       {s} ({how}): |⟨s|C|0⟩|² = {p:.6e}{}", if same { "  = the peak read here" } else { "" });
    }
    println!(
        "time            {secs:.1} s, discarded {:.3e}, largest bond {peak_bond}, most elements {peak_elems}, swaps kept {}",
        op.discarded(),
        op.swaps()
    );
    let request = json!({
        "instance": format!("peaked_circuit_{}", instance(name)),
        "source": super::common::source(path),
        "method": "quantum_peaked: the circuit as layers of two-qubit blocks; every centre probed by growing the circuit's operator a few layers each way (the smallest operator marks where a block meets its approximate inverse); the whole circuit grown outward from that centre as an operator on the graph its gates build (rows after the centre, columns before it), simple-update truncation; the peak read on |0…0⟩ by exact contraction of the network",
        "window": window,
        "probe_chi": probe_chi,
        "chi": chi,
        "cutoff": cutoff,
        "read_chi": read_chi,
        "read_cutoff": read_cutoff,
        "unswap": unswap.map(|u| json!({"threshold": u.threshold, "rounds": u.rounds})),
        "qubits": map,
    });
    let result = json!({
        "block_layers": layered.depth(),
        "blocks": layered.blocks(),
        "two_qubit_gates": layered.two_qubit_gates(),
        "scan_best": ranked.iter().take(5).map(|p| json!({"centre": p.centre, "log_size": p.step.log_size, "max_bond": p.step.max_bond, "bonds": p.step.bonds, "discarded": p.step.discarded})).collect::<Vec<_>>(),
        "scan_median_log_size": median,
        "centre": centre,
        "steps": op.steps().iter().map(|s| json!([s.right, s.left, s.max_bond, s.bonds, s.elements, s.discarded])).collect::<Vec<_>>(),
        "steps_columns": ["rows", "cols", "max_bond", "bonds", "elements", "discarded"],
        "discarded": op.discarded(),
        "swaps_kept": op.swaps(),
        "string": text(&peak.string),
        "string_order": "bit q is qubit q of the register as compacted (the device qubit map[q])",
        "marginals_one": peak.marginals,
        "probability_read": peak.probability,
        "amplitude": { "re": amp.re, "im": amp.im, "probability": amp.probability, "norm": amp.norm },
        "published": pub_amps.iter().map(|(s, p, how)| json!({"string": s, "order": how, "probability": p})).collect::<Vec<_>>(),
        "published_source": published(name).map(|(_, w)| w),
        "wall_seconds": secs,
        "energy_figure": energy,
    });
    let out = super::common::receipt_path(&format!("peaked_{name}_structure_chi{chi}.json"));
    assert!(super::common::verifies(&super::common::seal(request, result, Some(&out))));
    println!("receipt         {}", out.display());
}

/// `profile centre=H [steps=…] [chi=…] [cutoff=…] [unswap=T]`: the operator
/// grown from one centre for each number of steps in `steps=` (default
/// 2,4,8,16,32), each printed as it finishes, so a run that will not finish
/// shows where it grows.
pub fn profile(circuit: &RotCircuit, args: &[String]) {
    let layered = Layered::new(circuit).expect("one- and two-qubit gates");
    let centre: usize = args.iter().find_map(|a| a.strip_prefix("centre=").map(|v| v.parse().expect("a centre"))).unwrap_or(layered.depth());
    let chi: usize = arg(args, "chi", 16);
    let cutoff: f64 = arg(args, "cutoff", 1e-2);
    let threshold: usize = arg(args, "unswap", 0);
    let unswap = (threshold > 0).then_some(Unswap { threshold, rounds: arg(args, "rounds", 4) });
    let steps: Vec<usize> = args
        .iter()
        .find_map(|a| a.strip_prefix("steps=").map(|v| v.split(',').map(|x| x.parse().expect("a step count")).collect()))
        .unwrap_or_else(|| vec![2, 4, 8, 16, 32]);
    println!("layers          {} block layers, {} blocks; centre {centre}, χ {chi}, cutoff {cutoff:e}, unswap {threshold}", layered.depth(), layered.blocks());
    for k in steps {
        let started = std::time::Instant::now();
        match probe(&layered, centre, k, Truncation { chi, cutoff }, unswap) {
            Ok(op) => {
                let s = op.steps().last().copied();
                if let Some(s) = s {
                    println!(
                        "  {k:>3} steps  rows {:>3} cols {:>3}  max bond {:>3}  bonds {:>4}  elements {:>9}  log size {:>8.2}  swaps {:>4}  {:.1} s",
                        s.right,
                        s.left,
                        s.max_bond,
                        s.bonds,
                        s.elements,
                        s.log_size,
                        op.swaps(),
                        started.elapsed().as_secs_f64()
                    );
                }
            }
            Err(e) => println!("  {k:>3} steps  {e}  {:.1} s", started.elapsed().as_secs_f64()),
        }
    }
}

/// `junction [window=2] [chi=8] [cutoff=…] [unswap=T]`: the scan alone, every
/// centre's probe printed in order, then the ten smallest.
pub fn junction(circuit: &RotCircuit, args: &[String]) {
    let layered = Layered::new(circuit).expect("one- and two-qubit gates");
    let window: usize = arg(args, "window", 2);
    let trunc = Truncation { chi: arg(args, "chi", 8), cutoff: arg(args, "cutoff", 1e-2) };
    let threshold: usize = arg(args, "unswap", 0);
    let started = std::time::Instant::now();
    let probes = match threshold {
        0 => scan(&layered, 0..=2 * layered.depth(), window, trunc),
        t => scan_unswapping(&layered, 0..=2 * layered.depth(), window, trunc, Unswap { threshold: t, rounds: arg(args, "rounds", 4) }),
    }
    .expect("centres in range");
    for p in &probes {
        println!("centre {:>4}  log size {:>8.2}  max bond {:>3}  bonds {:>4}", p.centre, p.step.log_size, p.step.max_bond, p.step.bonds);
    }
    let mut ranked = probes.clone();
    ranked.sort_by(|a, b| a.step.log_size.total_cmp(&b.step.log_size).then(a.centre.cmp(&b.centre)));
    println!("smallest:");
    for p in ranked.iter().take(10) {
        println!("  centre {:>4}  log size {:>8.2}  max bond {:>3}  bonds {:>4}", p.centre, p.step.log_size, p.step.max_bond, p.step.bonds);
    }
    println!("{} centres in {:.1} s", probes.len(), started.elapsed().as_secs_f64());
}

/// The instance's file stem after `peaked_circuit_`.
pub fn instance(name: &str) -> String {
    if name.starts_with('P') { name.to_string() } else { format!("heavy_hex_{name}") }
}

/// `verify STRING`: the string's amplitude `⟨s|C|0⟩` in the whole circuit,
/// the operator grown from `centre=` (default the scan's) at each cutoff in
/// `cutoffs=` with bond cap `chi=`, the network closed and contracted exactly.
/// A cutoff of 0 truncates nothing below the cap.
pub fn verify(circuit: &RotCircuit, map: &[u32], name: &str, path: &str, string: &str, args: &[String]) -> Value {
    let chi: usize = arg(args, "chi", 64);
    let limit: usize = arg(args, "limit", 1 << 26);
    let cutoffs: Vec<f64> = args
        .iter()
        .find_map(|a| a.strip_prefix("cutoffs=").map(|v| v.split(',').map(|x| x.parse().expect("a cutoff")).collect()))
        .unwrap_or_else(|| vec![1e-2, 1e-3, 1e-4]);
    let layered = Layered::new(circuit).expect("one- and two-qubit gates");
    let centre: usize = args.iter().find_map(|a| a.strip_prefix("centre=").map(|v| v.parse().expect("a centre"))).unwrap_or_else(|| {
        let probes = scan(&layered, 0..=2 * layered.depth(), 4, Truncation { chi: 8, cutoff: 1e-2 }).expect("centres");
        best_centre(&layered, &probes).expect("a centre")
    });
    let b = bits(string);
    assert_eq!(b.len(), circuit.n as usize, "a string of {} bits", circuit.n);
    let mut rows = Vec::new();
    for &cut in &cutoffs {
        let (res, secs, energy) = super::common::metered(|| {
            let op = cancel(&layered, centre, Truncation { chi, cutoff: cut }, 1).expect("the centre");
            let amp = op.amplitude(&b, limit);
            // The largest bond over the first nine tenths of the steps (the
            // cancelled middle) and over the rest (the outer layers).
            let n = op.steps().len();
            let split = n - n / 10;
            let middle = op.steps()[..split].iter().map(|s| s.max_bond).max().unwrap_or(1);
            let outer = op.steps()[split..].iter().map(|s| s.max_bond).max().unwrap_or(1);
            (op.discarded(), middle.max(outer), (middle, outer), amp)
        });
        let (discarded, bond, (middle, outer), amp) = res;
        match amp {
            Ok(a) => {
                println!(
                    "cutoff {cut:<8e} χ ≤ {chi:<4} bond {middle:>3} middle, {outer:>3} outer  |⟨s|C|0⟩|² = {:.6e}  ‖C|0⟩‖² = {:.6e}  ratio {:.6}  discarded {discarded:.3e}  {secs:.1} s",
                    a.probability,
                    a.norm,
                    a.probability / a.norm
                );
                rows.push(json!({"cutoff": cut, "chi": chi, "largest_bond": bond, "middle_bond": middle, "outer_bond": outer, "re": a.re, "im": a.im, "probability": a.probability, "norm": a.norm, "discarded": discarded, "wall_seconds": secs, "energy_figure": energy}));
            }
            Err(e) => {
                println!("cutoff {cut:<8e} χ ≤ {chi:<4} largest bond {bond:>4}  the closing contraction: {e}  {secs:.1} s");
                rows.push(json!({"cutoff": cut, "chi": chi, "largest_bond": bond, "refused": e.to_string(), "discarded": discarded, "wall_seconds": secs, "energy_figure": energy}));
            }
        }
    }
    let request = json!({
        "instance": format!("peaked_circuit_{}", instance(name)),
        "source": super::common::source(path),
        "method": "the amplitude ⟨s|C|0…0⟩ of the given string in the whole circuit: its operator grown outward from the centre (quantum_peaked::cancel) at each cutoff, closed with ⟨s| on the rows and |0⟩ on the columns and contracted exactly, the scale carried through every update; no readout, no marginals",
        "string": string,
        "string_order": "bit q is qubit q of the register as compacted (the device qubit map[q])",
        "centre": centre,
        "qubits": map,
    });
    let result = json!({ "sweep": rows });
    let out = super::common::receipt_path(&format!("peaked_{name}_verify_centre{centre}_chi{chi}.json"));
    let doc = super::common::seal(request, result, Some(&out));
    assert!(super::common::verifies(&doc));
    println!("receipt         {}", out.display());
    doc
}