wai-quantum 0.3.35

//! Ground states of spin chains by DMRG — `wai.quantum.dmrg`.
//!
//! Most claims of a quantum advantage in simulating matter are claims that
//! classical methods cannot reach the answer. For one-dimensional systems the
//! density-matrix renormalization group usually can, with the error
//! controlled by the entanglement it keeps. This module is that method,
//! deterministic and held to exact answers.
//!
//! - **Hamiltonians.** [`Chain`] is a nearest-neighbour spin-½ chain: real
//!   two-site terms on every bond and a one-site term on every site. Presets
//!   are the XXZ Heisenberg chain in a field and the transverse-field Ising
//!   chain. [`Chain::mpo`] writes it as a matrix-product operator, the finite
//!   state machine of its terms.
//! - **Two-site DMRG.** [`dmrg`] sweeps a matrix-product state left and right.
//!   - At each bond it finds the lowest eigenvector of the effective
//!     Hamiltonian by Lanczos, with full re-orthogonalisation.
//!   - It splits that vector by the eigenvectors of its reduced density
//!     matrix. It keeps at most `max_bond` of them and drops those below a
//!     weight cutoff.
//!   - The discarded weight, and the entanglement entropy at every bond, are
//!     reported.
//! - **Exact referees.**
//!   - [`exact_ground_energy`] runs Lanczos on the full `2ⁿ` space for short
//!     chains.
//!   - [`ising_exact_energy`] is the transverse-field Ising chain solved as
//!     free fermions: the ground energy of an open chain of any length, from
//!     the eigenvalues of an `n × n` matrix.
//!
//! The engine is generic over [`Scalar`]: real here, complex for time
//! evolution. The eigensolver is Householder tridiagonalisation followed by
//! the implicit QL iteration, with hypotenuses taken in-crate;
//! [`hermitian_eigen`] reduces a Hermitian matrix to a real tridiagonal one
//! first and shares the iteration. The effective Hamiltonian
//! is applied in blocks of rows across threads, each row in the same order.
//! Every sum runs in a fixed order, so a run is the same bits on every
//! machine and at any thread count.
//!
//! # Checked
//!
//! - **The eigensolvers.** They reconstruct random symmetric and Hermitian
//!   matrices up to 40 × 40 to `10⁻¹²` with orthonormal vectors.
//! - **The free-fermion solution.** It equals exact diagonalisation of the
//!   Ising chain at three couplings, for 4, 7 and 10 sites.
//! - **DMRG against exact diagonalisation.** On Heisenberg chains (isotropic,
//!   and anisotropic in a field) and Ising chains of 10–12 sites, energies
//!   agree to `10⁻⁸`.
//! - **At scale** (`dmrg_chain` example; energy error per site against the
//!   exact answer):
//!
//!   | chain                  | sites | bond | error per site | discarded weight |
//!   |------------------------|-------|------|----------------|------------------|
//!   | Ising, critical        | 100   | ≤ 32 | 1.1·10⁻¹¹      | 2.5·10⁻¹²        |
//!   | Ising, critical        | 200   | ≤ 64 | 1.4·10⁻¹¹      | 3.0·10⁻¹²        |
//!   | Heisenberg             | 20    | 64   | 4.4·10⁻¹²      | 4.4·10⁻¹²        |
//!
//!   The 100-site Heisenberg chain has no exact solution at that length. Its
//!   energy falls from above as the bond grows: −44.127739250 at bond 64
//!   (discarded 4.9·10⁻⁹), then −44.127739889 at bond 128 (2.5·10⁻¹¹).
//! - **Variational and repeatable.** The energy never goes below the exact
//!   value, never rises from one sweep to the next, and repeats bit for bit at
//!   any thread count.
//!
//! # Honest boundaries
//!
//! - **One dimension, nearest neighbours, real Hamiltonians.** Longer-range
//!   couplings, complex terms and two-dimensional lattices are not handled.
//!   Time evolution is the TDVP module's, on this engine.
//! - **No symmetry sectors.** Conserved quantities are not used to block the
//!   matrices, so a run costs more than a symmetric code would.
//! - **Variational.** The energy is an upper bound, converging from above as
//!   `max_bond` grows. Critical chains converge slowest.

pub use crate::linalg::C;
use crate::repro::ln;

// ---------------------------------------------------------------------------
// Scalars
// ---------------------------------------------------------------------------

/// The numbers a matrix-product state is made of: real for the ground states
/// of real Hamiltonians, complex for time evolution. Real arithmetic runs in
/// exactly the order it did before the engine was generic, so its bits are
/// unchanged.
pub trait Scalar: Copy + PartialEq + core::fmt::Debug + Send + Sync + 'static {
    const ZERO: Self;
    const ONE: Self;
    fn add(self, o: Self) -> Self;
    fn sub(self, o: Self) -> Self;
    fn mul(self, o: Self) -> Self;
    fn conj(self) -> Self;
    /// `self · c` for a real `c`.
    fn scale(self, c: f64) -> Self;
    /// `|self|²`.
    fn norm2(self) -> f64;
    fn re(self) -> f64;
    fn is_zero(self) -> bool;
    /// Eigenvalues (ascending) and eigenvectors (column `i` of the returned
    /// row-major matrix) of a Hermitian `n × n` matrix.
    fn eigh(a: &[Self], n: usize) -> (Vec<f64>, Vec<Self>);
}

impl Scalar for f64 {
    const ZERO: f64 = 0.0;
    const ONE: f64 = 1.0;
    #[inline]
    fn add(self, o: f64) -> f64 {
        self + o
    }
    #[inline]
    fn sub(self, o: f64) -> f64 {
        self - o
    }
    #[inline]
    fn mul(self, o: f64) -> f64 {
        self * o
    }
    #[inline]
    fn conj(self) -> f64 {
        self
    }
    #[inline]
    fn scale(self, c: f64) -> f64 {
        self * c
    }
    #[inline]
    fn norm2(self) -> f64 {
        self * self
    }
    #[inline]
    fn re(self) -> f64 {
        self
    }
    #[inline]
    fn is_zero(self) -> bool {
        self == 0.0
    }
    fn eigh(a: &[f64], n: usize) -> (Vec<f64>, Vec<f64>) {
        symmetric_eigen(a, n)
    }
}

impl Scalar for C {
    const ZERO: C = C { re: 0.0, im: 0.0 };
    const ONE: C = C { re: 1.0, im: 0.0 };
    #[inline]
    fn add(self, o: C) -> C {
        C { re: self.re + o.re, im: self.im + o.im }
    }
    #[inline]
    fn sub(self, o: C) -> C {
        C { re: self.re - o.re, im: self.im - o.im }
    }
    #[inline]
    fn mul(self, o: C) -> C {
        C { re: self.re * o.re - self.im * o.im, im: self.re * o.im + self.im * o.re }
    }
    #[inline]
    fn conj(self) -> C {
        C { re: self.re, im: -self.im }
    }
    #[inline]
    fn scale(self, c: f64) -> C {
        C { re: self.re * c, im: self.im * c }
    }
    #[inline]
    fn norm2(self) -> f64 {
        self.re * self.re + self.im * self.im
    }
    #[inline]
    fn re(self) -> f64 {
        self.re
    }
    #[inline]
    fn is_zero(self) -> bool {
        self.re == 0.0 && self.im == 0.0
    }
    fn eigh(a: &[C], n: usize) -> (Vec<f64>, Vec<C>) {
        hermitian_eigen(a, n)
    }
}

// ---------------------------------------------------------------------------
// The symmetric and Hermitian eigensolvers
// ---------------------------------------------------------------------------

fn hypot(a: f64, b: f64) -> f64 {
    let (a, b) = (a.abs(), b.abs());
    let m = a.max(b);
    if m == 0.0 {
        return 0.0;
    }
    let (x, y) = (a / m, b / m);
    m * (x * x + y * y).sqrt()
}

/// Eigenvalues (ascending) and eigenvectors of a real symmetric `n × n`
/// matrix given row-major. Eigenvector `i` is column `i` of the returned
/// row-major matrix.
// The loops keep the index form of the published tred2/tql2 algorithms, so
// the code can be checked against them line by line.
#[allow(clippy::needless_range_loop, clippy::manual_memcpy)]
pub fn symmetric_eigen(a: &[f64], n: usize) -> (Vec<f64>, Vec<f64>) {
    if n == 0 {
        return (Vec::new(), Vec::new());
    }
    let mut v: Vec<Vec<f64>> = (0..n).map(|i| a[i * n..(i + 1) * n].to_vec()).collect();
    let mut d = vec![0.0; n];
    let mut e = vec![0.0; n];
    // Householder reduction to tridiagonal form.
    for j in 0..n {
        d[j] = v[n - 1][j];
    }
    for i in (1..n).rev() {
        let mut scale = 0.0;
        let mut h = 0.0;
        for k in 0..i {
            scale += d[k].abs();
        }
        if scale == 0.0 {
            e[i] = d[i - 1];
            for j in 0..i {
                d[j] = v[i - 1][j];
                v[i][j] = 0.0;
                v[j][i] = 0.0;
            }
        } else {
            for k in 0..i {
                d[k] /= scale;
                h += d[k] * d[k];
            }
            let mut f = d[i - 1];
            let mut g = h.sqrt();
            if f > 0.0 {
                g = -g;
            }
            e[i] = scale * g;
            h -= f * g;
            d[i - 1] = f - g;
            for item in e.iter_mut().take(i) {
                *item = 0.0;
            }
            for j in 0..i {
                f = d[j];
                v[j][i] = f;
                g = e[j] + v[j][j] * f;
                for k in j + 1..i {
                    g += v[k][j] * d[k];
                    e[k] += v[k][j] * f;
                }
                e[j] = g;
            }
            f = 0.0;
            for j in 0..i {
                e[j] /= h;
                f += e[j] * d[j];
            }
            let hh = f / (h + h);
            for j in 0..i {
                e[j] -= hh * d[j];
            }
            for j in 0..i {
                f = d[j];
                g = e[j];
                for k in j..i {
                    v[k][j] -= f * e[k] + g * d[k];
                }
                d[j] = v[i - 1][j];
                v[i][j] = 0.0;
            }
        }
        d[i] = h;
    }
    for i in 0..n - 1 {
        v[n - 1][i] = v[i][i];
        v[i][i] = 1.0;
        let h = d[i + 1];
        if h != 0.0 {
            for k in 0..=i {
                d[k] = v[k][i + 1] / h;
            }
            for j in 0..=i {
                let mut g = 0.0;
                for k in 0..=i {
                    g += v[k][i + 1] * v[k][j];
                }
                for k in 0..=i {
                    v[k][j] -= g * d[k];
                }
            }
        }
        for row in v.iter_mut().take(i + 1) {
            row[i + 1] = 0.0;
        }
    }
    for j in 0..n {
        d[j] = v[n - 1][j];
        v[n - 1][j] = 0.0;
    }
    v[n - 1][n - 1] = 1.0;
    e[0] = 0.0;
    tql2(&mut d, &mut e, &mut v);
    (d, v.into_iter().flatten().collect())
}

/// The implicit QL iteration on the tridiagonal matrix with diagonal `d` and
/// subdiagonal `e[i] = T[i][i − 1]` (`e[0]` unused), applied to the rows of
/// `v`. On return `d` holds the eigenvalues ascending and the columns of `v`
/// the eigenvectors, carried through `v`'s starting transformation.
#[allow(clippy::needless_range_loop)]
fn tql2(d: &mut [f64], e: &mut [f64], v: &mut [Vec<f64>]) {
    let n = d.len();
    for i in 1..n {
        e[i - 1] = e[i];
    }
    e[n - 1] = 0.0;
    let mut f = 0.0;
    let mut tst1: f64 = 0.0;
    let eps = f64::EPSILON;
    for l in 0..n {
        tst1 = tst1.max(d[l].abs() + e[l].abs());
        let mut m = l;
        while m < n {
            if e[m].abs() <= eps * tst1 {
                break;
            }
            m += 1;
        }
        if m > l {
            loop {
                let mut g = d[l];
                let mut p = (d[l + 1] - g) / (2.0 * e[l]);
                let mut r = hypot(p, 1.0);
                if p < 0.0 {
                    r = -r;
                }
                d[l] = e[l] / (p + r);
                d[l + 1] = e[l] * (p + r);
                let dl1 = d[l + 1];
                let mut h = g - d[l];
                for item in d.iter_mut().skip(l + 2) {
                    *item -= h;
                }
                f += h;
                p = d[m];
                let (mut c, mut c2, mut c3) = (1.0, 1.0, 1.0);
                let el1 = e[l + 1];
                let (mut s, mut s2) = (0.0, 0.0);
                for i in (l..m).rev() {
                    c3 = c2;
                    c2 = c;
                    s2 = s;
                    g = c * e[i];
                    h = c * p;
                    r = hypot(p, e[i]);
                    e[i + 1] = s * r;
                    s = e[i] / r;
                    c = p / r;
                    p = c * d[i] - s * g;
                    d[i + 1] = h + s * (c * g + s * d[i]);
                    for row in v.iter_mut() {
                        h = row[i + 1];
                        row[i + 1] = s * row[i] + c * h;
                        row[i] = c * row[i] - s * h;
                    }
                }
                p = -s * s2 * c3 * el1 * e[l] / dl1;
                e[l] = s * p;
                d[l] = c * p;
                if e[l].abs() <= eps * tst1 {
                    break;
                }
            }
        }
        d[l] += f;
        e[l] = 0.0;
    }
    // Sort ascending, carrying the vectors.
    for i in 0..n - 1 {
        let mut k = i;
        let mut p = d[i];
        for (j, &dj) in d.iter().enumerate().skip(i + 1) {
            if dj < p {
                k = j;
                p = dj;
            }
        }
        if k != i {
            d.swap(k, i);
            for row in v.iter_mut() {
                row.swap(i, k);
            }
        }
    }
}

/// Eigenvalues (ascending) and eigenvectors of a Hermitian `n × n` matrix
/// given row-major. Eigenvector `i` is column `i` of the returned row-major
/// matrix.
///
/// Householder reflections chosen to leave a real subdiagonal reduce the
/// matrix to a real symmetric tridiagonal one, `Q† A Q = T`. The QL iteration
/// diagonalises `T`, and the reflections carry its eigenvectors back.
#[allow(clippy::needless_range_loop)]
pub fn hermitian_eigen(a: &[C], n: usize) -> (Vec<f64>, Vec<C>) {
    if n == 0 {
        return (Vec::new(), Vec::new());
    }
    let mut a = a.to_vec();
    let mut d = vec![0.0; n];
    let mut e = vec![0.0; n];
    // (tau, v) for each reflection H = I − tau v v†, acting on rows i + 1..n.
    let mut reflections: Vec<(C, Vec<C>)> = Vec::with_capacity(n.saturating_sub(1));
    for i in 0..n - 1 {
        let m = n - i - 1;
        let alpha = a[(i + 1) * n + i];
        let xnorm2: f64 = (i + 2..n).map(|r| a[r * n + i].norm2()).sum();
        if xnorm2 == 0.0 && alpha.im == 0.0 {
            // Already real and tridiagonal in this column.
            e[i + 1] = alpha.re;
            reflections.push((C::ZERO, Vec::new()));
        } else {
            let norm = (alpha.re * alpha.re + alpha.im * alpha.im + xnorm2).sqrt();
            let beta = if alpha.re >= 0.0 { -norm } else { norm };
            let tau = C { re: (beta - alpha.re) / beta, im: -alpha.im / beta };
            // v = (1, x / (alpha − beta)).
            let z = C { re: alpha.re - beta, im: alpha.im };
            let zz = z.re * z.re + z.im * z.im;
            let inv = C { re: z.re / zz, im: -z.im / zz };
            let mut v = Vec::with_capacity(m);
            v.push(C::ONE);
            for r in i + 2..n {
                v.push(Scalar::mul(a[r * n + i], inv));
            }
            e[i + 1] = beta;
            // p = tau · B v over the trailing block B = A[i+1.., i+1..].
            let mut p = vec![C::ZERO; m];
            for r in 0..m {
                let row = &a[(i + 1 + r) * n + i + 1..(i + 1 + r) * n + n];
                let mut s = C::ZERO;
                for (x, y) in row.iter().zip(&v) {
                    s = Scalar::add(s, Scalar::mul(*x, *y));
                }
                p[r] = Scalar::mul(tau, s);
            }
            // w = p − ½ tau (p† v) v
            let mut pv = C::ZERO;
            for (x, y) in p.iter().zip(&v) {
                pv = Scalar::add(pv, Scalar::mul(Scalar::conj(*x), *y));
            }
            let k = Scalar::scale(Scalar::mul(tau, pv), -0.5);
            let w: Vec<C> = p.iter().zip(&v).map(|(x, y)| Scalar::add(*x, Scalar::mul(k, *y))).collect();
            // B −= v w† + w v†
            for r in 0..m {
                for c in 0..m {
                    let upd = Scalar::add(Scalar::mul(v[r], Scalar::conj(w[c])), Scalar::mul(w[r], Scalar::conj(v[c])));
                    let at = (i + 1 + r) * n + i + 1 + c;
                    a[at] = Scalar::sub(a[at], upd);
                }
            }
            reflections.push((tau, v));
        }
        d[i] = a[i * n + i].re;
    }
    d[n - 1] = a[(n - 1) * n + n - 1].re;
    let mut z: Vec<Vec<f64>> = (0..n).map(|r| (0..n).map(|c| if r == c { 1.0 } else { 0.0 }).collect()).collect();
    tql2(&mut d, &mut e, &mut z);
    // X = H_0 H_1 … H_{n−2} Z.
    let mut x: Vec<C> = z.into_iter().flatten().map(|re| C { re, im: 0.0 }).collect();
    for (i, (tau, v)) in reflections.iter().enumerate().rev() {
        if v.is_empty() {
            continue;
        }
        for c in 0..n {
            let mut s = C::ZERO;
            for (t, vt) in v.iter().enumerate() {
                s = Scalar::add(s, Scalar::mul(Scalar::conj(*vt), x[(i + 1 + t) * n + c]));
            }
            let ts = Scalar::mul(*tau, s);
            for (t, vt) in v.iter().enumerate() {
                let at = (i + 1 + t) * n + c;
                x[at] = Scalar::sub(x[at], Scalar::mul(*vt, ts));
            }
        }
    }
    (d, x)
}

