//! 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]
//!     [--project [--kpi] [--odd]] [--dmrg-bond D] [--threads T]
//!
//! `--project` projects onto translations of the periodic lattice at
//! momentum 0 (or π with `--kpi`) and onto even (or `--odd`) parity under a
//! global spin flip.
//! ```
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) {
    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());
}

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 t0 = std::time::Instant::now();
    let (reference, how) = match if n <= 20 { model.exact_ground_energy(sector) } else { None } {
        Some(e) => (e, "exact diagonalisation".to_string()),
        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 cfg = VmcConfig { samples: arg(&a, "--samples", 1000), chains: 16, magnetisation: sector, lr: arg(&a, "--lr", 0.02), shift: 1e-3, threads, ..VmcConfig::default() };
    let steps: usize = arg(&a, "--steps", 600);
    let marshall = a.iter().any(|x| x == "--marshall");
    let symmetry = a.iter().any(|x| x == "--project").then(|| {
        assert!(periodic, "--project needs --periodic");
        let kpi = a.iter().any(|x| x == "--kpi");
        let base = if ly == 1 { Symmetry::ring(n, kpi) } else { Symmetry::torus(lx, ly, kpi, kpi) };
        base.with_spin_flip(a.iter().any(|x| x == "--odd"))
    });
    if a.iter().any(|x| x == "--complex") {
        let mut rbm: Rbm<C> = Rbm::complex(n, alpha, 0.01, 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);
    } else {
        let mut rbm = Rbm::new(n, alpha, 0.01, 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);
    }
}