//! A neural quantum state for a spin model's ground state, by variational
//! Monte Carlo with stochastic reconfiguration, against exact
//! diagonalisation (up to 20 spins) or DMRG on the same operator sum.
//!
//! ```text
//! cargo run --release --example nqs_ground --features quantum_nqs -- \
//!     --model heisenberg|j1j2|ising --lx L --ly L [--j2 J] [--h H] [--periodic]
//!     [--complex] [--marshall] [--alpha A] [--steps N] [--samples S] [--lr R] [--shift λ]
//!     [--init W] [--exact] [--guide ε]
//!     [--project [--kpi] [--point-group a1|a2|b1|b2 | --reflections a1|a2|b1|b2]
//!     [--odd]] [--dmrg-bond D]
//!     [--threads T] [--exact-final]
//!
//! `--project` projects onto translations of the periodic lattice at
//! momentum 0 (or π with `--kpi`), with `--point-group` on a square lattice
//! onto a representation of its rotations and reflections too, with
//! `--reflections` on any lattice onto one of its half turn and axis
//! reflections (`RectIrrep`), and onto even
//! (or `--odd`) parity under a global spin flip. Up to 20 spins, the
//! reference is then the exact ground energy in that sector. `--exact-final`
//! also evaluates the trained network with every configuration enumerated
//! (at most 24 spins), separating what training left from what sampling
//! adds. `--guide ε` samples `|ψ|²` with a floor `ε` under it, weighted
//! back (`VmcConfig::guide`); projected networks default to `0.1`, others to
//! none.
//! ```
use wai_quantum::quantum_dmrg::{dmrg_mpo, DmrgConfig, Mps, C};
use wai_quantum::quantum_nqs::*;

fn arg<T: std::str::FromStr>(a: &[String], key: &str, default: T) -> T {
    a.iter().position(|x| x == key).and_then(|i| a.get(i + 1)).and_then(|v| v.parse().ok()).unwrap_or(default)
}

fn train<T: Amplitude>(model: SpinModel, rbm: Rbm<T>, cfg: VmcConfig, steps: usize, reference: f64, exact_final: bool) {
    println!("RBM α={}: {} parameters, {} samples a step", rbm.m / rbm.n, rbm.params(), cfg.samples);
    let mut vmc = Vmc::new(model, rbm, cfg);
    let t0 = std::time::Instant::now();
    for it in 1..=steps {
        let e = vmc.step();
        if it % (steps / 10).max(1) == 0 {
            println!("step {it:>5}  E {:.6} ± {:.6}  variance {:.2e}  acceptance {:.2}  relative error {:.2e}  {:.0}s", e.energy, e.error, e.variance, e.acceptance, (e.energy - reference) / reference.abs(), t0.elapsed().as_secs_f64());
        }
    }
    vmc.cfg.samples *= 8;
    let f = vmc.estimate();
    println!("final  E {:.6} ± {:.6}  variance {:.2e}  relative error {:.2e}", f.energy, f.error, f.variance, (f.energy - reference) / reference.abs());
    // The trained network's own energy and variance, every configuration
    // enumerated: what is left once sampling is taken out.
    if exact_final {
        vmc.cfg.exact = true;
        let x = vmc.estimate();
        println!("exact  E {:.8}  variance {:.2e}  relative error {:.2e}  (every configuration)", x.energy, x.variance, (x.energy - reference) / reference.abs());
    }
}

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let kind: String = arg(&a, "--model", "heisenberg".to_string());
    let (lx, ly): (usize, usize) = (arg(&a, "--lx", 4), arg(&a, "--ly", 4));
    let periodic = a.iter().any(|x| x == "--periodic");
    let threads: usize = arg(&a, "--threads", VmcConfig::default().threads);
    let n = lx * ly;
    let nn = if ly == 1 { chain_edges(n, periodic) } else { square_edges(lx, ly, periodic) };
    let nnn = if ly == 1 { chain_next_nearest(n, periodic) } else { square_diagonals(lx, ly, periodic) };
    let model = match kind.as_str() {
        "heisenberg" => SpinModel::heisenberg(n, &nn, 1.0),
        "j1j2" => SpinModel::j1j2(n, &nn, &nnn, 1.0, arg(&a, "--j2", 0.5)),
        "ising" => SpinModel::ising(n, &nn, 1.0, arg(&a, "--h", 3.044)),
        other => panic!("{other}"),
    };
    let sector = if model.conserves_magnetisation() { Some(0) } else { None };
    let point = a.iter().position(|x| x == "--point-group").map(|i| match a.get(i + 1).map(String::as_str) {
        Some("a1") => PointIrrep::A1,
        Some("a2") => PointIrrep::A2,
        Some("b1") => PointIrrep::B1,
        Some("b2") => PointIrrep::B2,
        other => panic!("--point-group: a1, a2, b1 or b2, not {other:?}"),
    });
    let reflections = a.iter().position(|x| x == "--reflections").map(|i| match a.get(i + 1).map(String::as_str) {
        Some("a1") => RectIrrep::A1,
        Some("a2") => RectIrrep::A2,
        Some("b1") => RectIrrep::B1,
        Some("b2") => RectIrrep::B2,
        other => panic!("--reflections: a1, a2, b1 or b2, not {other:?}"),
    });
    let symmetry = a.iter().any(|x| x == "--project").then(|| {
        assert!(periodic, "--project needs --periodic");
        let kpi = a.iter().any(|x| x == "--kpi");
        let base = match (point, reflections) {
            (Some(_), Some(_)) => panic!("--point-group or --reflections, not both"),
            (Some(irrep), None) => {
                assert_eq!(lx, ly, "--point-group needs a square lattice");
                Symmetry::square(lx, kpi, irrep)
            }
            (None, Some(irrep)) => Symmetry::rectangle(lx, ly, kpi, kpi, irrep),
            (None, None) if ly == 1 => Symmetry::ring(n, kpi),
            (None, None) => Symmetry::torus(lx, ly, kpi, kpi),
        };
        base.with_spin_flip(a.iter().any(|x| x == "--odd"))
    });
    let t0 = std::time::Instant::now();
    // A projected network is held to its own sector's ground energy.
    let in_sector = match &symmetry {
        Some(sym) if n <= 24 => Some(model.exact_ground_energy_in(sector, sym).expect("the model is invariant under the group, and the sector is not empty")),
        _ => None,
    };
    let (reference, how) = match (in_sector, if n <= 20 && in_sector.is_none() { model.exact_ground_energy(sector) } else { None }) {
        (Some(e), _) => (e, "exact diagonalisation in the projection's sector".to_string()),
        (None, Some(e)) => (e, "exact diagonalisation".to_string()),
        (None, None) => {
            let bond = arg(&a, "--dmrg-bond", 128usize);
            let r = dmrg_mpo(&model.opsum().mpo(), Mps::random(n, 8, 1), &DmrgConfig { max_bond: bond, sweeps: 6, threads, ..DmrgConfig::default() });
            (r.energy, format!("DMRG, bond {bond}, discarded {:.1e}", r.discarded))
        }
    };
    println!("{kind} {lx}×{ly}{}: reference {reference:.8} ({how}, {:.1}s)", if periodic { " periodic" } else { "" }, t0.elapsed().as_secs_f64());
    let alpha: usize = arg(&a, "--alpha", 2);
    let mut cfg = VmcConfig { samples: arg(&a, "--samples", 1000), chains: 16, magnetisation: sector, lr: arg(&a, "--lr", 0.02), shift: arg(&a, "--shift", 1e-3), exact: a.iter().any(|x| x == "--exact"), threads, ..VmcConfig::default() };
    let steps: usize = arg(&a, "--steps", 600);
    let marshall = a.iter().any(|x| x == "--marshall");
    let exact_final = a.iter().any(|x| x == "--exact-final");
    // Parameters start uniform in ±init. A projected network that starts
    // near zero is nearly the projection of a constant, the ferromagnet: an
    // eigenstate, which reconfiguration barely moves.
    let init: f64 = arg(&a, "--init", if symmetry.is_some() { 0.1 } else { 0.01 });
    // A projected network's amplitude is a sum that cancels near its nodes,
    // where the local energy is large and |ψ|² rarely samples; a guide floor
    // visits there, and training then corrects what it finds.
    cfg.guide = arg(&a, "--guide", if symmetry.is_some() { 0.1 } else { 0.0 });
    if a.iter().any(|x| x == "--complex") {
        let mut rbm: Rbm<C> = Rbm::complex(n, alpha, init, 7);
        if marshall {
            rbm = rbm.with_marshall_sign(checkerboard(lx, ly));
        }
        if let Some(sym) = symmetry {
            rbm = rbm.with_symmetry(sym);
        }
        train(model, rbm, cfg, steps, reference, exact_final);
    } else {
        let mut rbm = Rbm::new(n, alpha, init, 7);
        if marshall {
            rbm = rbm.with_marshall_sign(checkerboard(lx, ly));
        }
        if let Some(sym) = symmetry {
            rbm = rbm.with_symmetry(sym);
        }
        train(model, rbm, cfg, steps, reference, exact_final);
    }
}