// ---------------------------------------------------------------------------
// Hamiltonians
// ---------------------------------------------------------------------------

/// A real one-site operator, row-major in the basis `|↑⟩ = 0`, `|↓⟩ = 1`.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Op(pub [f64; 4]);

impl Op {
    pub fn id() -> Op {
        Op([1.0, 0.0, 0.0, 1.0])
    }
    pub fn x() -> Op {
        Op([0.0, 1.0, 1.0, 0.0])
    }
    pub fn z() -> Op {
        Op([1.0, 0.0, 0.0, -1.0])
    }
    /// `S⁺ = |↑⟩⟨↓|`.
    pub fn plus() -> Op {
        Op([0.0, 1.0, 0.0, 0.0])
    }
    /// `S⁻ = |↓⟩⟨↑|`.
    pub fn minus() -> Op {
        Op([0.0, 0.0, 1.0, 0.0])
    }
    /// `S^z = σ^z / 2`.
    pub fn sz() -> Op {
        Op([0.5, 0.0, 0.0, -0.5])
    }
    fn scaled(self, c: f64) -> Op {
        Op(self.0.map(|v| v * c))
    }
    fn add(self, o: Op) -> Op {
        Op([self.0[0] + o.0[0], self.0[1] + o.0[1], self.0[2] + o.0[2], self.0[3] + o.0[3]])
    }
}

/// A uniform nearest-neighbour chain: `H = Σ_i Σ_t c_t A_t(i) B_t(i+1) +
/// Σ_i Σ_f c_f O_f(i)`, open boundaries.
#[derive(Clone, Debug, PartialEq)]
pub struct Chain {
    pub n: usize,
    pub bonds: Vec<(f64, Op, Op)>,
    pub sites: Vec<(f64, Op)>,
}

impl Chain {
    /// The XXZ Heisenberg chain in a field:
    /// `H = Σ [J (S^x S^x + S^y S^y) + J_z S^z S^z] − h Σ S^z`.
    pub fn heisenberg(n: usize, j: f64, jz: f64, h: f64) -> Chain {
        let mut sites = Vec::new();
        if h != 0.0 {
            sites.push((-h, Op::sz()));
        }
        Chain { n, bonds: vec![(j / 2.0, Op::plus(), Op::minus()), (j / 2.0, Op::minus(), Op::plus()), (jz, Op::sz(), Op::sz())], sites }
    }

    /// The transverse-field Ising chain: `H = −J Σ σ^z σ^z − h Σ σ^x`.
    pub fn ising(n: usize, j: f64, h: f64) -> Chain {
        Chain { n, bonds: vec![(-j, Op::z(), Op::z())], sites: vec![(-h, Op::x())] }
    }

    /// The chain as a matrix-product operator. The finite-state machine is:
    /// state 0 is "nothing placed yet", state `t + 1` is "A_t placed, B_t
    /// owed", and the last state is "done".
    pub fn mpo(&self) -> Mpo {
        let m = self.bonds.len();
        let w = m + 2;
        let onsite = self.sites.iter().fold(Op([0.0; 4]), |acc, &(c, o)| acc.add(o.scaled(c)));
        let mut site = vec![Op([0.0; 4]); w * w];
        site[0] = Op::id();
        site[(w - 1) * w + (w - 1)] = Op::id();
        site[w - 1] = onsite;
        for (t, &(c, a, b)) in self.bonds.iter().enumerate() {
            site[t + 1] = a.scaled(c);
            site[(t + 1) * w + (w - 1)] = b;
        }
        Mpo { n: self.n, w, site }
    }
}

/// A uniform matrix-product operator: `site[a·w + b]` is the one-site
/// operator from state `a` to state `b`; the chain starts in state 0 and
/// ends in state `w − 1`.
#[derive(Clone, Debug, PartialEq)]
pub struct Mpo {
    pub n: usize,
    pub w: usize,
    site: Vec<Op>,
}

// ---------------------------------------------------------------------------
// Exact referees
// ---------------------------------------------------------------------------

/// Lanczos for the lowest eigenpair of a symmetric operator given by
/// `apply`, from `start`, with full re-orthogonalisation.
fn lanczos(apply: &dyn Fn(&[f64], &mut [f64]), start: &[f64], krylov: usize, restarts: usize, tol: f64) -> (f64, Vec<f64>) {
    let dim = start.len();
    let norm = |v: &[f64]| v.iter().map(|x| x * x).sum::<f64>().sqrt();
    let mut v = start.to_vec();
    let nv = norm(&v);
    if nv == 0.0 {
        v = (0..dim).map(|i| 1.0 + (i % 7) as f64 * 0.1).collect();
    }
    let nv = norm(&v);
    v.iter_mut().for_each(|x| *x /= nv);
    let mut energy = 0.0;
    let mut w = vec![0.0; dim];
    for _ in 0..restarts.max(1) {
        let k = krylov.min(dim).max(1);
        let mut basis: Vec<Vec<f64>> = vec![v.clone()];
        let (mut alpha, mut beta) = (Vec::new(), Vec::new());
        for j in 0..k {
            apply(&basis[j], &mut w);
            let a: f64 = w.iter().zip(&basis[j]).map(|(x, y)| x * y).sum();
            alpha.push(a);
            // Full re-orthogonalisation, twice.
            for _ in 0..2 {
                for b in &basis {
                    let c: f64 = w.iter().zip(b).map(|(x, y)| x * y).sum();
                    w.iter_mut().zip(b).for_each(|(x, y)| *x -= c * y);
                }
            }
            let bnorm = norm(&w);
            if j + 1 == k || bnorm < 1e-13 {
                break;
            }
            beta.push(bnorm);
            basis.push(w.iter().map(|x| x / bnorm).collect());
        }
        let m = alpha.len();
        let mut t = vec![0.0; m * m];
        for i in 0..m {
            t[i * m + i] = alpha[i];
            if i + 1 < m {
                t[i * m + i + 1] = beta[i];
                t[(i + 1) * m + i] = beta[i];
            }
        }
        let (vals, vecs) = symmetric_eigen(&t, m);
        energy = vals[0];
        let mut next = vec![0.0; dim];
        for (i, b) in basis.iter().enumerate().take(m) {
            let c = vecs[i * m];
            next.iter_mut().zip(b).for_each(|(x, y)| *x += c * y);
        }
        let nn = norm(&next);
        next.iter_mut().for_each(|x| *x /= nn);
        v = next;
        apply(&v, &mut w);
        let resid: f64 = w.iter().zip(&v).map(|(x, y)| (x - energy * y) * (x - energy * y)).sum::<f64>().sqrt();
        if resid < tol {
            break;
        }
    }
    (energy, v)
}

