# Restricted Hartree-Fock for a linear chain of N hydrogen atoms in STO-3G.
#
# A scaled-down but genuine quantum-chemistry workload: the same machinery a
# package like Psi4 runs, restricted to s-type Gaussians. The cost is dominated
# by the O(N^4) two-electron integral build, which is parallelized with
# `parallel for` (OpenMP on the native build).
#
# Change `n_atoms` (keep it even for closed-shell RHF) to scale the work.

pi = 3.141592653589793

# STO-3G hydrogen: one contracted s-shell, three primitives.
alpha_h = [3.42525091, 0.62391373, 0.16885540]
d_h     = [0.15432897, 0.53532814, 0.44463454]
n_prim = 3

# ---- Molecule: N hydrogens evenly spaced along z. ----
n_atoms = 8
R = 1.4
n_basis = n_atoms
n_occ = n_atoms / 2

centers = zeros(n_atoms, 3)
for i to n_atoms {
    centers[i, 2] = i * R
}
Z = zeros(n_atoms)
for i to n_atoms {
    Z[i] = 1.0
}
basis_atom = zeros(n_basis)
for i to n_basis {
    basis_atom[i] = i
}

# ---- Boys function F_0(T) by direct Simpson integration. ----
def boys_F0(T) {
    if T < 1.0e-8 {
        return 1.0 - T / 3.0 + T * T / 10.0
    }
    return simpson(fn(t) -> exp(-T * t * t), 0.0, 1.0, 40)
}

def dist_sq(A, B) {
    d0 = A[0] - B[0]
    d1 = A[1] - B[1]
    d2 = A[2] - B[2]
    return d0 * d0 + d1 * d1 + d2 * d2
}

def gauss_product_center(a, A, b, B) {
    p = a + b
    P = zeros(3)
    P[0] = (a * A[0] + b * B[0]) / p
    P[1] = (a * A[1] + b * B[1]) / p
    P[2] = (a * A[2] + b * B[2]) / p
    return P
}

def prim_overlap(a, A, b, B) {
    p = a + b
    r2 = dist_sq(A, B)
    K = exp(-a * b / p * r2)
    Na = pow(2.0 * a / pi, 0.75)
    Nb = pow(2.0 * b / pi, 0.75)
    return Na * Nb * pow(pi / p, 1.5) * K
}

def prim_kinetic(a, A, b, B) {
    p = a + b
    r2 = dist_sq(A, B)
    S = prim_overlap(a, A, b, B)
    return a * b / p * (3.0 - 2.0 * a * b / p * r2) * S
}

def prim_nuclear(a, A, b, B, C, Zc) {
    p = a + b
    Na = pow(2.0 * a / pi, 0.75)
    Nb = pow(2.0 * b / pi, 0.75)
    P = gauss_product_center(a, A, b, B)
    rAB = dist_sq(A, B)
    rPC = dist_sq(P, C)
    K = exp(-a * b / p * rAB)
    return -2.0 * pi / p * Zc * Na * Nb * K * boys_F0(p * rPC)
}

def prim_eri(a, A, b, B, c, C, d, D) {
    p = a + b
    q = c + d
    Na = pow(2.0 * a / pi, 0.75)
    Nb = pow(2.0 * b / pi, 0.75)
    Nc = pow(2.0 * c / pi, 0.75)
    Nd = pow(2.0 * d / pi, 0.75)
    rAB = dist_sq(A, B)
    rCD = dist_sq(C, D)
    P = gauss_product_center(a, A, b, B)
    Q = gauss_product_center(c, C, d, D)
    rPQ = dist_sq(P, Q)
    K1 = exp(-a * b / p * rAB)
    K2 = exp(-c * d / q * rCD)
    pref = 2.0 * pow(pi, 2.5) / (p * q * sqrt(p + q))
    return Na * Nb * Nc * Nd * pref * K1 * K2 * boys_F0(p * q / (p + q) * rPQ)
}

def center_of(atom_idx) {
    R = zeros(3)
    R[0] = centers[atom_idx, 0]
    R[1] = centers[atom_idx, 1]
    R[2] = centers[atom_idx, 2]
    return R
}

def contracted_one_elec(mu, nu, kind, C, Zc) {
    A = center_of(basis_atom[mu])
    B = center_of(basis_atom[nu])
    total = 0.0
    for p to n_prim {
        for q to n_prim {
            a = alpha_h[p]
            b = alpha_h[q]
            dp = d_h[p]
            dq = d_h[q]
            v = 0.0
            if kind == 0 { v = prim_overlap(a, A, b, B) }
            if kind == 1 { v = prim_kinetic(a, A, b, B) }
            if kind == 2 { v = prim_nuclear(a, A, b, B, C, Zc) }
            total += dp * dq * v
        }
    }
    return total
}

