//! The peaked circuits as a matrix-product operator on a chain
//! (`quantum_peaked::chain`): the circuit split at a cut, each side routed onto
//! the chain as it is absorbed, hidden permutations stripped by unswapping, the
//! peak read on `|0…0⟩` by exact conditional decoding and checked by sampling,
//! and a string's amplitude in the whole operator.

#![allow(dead_code)]

use serde_json::{Value, json};
use wai_quantum::quantum_peaked::{ChainOperator, ChainOptions, ChainStep, Layered, Truncation, grow_chain};
use wai_quantum::quantum_spd::RotCircuit;
use wai_quantum::quantum_svd::phases;

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

fn list<T: std::str::FromStr>(args: &[String], key: &str) -> Option<Vec<T>> {
    args.iter().find_map(|a| a.strip_prefix(&format!("{key}=")).map(|v| v.split(',').map(|x| x.parse().unwrap_or_else(|_| panic!("{key}={v}"))).collect()))
}

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

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

/// The options from the arguments, the published settings by default.
fn options(layered: &Layered, args: &[String]) -> ChainOptions {
    let cores = std::thread::available_parallelism().map_or(1, |p| p.get());
    let p = ChainOptions::published(layered, cores);
    ChainOptions {
        trunc: Truncation { chi: arg(args, "chi", p.trunc.chi), cutoff: arg(args, "cutoff", p.trunc.cutoff) },
        split: arg(args, "split", p.split),
        unswap_elements: arg(args, "unswap", p.unswap_elements),
        unswap_rounds: arg(args, "rounds", p.unswap_rounds),
        lookahead: arg(args, "lookahead", p.lookahead),
        route_both: arg(args, "route", "own".to_string()) == "both",
        threads: arg(args, "threads", p.threads),
    }
}

fn opts_json(o: &ChainOptions) -> Value {
    json!({
        "chi": o.trunc.chi, "cutoff": o.trunc.cutoff, "split_blocks": o.split,
        "unswap_elements": o.unswap_elements, "unswap_rounds": o.unswap_rounds, "lookahead": o.lookahead,
        "routing": if o.route_both { "whole sites (row and column together)" } else { "each side's own index (rows after the cut, columns before it)" },
        "threads": o.threads,
    })
}

/// Grow the operator, printing a line every `every` layers.
fn grow(layered: &Layered, o: &ChainOptions, every: usize) -> ChainOperator {
    let mut layer = 0usize;
    let mut biggest = (0usize, 0usize);
    let mut print = |s: &ChainStep| {
        layer += 1;
        biggest = (biggest.0.max(s.max_bond), biggest.1.max(s.elements));
        if layer.is_multiple_of(every.max(1)) || s.unswapped > 0 {
            println!(
                "  layer {layer:>5} {}  done cols {:>4} rows {:>4}  bond {:>5}  elements {:>9}  discarded {:.3e}{}  {:.1} s",
                if s.rows { "R" } else { "C" },
                s.cols_done,
                s.rows_done,
                s.max_bond,
                s.elements,
                s.discarded,
                if s.trials > 0 { format!("  unswapped {} of {} trials", s.unswapped, s.trials) } else { String::new() },
                s.seconds
            );
        }
        true
    };
    grow_chain(layered, o, &mut print).expect("the cut is within the circuit")
}

fn steps_json(op: &ChainOperator, every: usize) -> Value {
    let n = op.steps().len();
    let rows: Vec<Value> = op
        .steps()
        .iter()
        .enumerate()
        .filter(|(i, s)| i.is_multiple_of(every.max(1)) || i + 1 == n || s.unswapped > 0)
        .map(|(i, s)| json!([i + 1, if s.rows { "R" } else { "C" }, s.cols_done, s.rows_done, s.gates, s.swaps, s.max_bond, s.elements, s.discarded, s.unswapped, s.trials, s.seconds]))
        .collect();
    json!({ "columns": ["layer", "side", "cols_done", "rows_done", "gates", "swaps", "max_bond", "elements", "discarded", "unswapped", "unswap_trials", "seconds"], "every": every, "rows": rows })
}

fn svd_phases() -> Value {
    let (ph, rot) = phases();
    json!({ "seconds": ph.iter().map(|(n, s)| json!([n, s])).collect::<Vec<_>>(), "rotations": rot })
}