/// `out = H inp` for `chain` on all `2ⁿ` amplitudes, site 0 the most
/// significant bit.
pub(crate) fn apply_dense(chain: &Chain, inp: &[f64], out: &mut [f64]) {
    let n = chain.n;
    let bit = |x: usize, i: usize| (x >> (n - 1 - i)) & 1;
    out.iter_mut().for_each(|o| *o = 0.0);
    for (x, &amp) in inp.iter().enumerate() {
        if amp == 0.0 {
            continue;
        }
        for i in 0..n {
            let s = bit(x, i);
            for &(c, o) in &chain.sites {
                for sp in 0..2 {
                    let m = o.0[sp * 2 + s];
                    if m != 0.0 {
                        let y = (x & !(1 << (n - 1 - i))) | (sp << (n - 1 - i));
                        out[y] += c * m * amp;
                    }
                }
            }
            if i + 1 < n {
                let s2 = bit(x, i + 1);
                for &(c, a, b) in &chain.bonds {
                    for ap in 0..2 {
                        let ma = a.0[ap * 2 + s];
                        if ma == 0.0 {
                            continue;
                        }
                        for bp in 0..2 {
                            let mb = b.0[bp * 2 + s2];
                            if mb == 0.0 {
                                continue;
                            }
                            let mask = (1 << (n - 1 - i)) | (1 << (n - 2 - i));
                            let y = (x & !mask) | (ap << (n - 1 - i)) | (bp << (n - 2 - i));
                            out[y] += c * ma * mb * amp;
                        }
                    }
                }
            }
        }
    }
}

/// The exact ground energy of `chain` by Lanczos on all `2ⁿ` states (at most
/// 22 sites).
pub fn exact_ground_energy(chain: &Chain) -> Option<f64> {
    let n = chain.n;
    if n == 0 || n > 22 {
        return None;
    }
    let dim = 1usize << n;
    let apply = |inp: &[f64], out: &mut [f64]| apply_dense(chain, inp, out);
    // u64, not usize: the start vector must not depend on the pointer width.
    let start: Vec<f64> = (0..dim as u64).map(|x| 1.0 + (x.wrapping_mul(2_654_435_761) % 1000) as f64 / 1000.0).collect();
    Some(lanczos(&apply, &start, 60, 20, 1e-10).0)
}

/// The exact ground energy of the open transverse-field Ising chain
/// `−J Σ σ^z σ^z − h Σ σ^x` as free fermions. After a Jordan–Wigner
/// transformation, the chain is quadratic in fermions with
/// `A_ii = 2h`, `A_{i,i+1} = A_{i+1,i} = −J` and `B_{i,i+1} = −B_{i+1,i} = −J`.
/// The quasiparticle energies `Λ_k` are the square roots of the eigenvalues
/// of `(A − B)(A + B)`, and `E₀ = −½ Σ_k Λ_k`.
pub fn ising_exact_energy(n: usize, j: f64, h: f64) -> f64 {
    let mut a = vec![0.0; n * n];
    let mut b = vec![0.0; n * n];
    for i in 0..n {
        a[i * n + i] = 2.0 * h;
        if i + 1 < n {
            a[i * n + i + 1] = -j;
            a[(i + 1) * n + i] = -j;
            b[i * n + i + 1] = -j;
            b[(i + 1) * n + i] = j;
        }
    }
    // M = (A − B)(A + B), symmetric positive semidefinite.
    let mut m = vec![0.0; n * n];
    for r in 0..n {
        for c in 0..n {
            let mut s = 0.0;
            for k in 0..n {
                s += (a[r * n + k] - b[r * n + k]) * (a[k * n + c] + b[k * n + c]);
            }
            m[r * n + c] = s;
        }
    }
    // Symmetrise against rounding.
    for r in 0..n {
        for c in r + 1..n {
            let s = 0.5 * (m[r * n + c] + m[c * n + r]);
            m[r * n + c] = s;
            m[c * n + r] = s;
        }
    }
    let (vals, _) = symmetric_eigen(&m, n);
    -0.5 * vals.iter().map(|v| v.max(0.0).sqrt()).sum::<f64>()
}

// ---------------------------------------------------------------------------
// Matrix-product states and DMRG
// ---------------------------------------------------------------------------

/// A matrix-product state: site `i` is `left × 2 × right`, stored
/// `[(a·2 + s)·right + b]`. Real for ground states, complex ([`C`]) for time
/// evolution.
#[derive(Clone, Debug, PartialEq)]
pub struct Mps<T = f64> {
    pub dims: Vec<usize>,
    pub sites: Vec<Vec<T>>,
}

impl Mps<f64> {
    /// A deterministic pseudo-random state with bond dimension `bond` (capped
    /// by what the chain allows).
    pub fn random(n: usize, bond: usize, seed: u64) -> Mps {
        let mut s = seed;
        let mut next = move || {
            s = s.wrapping_add(0x9e37_79b9_7f4a_7c15);
            let mut z = s;
            z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
            z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
            ((z ^ (z >> 31)) >> 11) as f64 / 9_007_199_254_740_992.0 - 0.5
        };
        let mut dims = vec![1usize; n + 1];
        for (i, d) in dims.iter_mut().enumerate().take(n).skip(1) {
            let cap = 1usize.checked_shl(i.min(n - i).min(30) as u32).unwrap_or(usize::MAX);
            *d = bond.min(cap);
        }
        let sites = (0..n).map(|i| (0..dims[i] * 2 * dims[i + 1]).map(|_| next()).collect()).collect();
        Mps { dims, sites }
    }

    /// The same state with complex entries.
    pub fn to_complex(&self) -> Mps<C> {
        Mps { dims: self.dims.clone(), sites: self.sites.iter().map(|s| s.iter().map(|&re| C { re, im: 0.0 }).collect()).collect() }
    }
}

impl Mps<C> {
    /// The product state `⊗ (a_i |↑⟩ + b_i |↓⟩)`, each factor normalised.
    pub fn product(factors: &[[C; 2]]) -> Mps<C> {
        let sites = factors
            .iter()
            .map(|f| {
                let norm = (f[0].norm2() + f[1].norm2()).sqrt();
                assert!(norm > 0.0, "a product factor must be non-zero");
                vec![Scalar::scale(f[0], 1.0 / norm), Scalar::scale(f[1], 1.0 / norm)]
            })
            .collect();
        Mps { dims: vec![1; factors.len() + 1], sites }
    }

    /// The basis state with site `i` down where `down[i]`.
    pub fn basis(down: &[bool]) -> Mps<C> {
        let f: Vec<[C; 2]> = down.iter().map(|&d| if d { [C::ZERO, C::ONE] } else { [C::ONE, C::ZERO] }).collect();
        Mps::product(&f)
    }
}

impl<T: Scalar> Mps<T> {
    pub fn n(&self) -> usize {
        self.sites.len()
    }

    /// All `2ⁿ` amplitudes (at most 24 sites), site 0 the most significant
    /// bit and `|↑⟩ = 0`.
    pub fn to_dense(&self) -> Option<Vec<T>> {
        let n = self.n();
        if n == 0 || n > 24 {
            return None;
        }
        // Rows are the basis states of the sites so far; columns the bond.
        let mut cur = vec![T::ONE];
        let mut rows = 1usize;
        for i in 0..n {
            let (dl, dr) = (self.dims[i], self.dims[i + 1]);
            let site = &self.sites[i];
            let mut next = vec![T::ZERO; rows * 2 * dr];
            for r in 0..rows {
                for a in 0..dl {
                    let c = cur[r * dl + a];
                    if c.is_zero() {
                        continue;
                    }
                    for s in 0..2 {
                        for b in 0..dr {
                            let at = (r * 2 + s) * dr + b;
                            next[at] = next[at].add(c.mul(site[(a * 2 + s) * dr + b]));
                        }
                    }
                }
            }
            cur = next;
            rows *= 2;
        }
        Some(cur)
    }
}