def contracted_eri(mu, nu, lam, sig) {
    A = center_of(basis_atom[mu])
    B = center_of(basis_atom[nu])
    C = center_of(basis_atom[lam])
    D = center_of(basis_atom[sig])
    total = 0.0
    for p to n_prim {
        for q to n_prim {
            for r to n_prim {
                for s to n_prim {
                    a = alpha_h[p]
                    b = alpha_h[q]
                    c = alpha_h[r]
                    dd = alpha_h[s]
                    coef4 = d_h[p] * d_h[q] * d_h[r] * d_h[s]
                    total += coef4 * prim_eri(a, A, b, B, c, C, dd, D)
                }
            }
        }
    }
    return total
}

# ---- One-electron matrices S, T, V and core Hamiltonian. ----
S = zeros(n_basis, n_basis)
T = zeros(n_basis, n_basis)
V = zeros(n_basis, n_basis)
H_core = zeros(n_basis, n_basis)
z3 = zeros(3)

for mu to n_basis {
    for nu to n_basis {
        S[mu, nu] = contracted_one_elec(mu, nu, 0, z3, 0.0)
        T[mu, nu] = contracted_one_elec(mu, nu, 1, z3, 0.0)
        v_total = 0.0
        for atom to n_atoms {
            Rc = center_of(atom)
            v_total += contracted_one_elec(mu, nu, 2, Rc, Z[atom])
        }
        V[mu, nu] = v_total
        H_core[mu, nu] = T[mu, nu] + V[mu, nu]
    }
}

# ---- Two-electron integrals (the O(N^4) hot spot), parallelized. ----
narrate "Building two-electron integrals (parallel)..."
eri = zeros(n_basis, n_basis, n_basis, n_basis)
parallel for mu to n_basis {
    for nu to n_basis {
        for lam to n_basis {
            for sig to n_basis {
                eri[mu, nu, lam, sig] = contracted_eri(mu, nu, lam, sig)
            }
        }
    }
}

# ---- Symmetric orthogonalization X = S^{-1/2}. ----
s_vals = zeros(n_basis)
s_vecs = zeros(n_basis, n_basis)
eig_sym(S, s_vals, s_vecs)
D_inv_sqrt = zeros(n_basis, n_basis)
for i to n_basis {
    D_inv_sqrt[i, i] = 1.0 / sqrt(s_vals[i])
}
X = s_vecs @ D_inv_sqrt @ transpose(s_vecs)

# ---- Nuclear repulsion. ----
E_nuc = 0.0
for i to n_atoms {
    for j to n_atoms {
        if j > i {
            E_nuc += Z[i] * Z[j] / sqrt(dist_sq(center_of(i), center_of(j)))
        }
    }
}

# ---- SCF loop. ----
P = zeros(n_basis, n_basis)
E_total = 0.0
tol = 1.0e-8
max_iter = 100

for iter to max_iter {
    G = zeros(n_basis, n_basis)
    for mu to n_basis {
        for nu to n_basis {
            g = 0.0
            for lam to n_basis {
                for sig to n_basis {
                    j_int = eri[mu, nu, lam, sig]
                    k_int = eri[mu, lam, nu, sig]
                    g += P[lam, sig] * (j_int - 0.5 * k_int)
                }
            }
            G[mu, nu] = g
        }
    }
    F = H_core + G
    Fp = transpose(X) @ F @ X
    eps_vals = zeros(n_basis)
    Cp = zeros(n_basis, n_basis)
    eig_sym(Fp, eps_vals, Cp)
    C = X @ Cp

    P_new = zeros(n_basis, n_basis)
    for mu to n_basis {
        for nu to n_basis {
            psum = 0.0
            for a to n_occ {
                psum += C[mu, a] * C[nu, a]
            }
            P_new[mu, nu] = 2.0 * psum
        }
    }

    E_elec = 0.0
    for mu to n_basis {
        for nu to n_basis {
            E_elec += 0.5 * P_new[mu, nu] * (H_core[mu, nu] + F[mu, nu])
        }
    }
    E_total = E_elec + E_nuc

    delta = 0.0
    for mu to n_basis {
        for nu to n_basis {
            d = P_new[mu, nu] - P[mu, nu]
            if d < 0.0 { d = -d }
            if d > delta { delta = d }
        }
    }
    P = P_new
    if delta < tol { break }
}

show n_atoms
show E_total

# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors