# Polyatomic geometry optimization: relax water (H2O) to its equilibrium by
# numerical-gradient Newton on the RHF/STO-3G energy surface.
#
# To avoid the 6 translation/rotation zero modes (which make a full 9-Cartesian
# Hessian singular), water is optimized in 2 symmetry-reduced internal DOF: O is
# pinned at the origin and the two H's sit mirror-symmetric at (+/-x, y). The
# bond length and H-O-H angle are read off at the end. 2D Newton uses the
# numerical gradient and Hessian (central differences) of the energy.
#
# Validation: r_OH -> ~1.87 bohr (~0.99 A) and angle -> ~100 deg (HF/STO-3G water).

pi = 3.141592653589793
n_prim = 3

# ─── McMurchie-Davidson integral engine (reads bf* globals by index) ───
def boys_n(nn, T) {
    if T < 1.0e-12 { return 1.0 / (2.0 * nn + 1.0) }
    return simpson(fn(t) -> pow(t, 2.0 * nn) * exp(-T * t * t), 0.0, 1.0, 100)
}
def hermite_E(i, j, t, Qx, a, b) {
    if i < 0.0 { return 0.0 }
    if j < 0.0 { return 0.0 }
    if t < 0.0 { return 0.0 }
    if t > i + j { return 0.0 }
    p = a + b
    if i == 0.0 {
        if j == 0.0 {
            mu = a * b / p
            return exp(-mu * Qx * Qx)
        }
        xpb = a / p * Qx
        return (0.5 / p * hermite_E(i, j - 1.0, t - 1.0, Qx, a, b)
             + xpb * hermite_E(i, j - 1.0, t, Qx, a, b)
             + (t + 1.0) * hermite_E(i, j - 1.0, t + 1.0, Qx, a, b))
    }
    xpa = -b / p * Qx
    return (0.5 / p * hermite_E(i - 1.0, j, t - 1.0, Qx, a, b)
         + xpa * hermite_E(i - 1.0, j, t, Qx, a, b)
         + (t + 1.0) * hermite_E(i - 1.0, j, t + 1.0, Qx, a, b))
}
def hermite_R(t, u, v, nn, p, PCx, PCy, PCz, RPC2) {
    if t == 0.0 {
        if u == 0.0 {
            if v == 0.0 { return pow(-2.0 * p, nn) * boys_n(nn, p * RPC2) }
            res = PCz * hermite_R(t, u, v - 1.0, nn + 1.0, p, PCx, PCy, PCz, RPC2)
            if v > 1.0 { res += (v - 1.0) * hermite_R(t, u, v - 2.0, nn + 1.0, p, PCx, PCy, PCz, RPC2) }
            return res
        }
        res = PCy * hermite_R(t, u - 1.0, v, nn + 1.0, p, PCx, PCy, PCz, RPC2)
        if u > 1.0 { res += (u - 1.0) * hermite_R(t, u - 2.0, v, nn + 1.0, p, PCx, PCy, PCz, RPC2) }
        return res
    }
    res = PCx * hermite_R(t - 1.0, u, v, nn + 1.0, p, PCx, PCy, PCz, RPC2)
    if t > 1.0 { res += (t - 1.0) * hermite_R(t - 2.0, u, v, nn + 1.0, p, PCx, PCy, PCz, RPC2) }
    return res
}
def overlap_1d(i, j, Qx, a, b) {
    return hermite_E(i, j, 0.0, Qx, a, b) * sqrt(pi / (a + b))
}
def kinetic_1d(i, j, Qx, a, b) {
    term = b * (2.0 * j + 1.0) * overlap_1d(i, j, Qx, a, b)
    term += -2.0 * b * b * overlap_1d(i, j + 2.0, Qx, a, b)
    term += -0.5 * j * (j - 1.0) * overlap_1d(i, j - 2.0, Qx, a, b)
    return term
}
def prim_overlap(a, b, mu, nu) {
    sx = overlap_1d(bfL[mu], bfL[nu], bfCx[mu] - bfCx[nu], a, b)
    sy = overlap_1d(bfM[mu], bfM[nu], bfCy[mu] - bfCy[nu], a, b)
    sz = overlap_1d(bfN[mu], bfN[nu], bfCz[mu] - bfCz[nu], a, b)
    return sx * sy * sz
}
def prim_kinetic(a, b, mu, nu) {
    sx = overlap_1d(bfL[mu], bfL[nu], bfCx[mu] - bfCx[nu], a, b)
    sy = overlap_1d(bfM[mu], bfM[nu], bfCy[mu] - bfCy[nu], a, b)
    sz = overlap_1d(bfN[mu], bfN[nu], bfCz[mu] - bfCz[nu], a, b)
    tx = kinetic_1d(bfL[mu], bfL[nu], bfCx[mu] - bfCx[nu], a, b)
    ty = kinetic_1d(bfM[mu], bfM[nu], bfCy[mu] - bfCy[nu], a, b)
    tz = kinetic_1d(bfN[mu], bfN[nu], bfCz[mu] - bfCz[nu], a, b)
    return tx * sy * sz + sx * ty * sz + sx * sy * tz
}
def prim_nuclear(a, b, mu, nu, Cx, Cy, Cz) {
    la = bfL[mu]
    ma = bfM[mu]
    na = bfN[mu]
    lb = bfL[nu]
    mb = bfM[nu]
    nb = bfN[nu]
    p = a + b
    Px = (a * bfCx[mu] + b * bfCx[nu]) / p
    Py = (a * bfCy[mu] + b * bfCy[nu]) / p
    Pz = (a * bfCz[mu] + b * bfCz[nu]) / p
    PCx = Px - Cx
    PCy = Py - Cy
    PCz = Pz - Cz
    RPC2 = PCx * PCx + PCy * PCy + PCz * PCz
    qx = bfCx[mu] - bfCx[nu]
    qy = bfCy[mu] - bfCy[nu]
    qz = bfCz[mu] - bfCz[nu]
    val = 0.0
    for t to la + lb + 1.0 {
        for u to ma + mb + 1.0 {
            for v to na + nb + 1.0 {
                e = hermite_E(la, lb, t, qx, a, b) * hermite_E(ma, mb, u, qy, a, b) * hermite_E(na, nb, v, qz, a, b)
                val += e * hermite_R(t, u, v, 0.0, p, PCx, PCy, PCz, RPC2)
            }
        }
    }
    return 2.0 * pi / p * val
}
def prim_eri(a, b, c, d, mu, nu, lam, sig) {
    p = a + b
    q = c + d
    Px = (a * bfCx[mu] + b * bfCx[nu]) / p
    Py = (a * bfCy[mu] + b * bfCy[nu]) / p
    Pz = (a * bfCz[mu] + b * bfCz[nu]) / p
    Qx = (c * bfCx[lam] + d * bfCx[sig]) / q
    Qy = (c * bfCy[lam] + d * bfCy[sig]) / q
    Qz = (c * bfCz[lam] + d * bfCz[sig]) / q
    alpha = p * q / (p + q)
    PQx = Px - Qx
    PQy = Py - Qy
    PQz = Pz - Qz
    RPQ2 = PQx * PQx + PQy * PQy + PQz * PQz
    qx1 = bfCx[mu] - bfCx[nu]
    qy1 = bfCy[mu] - bfCy[nu]
    qz1 = bfCz[mu] - bfCz[nu]
    qx2 = bfCx[lam] - bfCx[sig]
    qy2 = bfCy[lam] - bfCy[sig]
    qz2 = bfCz[lam] - bfCz[sig]
    val = 0.0
    for t1 to bfL[mu] + bfL[nu] + 1.0 {
        for u1 to bfM[mu] + bfM[nu] + 1.0 {
            for v1 to bfN[mu] + bfN[nu] + 1.0 {
                e1 = hermite_E(bfL[mu], bfL[nu], t1, qx1, a, b) * hermite_E(bfM[mu], bfM[nu], u1, qy1, a, b) * hermite_E(bfN[mu], bfN[nu], v1, qz1, a, b)
                for t2 to bfL[lam] + bfL[sig] + 1.0 {
                    for u2 to bfM[lam] + bfM[sig] + 1.0 {
                        for v2 to bfN[lam] + bfN[sig] + 1.0 {
                            e2 = hermite_E(bfL[lam], bfL[sig], t2, qx2, c, d) * hermite_E(bfM[lam], bfM[sig], u2, qy2, c, d) * hermite_E(bfN[lam], bfN[sig], v2, qz2, c, d)
                            sign = pow(-1.0, t2 + u2 + v2)
                            val += e1 * e2 * sign * hermite_R(t1 + t2, u1 + u2, v1 + v2, 0.0, alpha, PQx, PQy, PQz, RPQ2)
                        }
                    }
                }
            }
        }
    }
    return 2.0 * pow(pi, 2.5) / (p * q * sqrt(p + q)) * val
}
def cont_overlap(mu, nu) {
    s = 0.0
    for ka to n_prim { for kb to n_prim { s += bfC[mu, ka] * bfC[nu, kb] * prim_overlap(bfA[mu, ka], bfA[nu, kb], mu, nu) } }
    return s
}
def cont_kinetic(mu, nu) {
    s = 0.0
    for ka to n_prim { for kb to n_prim { s += bfC[mu, ka] * bfC[nu, kb] * prim_kinetic(bfA[mu, ka], bfA[nu, kb], mu, nu) } }
    return s
}
def cont_nuclear(mu, nu, Cx, Cy, Cz) {
    s = 0.0
    for ka to n_prim { for kb to n_prim { s += bfC[mu, ka] * bfC[nu, kb] * prim_nuclear(bfA[mu, ka], bfA[nu, kb], mu, nu, Cx, Cy, Cz) } }
    return s
}
def cont_eri(mu, nu, lam, sig) {
    s = 0.0
    for ka to n_prim { for kb to n_prim { for kc to n_prim { for kd to n_prim {
        coef = bfC[mu, ka] * bfC[nu, kb] * bfC[lam, kc] * bfC[sig, kd]
        s += coef * prim_eri(bfA[mu, ka], bfA[nu, kb], bfA[lam, kc], bfA[sig, kd], mu, nu, lam, sig)
    } } } }
    return s
}
def set_bf(idx, cx, cy, cz, l, m, n, a0, a1, a2, c0, c1, c2) {
    bfCx[idx] = cx
    bfCy[idx] = cy
    bfCz[idx] = cz
    bfL[idx] = l
    bfM[idx] = m
    bfN[idx] = n
    bfA[idx, 0] = a0
    bfA[idx, 1] = a1
    bfA[idx, 2] = a2
    lt = l + m + n
    bfC[idx, 0] = c0 * pow(2.0 * a0 / pi, 0.75) * pow(4.0 * a0, lt / 2.0)
    bfC[idx, 1] = c1 * pow(2.0 * a1 / pi, 0.75) * pow(4.0 * a1, lt / 2.0)
    bfC[idx, 2] = c2 * pow(2.0 * a2 / pi, 0.75) * pow(4.0 * a2, lt / 2.0)
}