/// Settings for [`dmrg`].
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct DmrgConfig {
    pub max_bond: usize,
    /// Discard density-matrix weight below this.
    pub cutoff: f64,
    /// Full sweeps (left to right and back).
    pub sweeps: usize,
    /// Lanczos vectors per restart, and restarts.
    pub krylov: usize,
    pub restarts: usize,
    /// Stop Lanczos once the residual `‖Hv − Ev‖` falls below this. The
    /// energy error goes as its square.
    pub tol: f64,
    pub seed: u64,
    /// Threads for the effective Hamiltonian; the result does not depend on
    /// it.
    pub threads: usize,
}

/// The machine's thread count; one in WebAssembly.
pub(crate) fn default_threads() -> usize {
    if cfg!(target_arch = "wasm32") { 1 } else { std::thread::available_parallelism().map_or(1, |n| n.get()) }
}

impl Default for DmrgConfig {
    fn default() -> DmrgConfig {
        DmrgConfig { max_bond: 64, cutoff: 1e-12, sweeps: 6, krylov: 24, restarts: 4, tol: 1e-8, seed: 1, threads: default_threads() }
    }
}

/// The outcome of a DMRG run.
#[derive(Clone, Debug, PartialEq)]
pub struct DmrgResult {
    pub energy: f64,
    /// The energy after each sweep.
    pub sweep_energies: Vec<f64>,
    /// The largest discarded weight in the last sweep.
    pub discarded: f64,
    /// Von Neumann entanglement entropy at each of the `n − 1` bonds (last
    /// sweep).
    pub entropies: Vec<f64>,
    pub mps: Mps,
}

/// Contract environments and effective Hamiltonians. Environments are
/// `L[a', w, a]` and `R[b', w, b]`, bra index first.
pub(crate) struct Engine<'a> {
    mpo: &'a Mpo,
    w: usize,
}

impl Engine<'_> {
    pub(crate) fn new(mpo: &Mpo) -> Engine<'_> {
        Engine { mpo, w: mpo.w }
    }

    fn op(&self, a: usize, b: usize) -> &Op {
        &self.mpo.site[a * self.w + b]
    }

    fn nonzero(&self, a: usize, b: usize) -> bool {
        self.op(a, b).0.iter().any(|&v| v != 0.0)
    }

    /// The environment at the left edge.
    pub(crate) fn left_edge<T: Scalar>(&self) -> Vec<T> {
        let mut l = vec![T::ZERO; self.w];
        l[0] = T::ONE;
        l
    }

    /// The environment at the right edge.
    pub(crate) fn right_edge<T: Scalar>(&self) -> Vec<T> {
        let mut r = vec![T::ZERO; self.w];
        r[self.w - 1] = T::ONE;
        r
    }

    /// `L'[b', w2, b] = Σ A*[a',s',b'] L[a',w1,a] W[w1,w2,s',s] A[a,s,b]`.
    pub(crate) fn grow_left<T: Scalar>(&self, l: &[T], a: &[T], dl: usize, dr: usize) -> Vec<T> {
        let w = self.w;
        // T[a', w1, s, b] = Σ_a L[a', w1, a] A[a, s, b]
        let mut t = vec![T::ZERO; dl * w * 2 * dr];
        for ap in 0..dl {
            for w1 in 0..w {
                let lrow = &l[(ap * w + w1) * dl..(ap * w + w1 + 1) * dl];
                for (aa, &lv) in lrow.iter().enumerate() {
                    if lv.is_zero() {
                        continue;
                    }
                    for s in 0..2 {
                        let src = &a[(aa * 2 + s) * dr..(aa * 2 + s + 1) * dr];
                        let dst = &mut t[((ap * w + w1) * 2 + s) * dr..((ap * w + w1) * 2 + s + 1) * dr];
                        for (d, &x) in dst.iter_mut().zip(src) {
                            *d = d.add(lv.mul(x));
                        }
                    }
                }
            }
        }
        // U[a', w2, s', b] = Σ_{w1, s} W[w1, w2, s', s] T[a', w1, s, b]
        let mut u = vec![T::ZERO; dl * w * 2 * dr];
        for ap in 0..dl {
            for w1 in 0..w {
                for w2 in 0..w {
                    if !self.nonzero(w1, w2) {
                        continue;
                    }
                    let o = self.op(w1, w2).0;
                    for sp in 0..2 {
                        for s in 0..2 {
                            let m = o[sp * 2 + s];
                            if m == 0.0 {
                                continue;
                            }
                            let src = ((ap * w + w1) * 2 + s) * dr;
                            let dst = ((ap * w + w2) * 2 + sp) * dr;
                            for b in 0..dr {
                                u[dst + b] = u[dst + b].add(t[src + b].scale(m));
                            }
                        }
                    }
                }
            }
        }
        // L'[b', w2, b] = Σ_{a', s'} A*[a', s', b'] U[a', w2, s', b]
        let mut out = vec![T::ZERO; dr * w * dr];
        for ap in 0..dl {
            for sp in 0..2 {
                let arow = &a[(ap * 2 + sp) * dr..(ap * 2 + sp + 1) * dr];
                for (bp, &av) in arow.iter().enumerate() {
                    if av.is_zero() {
                        continue;
                    }
                    let av = av.conj();
                    for w2 in 0..w {
                        let src = &u[((ap * w + w2) * 2 + sp) * dr..((ap * w + w2) * 2 + sp + 1) * dr];
                        let dst = &mut out[(bp * w + w2) * dr..(bp * w + w2 + 1) * dr];
                        for (d, &x) in dst.iter_mut().zip(src) {
                            *d = d.add(av.mul(x));
                        }
                    }
                }
            }
        }
        out
    }

    /// `R'[a', w1, a] = Σ B*[a',s',b'] W[w1,w2,s',s] R[b',w2,b] B[a,s,b]`.
    pub(crate) fn grow_right<T: Scalar>(&self, r: &[T], bt: &[T], dl: usize, dr: usize) -> Vec<T> {
        let w = self.w;
        // T[a, s, b', w2] = Σ_b B[a, s, b] R[b', w2, b]
        let mut t = vec![T::ZERO; dl * 2 * dr * w];
        for a in 0..dl {
            for s in 0..2 {
                let brow = &bt[(a * 2 + s) * dr..(a * 2 + s + 1) * dr];
                for bp in 0..dr {
                    for w2 in 0..w {
                        let rrow = &r[(bp * w + w2) * dr..(bp * w + w2 + 1) * dr];
                        let mut acc = T::ZERO;
                        for (x, y) in brow.iter().zip(rrow) {
                            acc = acc.add(x.mul(*y));
                        }
                        t[((a * 2 + s) * dr + bp) * w + w2] = acc;
                    }
                }
            }
        }
        // U[a, s', b', w1] = Σ_{w2, s} W[w1, w2, s', s] T[a, s, b', w2]
        let mut u = vec![T::ZERO; dl * 2 * dr * w];
        for w1 in 0..w {
            for w2 in 0..w {
                if !self.nonzero(w1, w2) {
                    continue;
                }
                let o = self.op(w1, w2).0;
                for sp in 0..2 {
                    for s in 0..2 {
                        let m = o[sp * 2 + s];
                        if m == 0.0 {
                            continue;
                        }
                        for a in 0..dl {
                            for bp in 0..dr {
                                let at = ((a * 2 + sp) * dr + bp) * w + w1;
                                u[at] = u[at].add(t[((a * 2 + s) * dr + bp) * w + w2].scale(m));
                            }
                        }
                    }
                }
            }
        }
        // R'[a', w1, a] = Σ_{s', b'} B*[a', s', b'] U[a, s', b', w1]
        let mut out = vec![T::ZERO; dl * w * dl];
        for ap in 0..dl {
            for sp in 0..2 {
                for bp in 0..dr {
                    let bv = bt[(ap * 2 + sp) * dr + bp];
                    if bv.is_zero() {
                        continue;
                    }
                    let bv = bv.conj();
                    for a in 0..dl {
                        for w1 in 0..w {
                            let at = (ap * w + w1) * dl + a;
                            out[at] = out[at].add(bv.mul(u[((a * 2 + sp) * dr + bp) * w + w1]));
                        }
                    }
                }
            }
        }
        out
    }

    /// Rows `rows` (bra indices `a'`) of `H_eff θ` for the two-site tensor
    /// `θ[a, s1, s2, b]`, written to `out` (which holds just those rows).
    #[allow(clippy::too_many_arguments)]
    fn apply_rows<T: Scalar>(&self, l: &[T], r: &[T], theta: &[T], out: &mut [T], dl: usize, dr: usize, rows: core::ops::Range<usize>) {
        let w = self.w;
        let first = rows.start;
        let nrows = rows.len();
        // T1[a', w0, s1, s2, b] = Σ_a L[a', w0, a] θ[a, s1, s2, b]
        let blk = 4 * dr;
        let mut t1 = vec![T::ZERO; nrows * w * blk];
        for ap in 0..nrows {
            for w0 in 0..w {
                let lrow = &l[((first + ap) * w + w0) * dl..((first + ap) * w + w0 + 1) * dl];
                let dst = &mut t1[(ap * w + w0) * blk..(ap * w + w0 + 1) * blk];
                for (a, &lv) in lrow.iter().enumerate() {
                    if lv.is_zero() {
                        continue;
                    }
                    for (d, &x) in dst.iter_mut().zip(&theta[a * blk..(a + 1) * blk]) {
                        *d = d.add(lv.mul(x));
                    }
                }
            }
        }
        // T2[a', w1, s1', s2, b] = Σ_{w0, s1} W[w0, w1, s1', s1] T1[a', w0, s1, s2, b]
        let half = 2 * dr;
        let mut t2 = vec![T::ZERO; nrows * w * blk];
        for ap in 0..nrows {
            for w0 in 0..w {
                for w1 in 0..w {
                    if !self.nonzero(w0, w1) {
                        continue;
                    }
                    let o = self.op(w0, w1).0;
                    for s1p in 0..2 {
                        for s1 in 0..2 {
                            let m = o[s1p * 2 + s1];
                            if m == 0.0 {
                                continue;
                            }
                            let src = (ap * w + w0) * blk + s1 * half;
                            let dst = (ap * w + w1) * blk + s1p * half;
                            for k in 0..half {
                                t2[dst + k] = t2[dst + k].add(t1[src + k].scale(m));
                            }
                        }
                    }
                }
            }
        }
        // T3[a', s1', w2, s2', b] = Σ_{w1, s2} W[w1, w2, s2', s2] T2[a', w1, s1', s2, b]
        let mut t3 = vec![T::ZERO; nrows * 2 * w * 2 * dr];
        for ap in 0..nrows {
            for w1 in 0..w {
                for w2 in 0..w {
                    if !self.nonzero(w1, w2) {
                        continue;
                    }
                    let o = self.op(w1, w2).0;
                    for s2p in 0..2 {
                        for s2 in 0..2 {
                            let m = o[s2p * 2 + s2];
                            if m == 0.0 {
                                continue;
                            }
                            for s1p in 0..2 {
                                let src = (ap * w + w1) * blk + s1p * half + s2 * dr;
                                let dst = (((ap * 2 + s1p) * w + w2) * 2 + s2p) * dr;
                                for b in 0..dr {
                                    t3[dst + b] = t3[dst + b].add(t2[src + b].scale(m));
                                }
                            }
                        }
                    }
                }
            }
        }
        // out[a', s1', s2', b'] = Σ_{w2, b} T3[a', s1', w2, s2', b] R[b', w2, b]
        for ap in 0..nrows {
            for s1p in 0..2 {
                for s2p in 0..2 {
                    for bp in 0..dr {
                        let mut acc = T::ZERO;
                        for w2 in 0..w {
                            let trow = &t3[(((ap * 2 + s1p) * w + w2) * 2 + s2p) * dr..(((ap * 2 + s1p) * w + w2) * 2 + s2p + 1) * dr];
                            let rrow = &r[(bp * w + w2) * dr..(bp * w + w2 + 1) * dr];
                            for (x, y) in trow.iter().zip(rrow) {
                                acc = acc.add(x.mul(*y));
                            }
                        }
                        out[((ap * 2 + s1p) * 2 + s2p) * dr + bp] = acc;
                    }
                }
            }
        }
    }

    /// Rows `rows` of `H_eff C` for the one-site tensor `C[a, s, b]`.
    #[cfg(feature = "quantum_tdvp")]
    #[allow(clippy::too_many_arguments)]
    fn apply1_rows<T: Scalar>(&self, l: &[T], r: &[T], c: &[T], out: &mut [T], dl: usize, dr: usize, rows: core::ops::Range<usize>) {
        let w = self.w;
        let first = rows.start;
        let nrows = rows.len();
        // T1[a', w0, s, b] = Σ_a L[a', w0, a] C[a, s, b]
        let blk = 2 * dr;
        let mut t1 = vec![T::ZERO; nrows * w * blk];
        for ap in 0..nrows {
            for w0 in 0..w {
                let lrow = &l[((first + ap) * w + w0) * dl..((first + ap) * w + w0 + 1) * dl];
                let dst = &mut t1[(ap * w + w0) * blk..(ap * w + w0 + 1) * blk];
                for (a, &lv) in lrow.iter().enumerate() {
                    if lv.is_zero() {
                        continue;
                    }
                    for (d, &x) in dst.iter_mut().zip(&c[a * blk..(a + 1) * blk]) {
                        *d = d.add(lv.mul(x));
                    }
                }
            }
        }
        // T2[a', s', w1, b] = Σ_{w0, s} W[w0, w1, s', s] T1[a', w0, s, b]
        let mut t2 = vec![T::ZERO; nrows * 2 * w * dr];
        for ap in 0..nrows {
            for w0 in 0..w {
                for w1 in 0..w {
                    if !self.nonzero(w0, w1) {
                        continue;
                    }
                    let o = self.op(w0, w1).0;
                    for sp in 0..2 {
                        for s in 0..2 {
                            let m = o[sp * 2 + s];
                            if m == 0.0 {
                                continue;
                            }
                            let src = (ap * w + w0) * blk + s * dr;
                            let dst = ((ap * 2 + sp) * w + w1) * dr;
                            for b in 0..dr {
                                t2[dst + b] = t2[dst + b].add(t1[src + b].scale(m));
                            }
                        }
                    }
                }
            }
        }
        // out[a', s', b'] = Σ_{w1, b} T2[a', s', w1, b] R[b', w1, b]
        for ap in 0..nrows {
            for sp in 0..2 {
                for bp in 0..dr {
                    let mut acc = T::ZERO;
                    for w1 in 0..w {
                        let trow = &t2[((ap * 2 + sp) * w + w1) * dr..((ap * 2 + sp) * w + w1 + 1) * dr];
                        let rrow = &r[(bp * w + w1) * dr..(bp * w + w1 + 1) * dr];
                        for (x, y) in trow.iter().zip(rrow) {
                            acc = acc.add(x.mul(*y));
                        }
                    }
                    out[(ap * 2 + sp) * dr + bp] = acc;
                }
            }
        }
    }

    /// `out = H_eff θ` for the two-site `θ[a, s1, s2, b]` between `l` (dl)
    /// and `r` (dr). Blocks of bra rows run on separate threads; each row is
    /// computed in the same order whatever the split, so the result is the
    /// same bits.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn apply<T: Scalar>(&self, l: &[T], r: &[T], theta: &[T], out: &mut [T], dl: usize, dr: usize, threads: usize) {
        let threads = Self::split_threads(threads, dl, dr);
        if threads <= 1 {
            self.apply_rows(l, r, theta, out, dl, dr, 0..dl);
            return;
        }
        let row = 4 * dr;
        let per = dl.div_ceil(threads);
        std::thread::scope(|scope| {
            for (c, chunk) in out.chunks_mut(per * row).enumerate() {
                let start = c * per;
                let end = (start + per).min(dl);
                scope.spawn(move || self.apply_rows(l, r, theta, chunk, dl, dr, start..end));
            }
        });
    }

    /// `out = H_eff C` for the one-site `C[a, s, b]`, threaded like
    /// [`Engine::apply`].
    #[cfg(feature = "quantum_tdvp")]
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn apply1<T: Scalar>(&self, l: &[T], r: &[T], c: &[T], out: &mut [T], dl: usize, dr: usize, threads: usize) {
        let threads = Self::split_threads(threads, dl, dr);
        if threads <= 1 {
            self.apply1_rows(l, r, c, out, dl, dr, 0..dl);
            return;
        }
        let row = 2 * dr;
        let per = dl.div_ceil(threads);
        std::thread::scope(|scope| {
            for (k, chunk) in out.chunks_mut(per * row).enumerate() {
                let start = k * per;
                let end = (start + per).min(dl);
                scope.spawn(move || self.apply1_rows(l, r, c, chunk, dl, dr, start..end));
            }
        });
    }

    /// At least eight rows a thread; small bonds stay on one thread, where
    /// spawning would cost more than it saves.
    fn split_threads(threads: usize, dl: usize, dr: usize) -> usize {
        let threads = if cfg!(target_arch = "wasm32") { 1 } else { threads.max(1).min(dl / 8) };
        if dl * dr < 2048 { 1 } else { threads }
    }
}