/// `chain [cutoff=2e-3] [chi=8192] [split=…] [unswap=1000000] [rounds=20]
/// [lookahead=20] [threads=…] [shots=1000] [seed=123] [every=100]`: grow the
/// whole circuit's operator on a chain from the cut, read the peak, sample it,
/// and give the string's amplitude, and the published string's. Writes a receipt.
pub fn chain(circuit: &RotCircuit, map: &[u32], name: &str, path: &str, args: &[String]) {
    let layered = Layered::new(circuit).expect("one- and two-qubit gates");
    let o = options(&layered, args);
    let shots: usize = arg(args, "shots", 1000);
    let seed: u64 = arg(args, "seed", 123);
    let every: usize = arg(args, "every", 100);
    println!("layers          {} block layers, {} blocks; cut after {} blocks", layered.depth(), layered.blocks(), o.split);
    println!(
        "options         χ ≤ {}, cutoff {:e}, unswapping past {} numbers ({} rounds), lookahead {}, routing by {}, {} threads",
        o.trunc.chi,
        o.trunc.cutoff,
        o.unswap_elements,
        o.unswap_rounds,
        o.lookahead,
        if o.route_both { "whole sites" } else { "each side's index" },
        o.threads
    );
    let ((op, peak, amp, counts, pub_amps, flips), secs, energy) = super::common::metered(|| {
        let op = grow(&layered, &o, every);
        let peak = op.peak();
        let amp = op.amplitude(&peak.string).expect("the amplitude");
        let samples = op.sample(shots, seed);
        let mut counts: Vec<(Vec<u8>, usize)> = Vec::new();
        for s in samples {
            match counts.iter_mut().find(|(x, _)| *x == s) {
                Some(c) => c.1 += 1,
                None => counts.push((s, 1)),
            }
        }
        counts.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0)));
        let pub_amps: Vec<(String, f64, &'static str)> = super::structure::published(name)
            .map(|(s, _)| super::structure::orient(s).into_iter().map(|(b, how)| (text(&b), op.amplitude(&b).expect("the amplitude").probability, how)).collect())
            .unwrap_or_default();
        // Each single-bit change of the peak: how far below it the neighbours are.
        let flips: Vec<f64> = (0..peak.string.len())
            .map(|q| {
                let mut t = peak.string.clone();
                t[q] ^= 1;
                op.amplitude(&t).expect("the amplitude").probability
            })
            .collect();
        (op, peak, amp, counts, pub_amps, flips)
    });
    let biggest_bond = op.steps().iter().map(|s| s.max_bond).max().unwrap_or(1);
    let biggest_elements = op.steps().iter().map(|s| s.elements).max().unwrap_or(0);
    let (routed, unswapped, trials) = op.swaps();
    let top = counts.first().map(|(s, c)| (text(s), *c)).unwrap_or_default();
    let flip_max = flips.iter().copied().fold(0.0, f64::max);
    let peak_rss = super::common::peak_rss_bytes();
    println!("operator        {} layers, largest bond {biggest_bond}, most numbers {biggest_elements}, now bond {} and {} numbers", op.steps().len(), op.max_bond(), op.elements());
    println!("swaps           {routed} routed, {unswapped} kept by unswapping of {trials} trials; discarded {:.3e}", op.discarded());
    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} by conditional decoding; weakest bit {:.3}", peak.probability, peak.conditionals.iter().copied().fold(1.0, f64::min));
    println!("amplitude       |⟨s|C|0⟩|² = {:.6e}, ‖C|0⟩‖² = {:.6e}, ratio {:.6}", amp.probability, amp.norm, amp.probability / amp.norm);
    println!("neighbours      the largest single-bit change: |⟨s'|C|0⟩|² = {flip_max:.3e}");
    println!("samples         {shots} shots (seed {seed}): most common {} × {}{}", top.1, top.0, if top.0 == text(&peak.string) { "  = the peak" } else { "" });
    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, peak memory {} MB", peak_rss.map_or("?".into(), |b| format!("{:.0}", b as f64 / 1e6)));
    let (ph, rot) = phases();
    println!("svd phases      {} (worker-seconds); {rot} rotations", ph.iter().map(|(n, t)| format!("{n} {t:.1}")).collect::<Vec<_>>().join(", "));
    let request = json!({
        "instance": format!("peaked_circuit_{}", super::structure::instance(name)),
        "source": super::common::source(path),
        "method": "quantum_peaked::chain: the circuit as two-qubit blocks split at a cut; the blocks before it absorbed on the operator's columns from the cut outward, those after it on its rows, the side with fewer gates applied going next; each side routed onto a chain a layer at a time (the front's gates on neighbouring sites applied, otherwise SWAPs of rows or columns chosen by a lookahead distance heuristic); every two-site operation a TEBD update of the operator as a matrix-product state of local dimension 4, decomposed by quantum_svd and truncated at the cutoff (relative to the largest singular value) and bond cap; when the operator holds more than the threshold of numbers, a SWAP of both indices, the columns or the rows tried on every bond (alternating parities, singular values alone) and kept as a relabelling where it shrinks the bond; the peak read on |0…0⟩ by exact conditional decoding along the chain, checked by exact sampling",
        "options": opts_json(&o),
        "shots": shots,
        "seed": seed,
        "qubits": map,
    });
    let result = json!({
        "block_layers": layered.depth(),
        "blocks": layered.blocks(),
        "two_qubit_gates": layered.two_qubit_gates(),
        "largest_bond": biggest_bond,
        "most_elements": biggest_elements,
        "final_bonds": op.bonds(),
        "swaps_routed": routed,
        "swaps_unswapped": unswapped,
        "unswap_trials": trials,
        "discarded": op.discarded(),
        "string": text(&peak.string),
        "string_order": "bit q is qubit q of the register as compacted (the device qubit map[q])",
        "probability_decoded": peak.probability,
        "conditionals": peak.conditionals,
        "amplitude": { "re": amp.re, "im": amp.im, "probability": amp.probability, "norm": amp.norm },
        "single_bit_neighbours_probability": flips,
        "samples": { "shots": shots, "most_common": top.0, "count": top.1, "distinct": counts.len() },
        "published": pub_amps.iter().map(|(s, p, how)| json!({"string": s, "order": how, "probability": p})).collect::<Vec<_>>(),
        "published_source": super::structure::published(name).map(|(_, w)| w),
        "scale": { "mantissa": op.scale().mantissa, "exponent": op.scale().exponent },
        "steps": steps_json(&op, every),
        "svd_phases": svd_phases(),
        "wall_seconds": secs,
        "peak_rss_bytes": peak_rss,
        "energy_figure": energy,
    });
    let out = super::common::receipt_path(&format!("peaked_{name}_chain_cut{}_cutoff{:e}_{}.json", o.split, o.trunc.cutoff, if o.route_both { "both" } else { "own" }));
    assert!(super::common::verifies(&super::common::seal(request, result, Some(&out))));
    println!("receipt         {}", out.display());
}