# STO-3G exponents for H and O.
e1s = zeros(11, 3)
esp = zeros(11, 3)
e1s[1, 0] = 3.42525091
e1s[1, 1] = 0.62391373
e1s[1, 2] = 0.16885540
e1s[8, 0] = 130.7093200
e1s[8, 1] = 23.8088610
e1s[8, 2] = 6.4436083
esp[8, 0] = 5.0331513
esp[8, 1] = 1.1695961
esp[8, 2] = 0.3803890

# Water: O (atom 0) + 2 H (atoms 1,2). n_basis = 5 + 1 + 1 = 7, n_occ = 5.
n_atoms = 3
atom_Z = zeros(n_atoms)
atom_x = zeros(n_atoms)
atom_y = zeros(n_atoms)
atom_z = zeros(n_atoms)
atom_Z[0] = 8.0
atom_Z[1] = 1.0
atom_Z[2] = 1.0
n_basis = 7
n_occ = 5

bfL = zeros(n_basis)
bfM = zeros(n_basis)
bfN = zeros(n_basis)
bfCx = zeros(n_basis)
bfCy = zeros(n_basis)
bfCz = zeros(n_basis)
bfA = zeros(n_basis, n_prim)
bfC = zeros(n_basis, n_prim)

# RHF/STO-3G energy at the current nuclear geometry (reads atom_* globals).
def energy() {
    idx = 0.0
    for i to n_atoms {
        zz = atom_Z[i]
        cx = atom_x[i]
        cy = atom_y[i]
        cz = atom_z[i]
        set_bf(idx, cx, cy, cz, 0.0, 0.0, 0.0, e1s[zz, 0], e1s[zz, 1], e1s[zz, 2], 0.15432897, 0.53532814, 0.44463454)
        idx = idx + 1.0
        if zz > 2.0 {
            set_bf(idx, cx, cy, cz, 0.0, 0.0, 0.0, esp[zz, 0], esp[zz, 1], esp[zz, 2], -0.09996723, 0.39951283, 0.70011547)
            idx = idx + 1.0
            set_bf(idx, cx, cy, cz, 1.0, 0.0, 0.0, esp[zz, 0], esp[zz, 1], esp[zz, 2], 0.15591627, 0.60768372, 0.39195739)
            idx = idx + 1.0
            set_bf(idx, cx, cy, cz, 0.0, 1.0, 0.0, esp[zz, 0], esp[zz, 1], esp[zz, 2], 0.15591627, 0.60768372, 0.39195739)
            idx = idx + 1.0
            set_bf(idx, cx, cy, cz, 0.0, 0.0, 1.0, esp[zz, 0], esp[zz, 1], esp[zz, 2], 0.15591627, 0.60768372, 0.39195739)
            idx = idx + 1.0
        }
    }
    for i to n_basis {
        nrm = 1.0 / sqrt(cont_overlap(i, i))
        for k to n_prim { bfC[i, k] = bfC[i, k] * nrm }
    }
    S = zeros(n_basis, n_basis)
    H_core = zeros(n_basis, n_basis)
    for mu to n_basis {
        for nu to n_basis {
            S[mu, nu] = cont_overlap(mu, nu)
            v = 0.0
            for at to n_atoms { v += -atom_Z[at] * cont_nuclear(mu, nu, atom_x[at], atom_y[at], atom_z[at]) }
            H_core[mu, nu] = cont_kinetic(mu, nu) + v
        }
    }
    eri = zeros(n_basis, n_basis, n_basis, n_basis)
    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] = cont_eri(mu, nu, lam, sig) }
            }
        }
    }
    sv = zeros(n_basis)
    svec = zeros(n_basis, n_basis)
    eig_sym(S, sv, svec)
    dih = zeros(n_basis, n_basis)
    for i to n_basis { dih[i, i] = 1.0 / sqrt(sv[i]) }
    X = svec @ dih @ transpose(svec)
    # Nuclear repulsion (general pairwise sum).
    Enuc = 0.0
    for i to n_atoms {
        for j to n_atoms {
            if j > i {
                dx = atom_x[i] - atom_x[j]
                dy = atom_y[i] - atom_y[j]
                dz = atom_z[i] - atom_z[j]
                Enuc += atom_Z[i] * atom_Z[j] / sqrt(dx * dx + dy * dy + dz * dz)
            }
        }
    }
    P = zeros(n_basis, n_basis)
    Etot = 0.0
    Eold = 0.0
    for iter to 100 {
        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 { g += P[lam, sig] * (eri[mu, nu, lam, sig] - 0.5 * eri[mu, lam, nu, sig]) }
                }
                G[mu, nu] = g
            }
        }
        F = H_core + G
        Fp = transpose(X) @ F @ X
        ev = zeros(n_basis)
        Cp = zeros(n_basis, n_basis)
        eig_sym(Fp, ev, Cp)
        C = X @ Cp
        Pn = zeros(n_basis, n_basis)
        for mu to n_basis {
            for nu to n_basis {
                ps = 0.0
                for a to n_occ { ps += C[mu, a] * C[nu, a] }
                Pn[mu, nu] = 2.0 * ps
            }
        }
        Ee = 0.0
        for mu to n_basis {
            for nu to n_basis { Ee += 0.5 * Pn[mu, nu] * (H_core[mu, nu] + F[mu, nu]) }
        }
        Etot = Ee + Enuc
        for mu to n_basis {
            for nu to n_basis { P[mu, nu] = 0.5 * Pn[mu, nu] + 0.5 * P[mu, nu] }
        }
        dd = Etot - Eold
        if dd < 0.0 { dd = -dd }
        if dd < 1.0e-10 { break }
        Eold = Etot
    }
    return Etot
}