/// Split `m` (rows × cols) by the eigenvectors of `m m†`, keeping at most
/// `max_bond` (and never more than `m` has rank for) and dropping weight
/// below `cutoff`. Returns the kept vectors (rows × k), `u† m` (k × cols),
/// the discarded weight and the kept normalised weights.
pub(crate) fn split<T: Scalar>(m: &[T], rows: usize, cols: usize, max_bond: usize, cutoff: f64) -> (Vec<T>, Vec<T>, f64, Vec<f64>) {
    let mut rho = vec![T::ZERO; rows * rows];
    for i in 0..rows {
        for j in i..rows {
            let mut s = T::ZERO;
            for c in 0..cols {
                s = s.add(m[i * cols + c].mul(m[j * cols + c].conj()));
            }
            rho[i * rows + j] = s;
            rho[j * rows + i] = s.conj();
        }
    }
    let (vals, vecs) = T::eigh(&rho, rows);
    let total: f64 = vals.iter().map(|v| v.max(0.0)).sum();
    // Largest first.
    let order: Vec<usize> = (0..rows).rev().collect();
    let mut k = 0;
    let mut kept_weight = 0.0;
    for &idx in &order {
        let w = vals[idx].max(0.0) / total;
        if k >= max_bond.min(cols) || (k > 0 && w < cutoff) {
            break;
        }
        kept_weight += w;
        k += 1;
    }
    let k = k.max(1);
    let mut u = vec![T::ZERO; rows * k];
    let mut weights = Vec::with_capacity(k);
    for (c, &idx) in order.iter().take(k).enumerate() {
        weights.push(vals[idx].max(0.0) / total);
        for r in 0..rows {
            u[r * k + c] = vecs[r * rows + idx];
        }
    }
    let mut rest = vec![T::ZERO; k * cols];
    for c in 0..k {
        for r in 0..rows {
            let uv = u[r * k + c];
            if uv.is_zero() {
                continue;
            }
            let uv = uv.conj();
            for j in 0..cols {
                rest[c * cols + j] = rest[c * cols + j].add(uv.mul(m[r * cols + j]));
            }
        }
    }
    (u, rest, (1.0 - kept_weight).max(0.0), weights)
}

pub(crate) fn transpose<T: Scalar>(m: &[T], rows: usize, cols: usize) -> Vec<T> {
    let mut t = vec![T::ZERO; rows * cols];
    for r in 0..rows {
        for c in 0..cols {
            t[c * rows + r] = m[r * cols + c];
        }
    }
    t
}

pub(crate) fn entropy(weights: &[f64]) -> f64 {
    -weights.iter().filter(|&&p| p > 0.0).map(|&p| p * ln(p)).sum::<f64>()
}

/// Bring `mps` into right-canonical form: every site but the first has
/// orthonormal rows, and the first carries the whole state.
pub(crate) fn right_canonicalise<T: Scalar>(mps: &mut Mps<T>) {
    let n = mps.n();
    for i in (1..n).rev() {
        let (dl, dr) = (mps.dims[i], mps.dims[i + 1]);
        // Site as dl × (2·dr); orthonormalise its rows through the
        // eigenvectors of its transpose.
        let m = mps.sites[i].clone();
        let (u, rest, _, _) = split(&transpose(&m, dl, 2 * dr), 2 * dr, dl, dl, 0.0);
        let k = rest.len() / dl;
        let new_site = transpose(&u, 2 * dr, k);
        let carry = transpose(&rest, k, dl); // dl × k
        // The previous site absorbs `carry` on its right index.
        let (pl, pr) = (mps.dims[i - 1], mps.dims[i]);
        let prev = &mps.sites[i - 1];
        let mut np = vec![T::ZERO; pl * 2 * k];
        for a in 0..pl * 2 {
            for c in 0..k {
                let mut s = T::ZERO;
                for b in 0..pr {
                    s = s.add(prev[a * pr + b].mul(carry[b * k + c]));
                }
                np[a * k + c] = s;
            }
        }
        mps.sites[i - 1] = np;
        mps.sites[i] = new_site;
        mps.dims[i] = k;
    }
}

/// Two-site DMRG for the ground state of `chain`.
pub fn dmrg(chain: &Chain, cfg: &DmrgConfig) -> DmrgResult {
    let n = chain.n;
    assert!(n >= 2, "DMRG needs at least two sites");
    let mpo = chain.mpo();
    let eng = Engine::new(&mpo);
    let mut mps = Mps::random(n, cfg.max_bond.min(8), cfg.seed);
    // Bring the state into right-canonical form from the right end, and
    // build the right environments.
    right_canonicalise(&mut mps);
    let mut rights: Vec<Vec<f64>> = vec![Vec::new(); n + 1];
    rights[n] = eng.right_edge();
    for i in (1..n).rev() {
        rights[i] = eng.grow_right(&rights[i + 1], &mps.sites[i], mps.dims[i], mps.dims[i + 1]);
    }
    let mut lefts: Vec<Vec<f64>> = vec![Vec::new(); n + 1];
    lefts[0] = eng.left_edge();
    let mut energy = 0.0;
    let mut sweep_energies = Vec::new();
    let mut entropies = vec![0.0; n - 1];
    let mut discarded = 0.0;
    for _sweep in 0..cfg.sweeps {
        discarded = 0.0;
        // Left to right: optimise (i, i+1), keep the left part.
        for i in 0..n - 1 {
            let (dl, dm, dr) = (mps.dims[i], mps.dims[i + 1], mps.dims[i + 2]);
            let theta = two_site(&mps.sites[i], &mps.sites[i + 1], dl, dm, dr);
            let (l, r) = (&lefts[i], &rights[i + 2]);
            let apply = |x: &[f64], y: &mut [f64]| eng.apply(l, r, x, y, dl, dr, cfg.threads);
            let (e, v) = lanczos(&apply, &theta, cfg.krylov, cfg.restarts, cfg.tol);
            energy = e;
            let (u, rest, disc, weights) = split(&v, dl * 2, 2 * dr, cfg.max_bond, cfg.cutoff);
            let k = weights.len();
            discarded = f64::max(discarded, disc);
            entropies[i] = entropy(&weights);
            mps.sites[i] = u;
            mps.sites[i + 1] = rest;
            mps.dims[i + 1] = k;
            lefts[i + 1] = eng.grow_left(&lefts[i], &mps.sites[i], dl, k);
        }
        // Right to left: optimise (i, i+1), keep the right part.
        for i in (0..n - 1).rev() {
            let (dl, dm, dr) = (mps.dims[i], mps.dims[i + 1], mps.dims[i + 2]);
            let theta = two_site(&mps.sites[i], &mps.sites[i + 1], dl, dm, dr);
            let (l, r) = (&lefts[i], &rights[i + 2]);
            let apply = |x: &[f64], y: &mut [f64]| eng.apply(l, r, x, y, dl, dr, cfg.threads);
            let (e, v) = lanczos(&apply, &theta, cfg.krylov, cfg.restarts, cfg.tol);
            energy = e;
            // Split on the right: rows of the right factor are eigenvectors
            // of θᵀθ.
            let vt = transpose(&v, dl * 2, 2 * dr);
            let (u, rest, disc, weights) = split(&vt, 2 * dr, dl * 2, cfg.max_bond, cfg.cutoff);
            let k = weights.len();
            discarded = f64::max(discarded, disc);
            entropies[i] = entropy(&weights);
            mps.sites[i + 1] = transpose(&u, 2 * dr, k);
            mps.sites[i] = transpose(&rest, k, dl * 2);
            mps.dims[i + 1] = k;
            rights[i + 1] = eng.grow_right(&rights[i + 2], &mps.sites[i + 1], k, dr);
        }
        sweep_energies.push(energy);
    }
    DmrgResult { energy, sweep_energies, discarded, entropies, mps }
}

