//! The tracker's 49-qubit guided sparse ground-state problem by DMRG, against
//! the energies published for it.
//!
//! `cargo run --release --example tracker_guided --features quantum_engine,quantum_receipt,quantum_dmrg -- <tracker checkout> [guide | random:SEED] [bond ...]`
//!
//! The Hamiltonian is 419 Pauli terms on 49 qubits, at most two qubits each;
//! the problem gives a basis state that overlaps the ground state. DMRG runs on
//! a chain ordered breadth first through the coupling graph, from that state or
//! from a pseudo-random one, at each bond dimension given: the energy is
//! variational, an upper
//! bound on the ground energy up to rounding, and the sweep energies, the
//! discarded weight and the change between bond dimensions are the
//! convergence evidence.
//!
//! Each run writes a signed receipt to `RECEIPTS_DIR` (default `receipts`).

#[path = "tracker/common.rs"]
mod common;
#[path = "tracker/guided.rs"]
mod guided;

use serde_json::json;
use wai_quantum::quantum_dmrg::{DmrgConfig, Mps, dmrg_mpo, expectation};

const HAMILTONIAN: &str = "data/variational-problems/hamiltonians/guided_sparse_ground_state_problem/49Q_v1/49Q_v1_hamiltonian.npy";
const GUIDE: &str = "data/variational-problems/hamiltonians/guided_sparse_ground_state_problem/49Q_v1/49Q_v1_initial_state_index.npy";

/// (what produced it, the energy, the tracker submission it is from)
const PUBLISHED: &[(&str, f64, u32)] = &[
    ("device: sample-based Krylov diagonalization", -85.58662506986455, 136),
    ("DMRG", -85.58662506986444, 142),
    ("truncated Arnoldi", -85.58662506986455, 137),
    ("selected CI (heat-bath)", -85.57662506986455, 144),
    ("selected CI (perturbative)", -85.57662506986455, 143),
    ("selected CI (adaptive sampling)", -85.55662506986455, 145),
];

fn main() {
    let mut args = std::env::args().skip(1);
    let root = common::tracker_root(args.next());
    // `guide` starts from the guiding basis state; `random:SEED` from a
    // pseudo-random state of the run's bond dimension.
    let from = args.next().unwrap_or_else(|| "random:1".into());
    let bonds: Vec<usize> = args.map(|a| a.parse().expect("a bond dimension")).collect();
    let bonds = if bonds.is_empty() { vec![16, 32, 64] } else { bonds };
    let (n, terms) = guided::hamiltonian(&common::read_pinned(&root, HAMILTONIAN));
    let index = guided::guide(&common::read_pinned(&root, GUIDE));
    let order = guided::bfs_order(n, &terms);
    let mpo = guided::op_sum(n, &terms, &order).mpo();
    let start = guided::basis_mps(index, &order);
    println!("instance        guided_sparse_ground_state_problem_49Q_v1");
    println!("qubits          {n}, {} terms, MPO bond {}", terms.len(), mpo.max_bond());
    println!("guide           {index}: ⟨H⟩ = {:.12} (by its diagonal terms), {:.12} (by the MPO)", guided::diagonal_energy(&terms, index), expectation(&mpo, &start));
    for (label, e, issue) in PUBLISHED {
        println!("published       {e:.14} {label} (submission {issue})");
    }
    let threads = std::thread::available_parallelism().map_or(1, |t| t.get());
    for bond in bonds {
        let cfg = DmrgConfig { max_bond: bond, cutoff: 1e-14, sweeps: 8, krylov: 32, restarts: 4, tol: 1e-10, seed: 1, threads };
        let initial = match from.strip_prefix("random:") {
            Some(seed) => Mps::random(n, bond, seed.parse().expect("a seed")),
            None if from == "guide" => start.clone(),
            None => panic!("start from `guide` or `random:SEED`, not {from}"),
        };
        let (r, secs, energy) = common::metered(|| dmrg_mpo(&mpo, initial, &cfg));
        let bond_used = r.mps.dims.iter().copied().max().unwrap_or(1);
        let max_entropy = r.entropies.iter().copied().fold(0.0, f64::max);
        println!(
            "bond {bond:<5} E = {:.14}  E − published = {:+.3e}  sweeps {:?}  discarded {:.1e}  bond used {bond_used}  S_max {max_entropy:.3}  {secs:.1} s  {}",
            r.energy,
            r.energy - PUBLISHED[0].1,
            r.sweep_energies.iter().map(|e| format!("{e:.12}")).collect::<Vec<_>>(),
            r.discarded,
            energy.get("joules_micro").map_or("unmetered".to_string(), |j| format!("{:.1} J", j.as_u64().unwrap_or(0) as f64 / 1e6)),
        );
        let request = json!({
            "instance": "guided_sparse_ground_state_problem_49Q_v1",
            "source": [common::source(HAMILTONIAN), common::source(GUIDE)],
            "method": "two-site DMRG on the exact MPO (quantum_dmrg::dmrg_mpo)",
            "start": from,
            "chain_order": order,
            "max_bond": bond,
            "cutoff": cfg.cutoff,
            "sweeps": cfg.sweeps,
            "krylov": cfg.krylov,
            "restarts": cfg.restarts,
            "tol": cfg.tol,
            "threads": threads,
            "published": PUBLISHED.iter().map(|(label, e, issue)| json!({ "label": label, "value": e, "submission": issue })).collect::<Vec<_>>(),
        });
        let result = json!({
            "energy": r.energy,
            "energy_bits": common::bits(r.energy),
            "sweep_energies": r.sweep_energies,
            "discarded": r.discarded,
            "bond_used": bond_used,
            "max_entropy": max_entropy,
            "wall_seconds": secs,
            "energy_figure": energy,
        });
        let out = common::receipt_path(&format!("guided_49Q_v1_{}_bond{bond}.json", from.replace(':', "")));
        let doc = common::seal(request, result, Some(&out));
        assert!(common::verifies(&doc));
        println!("receipt         {}", out.display());
    }
}