/// `chainverify STRING [cutoffs=2e-3,1e-3] [splits=…] [chi=…] …`: the
/// string's amplitude `⟨s|C|0…0⟩` in the whole circuit, the operator grown on
/// a chain at each cutoff and from each cut, closed with the string: no
/// readout, no sampling. Writes a receipt.
pub fn verify(circuit: &RotCircuit, map: &[u32], name: &str, path: &str, string: &str, args: &[String]) {
    let layered = Layered::new(circuit).expect("one- and two-qubit gates");
    let base = options(&layered, args);
    let cutoffs: Vec<f64> = list(args, "cutoffs").unwrap_or_else(|| vec![base.trunc.cutoff]);
    let splits: Vec<usize> = list(args, "splits").unwrap_or_else(|| vec![base.split]);
    let every: usize = arg(args, "every", 500);
    let b = bits(string);
    assert_eq!(b.len(), circuit.n as usize, "a string of {} bits", circuit.n);
    let mut rows = Vec::new();
    for &split in &splits {
        for &cutoff in &cutoffs {
            let o = ChainOptions { trunc: Truncation { chi: base.trunc.chi, cutoff }, split, ..base };
            println!("cut after {split} blocks, cutoff {cutoff:e}");
            let ((op, amp, flips), secs, energy) = super::common::metered(|| {
                let op = grow(&layered, &o, every);
                let amp = op.amplitude(&b).expect("the amplitude");
                let flips: Vec<f64> = (0..b.len())
                    .map(|q| {
                        let mut t = b.clone();
                        t[q] ^= 1;
                        op.amplitude(&t).expect("the amplitude").probability
                    })
                    .collect();
                (op, amp, flips)
            });
            let biggest = op.steps().iter().map(|s| s.max_bond).max().unwrap_or(1);
            let most = op.steps().iter().map(|s| s.elements).max().unwrap_or(0);
            let flip_max = flips.iter().copied().fold(0.0, f64::max);
            println!(
                "  |⟨s|C|0⟩|² = {:.6e}  ‖C|0⟩‖² = {:.6e}  ratio {:.6}  largest neighbour {flip_max:.3e}  bond ≤ {biggest}  numbers ≤ {most}  discarded {:.3e}  {secs:.1} s",
                amp.probability,
                amp.norm,
                amp.probability / amp.norm,
                op.discarded()
            );
            let (routed, unswapped, trials) = op.swaps();
            rows.push(json!({
                "split_blocks": split, "cutoff": cutoff, "chi": o.trunc.chi,
                "re": amp.re, "im": amp.im, "probability": amp.probability, "norm": amp.norm,
                "largest_single_bit_neighbour": flip_max,
                "largest_bond": biggest, "most_elements": most, "discarded": op.discarded(),
                "swaps_routed": routed, "swaps_unswapped": unswapped, "unswap_trials": trials,
                "wall_seconds": secs, "peak_rss_bytes": super::common::peak_rss_bytes(), "energy_figure": energy,
            }));
        }
    }
    let request = json!({
        "instance": format!("peaked_circuit_{}", super::structure::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 as a matrix-product operator on a chain (quantum_peaked::chain) from each cut 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 sampling",
        "string": string,
        "string_order": "bit q is qubit q of the register as compacted (the device qubit map[q])",
        "options": opts_json(&base),
        "cutoffs": cutoffs,
        "splits": splits,
        "qubits": map,
    });
    let result = json!({ "sweep": rows, "svd_phases": svd_phases() });
    let out = super::common::receipt_path(&format!("peaked_{name}_chainverify.json"));
    assert!(super::common::verifies(&super::common::seal(request, result, Some(&out))));
    println!("receipt         {}", out.display());
}