/// `θ[a, s1, s2, b] = Σ_c A[a, s1, c] B[c, s2, b]`.
pub(crate) fn two_site<T: Scalar>(a: &[T], b: &[T], dl: usize, dm: usize, dr: usize) -> Vec<T> {
    let mut t = vec![T::ZERO; dl * 4 * dr];
    for x in 0..dl * 2 {
        for c in 0..dm {
            let av = a[x * dm + c];
            if av.is_zero() {
                continue;
            }
            for s2 in 0..2 {
                let src = &b[(c * 2 + s2) * dr..(c * 2 + s2 + 1) * dr];
                let dst = &mut t[(x * 2 + s2) * dr..(x * 2 + s2 + 1) * dr];
                for (d, &v) in dst.iter_mut().zip(src) {
                    *d = d.add(av.mul(v));
                }
            }
        }
    }
    t
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn the_eigensolver_reconstructs_its_matrix() {
        let mut s = 7u64;
        for n in [1usize, 2, 5, 17, 40] {
            let mut a = vec![0.0; n * n];
            for i in 0..n {
                for j in i..n {
                    s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
                    let v = ((s >> 11) as f64 / 9_007_199_254_740_992.0) - 0.5;
                    a[i * n + j] = v;
                    a[j * n + i] = v;
                }
            }
            let (vals, vecs) = symmetric_eigen(&a, n);
            for w in vals.windows(2) {
                assert!(w[0] <= w[1]);
            }
            for r in 0..n {
                for c in 0..n {
                    let rec: f64 = (0..n).map(|k| vecs[r * n + k] * vals[k] * vecs[c * n + k]).sum();
                    assert!((rec - a[r * n + c]).abs() < 1e-12, "n={n}");
                    let orth: f64 = (0..n).map(|k| vecs[k * n + r] * vecs[k * n + c]).sum();
                    assert!((orth - if r == c { 1.0 } else { 0.0 }).abs() < 1e-12, "n={n}");
                }
            }
        }
    }

    #[test]
    fn the_hermitian_eigensolver_reconstructs_its_matrix() {
        let mut s = 11u64;
        let mut next = move || {
            s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
            ((s >> 11) as f64 / 9_007_199_254_740_992.0) - 0.5
        };
        for n in [1usize, 2, 3, 6, 17, 40] {
            let mut a = vec![C::ZERO; n * n];
            for i in 0..n {
                a[i * n + i] = C { re: next(), im: 0.0 };
                for j in i + 1..n {
                    let v = C { re: next(), im: next() };
                    a[i * n + j] = v;
                    a[j * n + i] = Scalar::conj(v);
                }
            }
            // A column already tridiagonal and real must pass through too.
            if n == 6 {
                for r in 2..n {
                    a[r * n] = C::ZERO;
                    a[r] = C::ZERO;
                }
                a[n] = C { re: 0.25, im: 0.0 };
                a[1] = C { re: 0.25, im: 0.0 };
            }
            let (vals, vecs) = hermitian_eigen(&a, n);
            for w in vals.windows(2) {
                assert!(w[0] <= w[1]);
            }
            for r in 0..n {
                for c in 0..n {
                    let mut rec = C::ZERO;
                    let mut orth = C::ZERO;
                    for k in 0..n {
                        rec = Scalar::add(rec, Scalar::scale(Scalar::mul(vecs[r * n + k], Scalar::conj(vecs[c * n + k])), vals[k]));
                        orth = Scalar::add(orth, Scalar::mul(Scalar::conj(vecs[k * n + r]), vecs[k * n + c]));
                    }
                    let d = Scalar::sub(rec, a[r * n + c]);
                    assert!(d.norm2().sqrt() < 1e-12, "n={n} reconstruct ({r},{c})");
                    let id = if r == c { 1.0 } else { 0.0 };
                    assert!((orth.re - id).abs() < 1e-12 && orth.im.abs() < 1e-12, "n={n} orthonormal ({r},{c})");
                }
            }
        }
        // On a real symmetric matrix it agrees with the real solver.
        let n = 9;
        let mut a = vec![0.0; n * n];
        for i in 0..n {
            for j in i..n {
                let v = next();
                a[i * n + j] = v;
                a[j * n + i] = v;
            }
        }
        let (real, _) = symmetric_eigen(&a, n);
        let (herm, _) = hermitian_eigen(&a.iter().map(|&re| C { re, im: 0.0 }).collect::<Vec<_>>(), n);
        for (x, y) in real.iter().zip(&herm) {
            assert!((x - y).abs() < 1e-13);
        }
    }

    #[test]
    fn free_fermions_solve_the_ising_chain() {
        for n in [4usize, 7, 10] {
            for (j, h) in [(1.0, 0.3), (1.0, 1.0), (0.7, 1.8)] {
                let exact = exact_ground_energy(&Chain::ising(n, j, h)).unwrap();
                let ff = ising_exact_energy(n, j, h);
                assert!((exact - ff).abs() < 1e-9, "n={n} J={j} h={h}: {exact} vs {ff}");
            }
        }
    }

    #[test]
    fn dmrg_matches_exact_diagonalisation() {
        let cfg = DmrgConfig { max_bond: 32, sweeps: 4, ..DmrgConfig::default() };
        for chain in [Chain::heisenberg(12, 1.0, 1.0, 0.0), Chain::heisenberg(10, 1.0, 0.5, 0.3), Chain::ising(12, 1.0, 1.0), Chain::ising(11, 1.0, 0.6)] {
            let exact = exact_ground_energy(&chain).unwrap();
            let r = dmrg(&chain, &cfg);
            assert!((r.energy - exact).abs() < 1e-8, "{chain:?}: {} vs {exact}", r.energy);
        }
    }

    #[test]
    fn dmrg_reaches_the_critical_ising_chain_at_forty_sites() {
        let n = 40;
        let exact = ising_exact_energy(n, 1.0, 1.0);
        let r = dmrg(&Chain::ising(n, 1.0, 1.0), &DmrgConfig { max_bond: 24, sweeps: 3, ..DmrgConfig::default() });
        assert!(r.energy >= exact - 1e-9, "variational: {} vs {exact}", r.energy);
        assert!((r.energy - exact).abs() / (n as f64) < 1e-8, "{} vs {exact}", r.energy);
        // Entanglement peaks mid-chain on an open critical chain.
        let mid = r.entropies[n / 2 - 1];
        assert!(mid > r.entropies[2] && mid > r.entropies[n - 4]);
    }

    #[test]
    fn threaded_effective_hamiltonian_is_bit_identical() {
        // Environments big enough to split across threads.
        let mpo = Chain::heisenberg(8, 1.0, 0.7, 0.2).mpo();
        let eng = Engine { mpo: &mpo, w: mpo.w };
        let (dl, dr, w) = (56usize, 48usize, mpo.w);
        let mut s = 3u64;
        let mut next = move || {
            s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
            ((s >> 11) as f64 / 9_007_199_254_740_992.0) - 0.5
        };
        let l: Vec<f64> = (0..dl * w * dl).map(|_| next()).collect();
        let r: Vec<f64> = (0..dr * w * dr).map(|_| next()).collect();
        let theta: Vec<f64> = (0..dl * 4 * dr).map(|_| next()).collect();
        let mut one = vec![0.0; theta.len()];
        eng.apply(&l, &r, &theta, &mut one, dl, dr, 1);
        for t in [2, 3, 7] {
            let mut many = vec![0.0; theta.len()];
            eng.apply(&l, &r, &theta, &mut many, dl, dr, t);
            assert!(one.iter().zip(&many).all(|(a, b)| a.to_bits() == b.to_bits()), "threads={t}");
        }
    }

    #[test]
    fn energy_falls_sweep_by_sweep_and_runs_repeat() {
        let chain = Chain::heisenberg(16, 1.0, 1.0, 0.0);
        let cfg = DmrgConfig { max_bond: 16, sweeps: 3, ..DmrgConfig::default() };
        let a = dmrg(&chain, &cfg);
        for w in a.sweep_energies.windows(2) {
            assert!(w[1] <= w[0] + 1e-10, "{:?}", a.sweep_energies);
        }
        assert_eq!(a, dmrg(&chain, &cfg));
        // Threads split the effective Hamiltonian by rows; no bit moves.
        let one = dmrg(&chain, &DmrgConfig { threads: 1, ..cfg });
        let many = dmrg(&chain, &DmrgConfig { threads: 5, ..cfg });
        assert_eq!(one, many);
    }
}