# Energy with the H's placed mirror-symmetric at (+/-px, py); O at the origin.
def e_at(px, py) {
    atom_x[0] = 0.0
    atom_y[0] = 0.0
    atom_z[0] = 0.0
    atom_x[1] = px
    atom_y[1] = py
    atom_z[1] = 0.0
    atom_x[2] = -px
    atom_y[2] = py
    atom_z[2] = 0.0
    return energy()
}

# ─── 2D Newton on (px, py): numerical gradient + Hessian (central differences). ───
hh = 0.0005
px = 1.40
py = 1.05
for it to 4 {
    e0 = e_at(px, py)
    axp = e_at(px + hh, py)
    axm = e_at(px - hh, py)
    ayp = e_at(px, py + hh)
    aym = e_at(px, py - hh)
    cpp = e_at(px + hh, py + hh)
    cpm = e_at(px + hh, py - hh)
    cmp = e_at(px - hh, py + hh)
    cmm = e_at(px - hh, py - hh)
    gx = (axp - axm) / (2.0 * hh)
    gy = (ayp - aym) / (2.0 * hh)
    hxx = (axp - 2.0 * e0 + axm) / (hh * hh)
    hyy = (ayp - 2.0 * e0 + aym) / (hh * hh)
    hxy = (cpp - cpm - cmp + cmm) / (4.0 * hh * hh)
    det = hxx * hyy - hxy * hxy
    dpx = -(hyy * gx - hxy * gy) / det
    dpy = -(hxx * gy - hxy * gx) / det
    px = px + dpx
    py = py + dpy
}
# Read off bond length and H-O-H angle (cosine; the test converts to degrees).
r_oh = sqrt(px * px + py * py)
cos_hoh = (py * py - px * px) / (px * px + py * py)
e_opt = e_at(px, py)
show r_oh
show cos_hoh
show e_opt

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