#!/usr/bin/env python3
"""Check the RHF formal specification against the imported native implementation.

Default: bounded synthetic-tensor checks, no molecular SCF. Use --molecule
for a complete H2O/STO-3G calculation; submit that mode through the project's
calculation scheduler when collecting scientific acceptance evidence.
This is numerical testing, not a proof of the compiled C++ program.
"""
from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path
import sys

import numpy as np
import vibeqc
from vibeqc import _vibeqc_core as core


class Checks:
    def __init__(self):
        self.residuals = {}
        self.negative_controls = []

    def close(self, name, actual, expected, tolerance=2e-10):
        actual, expected = np.asarray(actual), np.asarray(expected)
        if actual.shape != expected.shape:
            raise AssertionError(f"{name}: shape {actual.shape} != {expected.shape}")
        if not (np.isfinite(actual).all() and np.isfinite(expected).all()):
            raise AssertionError(f"{name}: nonfinite value")
        error = float(np.max(np.abs(actual - expected), initial=0.0))
        self.residuals[name] = max(error, self.residuals.get(name, 0.0))
        if error > tolerance:
            raise AssertionError(f"{name}: residual {error:.6g} > {tolerance:.6g}")

    def rejects(self, name, actual, expected, tolerance=2e-10):
        probe = Checks()
        try:
            probe.close(name, actual, expected, tolerance)
        except AssertionError:
            self.negative_controls.append(name)
            return
        raise AssertionError(f"negative control not detected: {name}")


def density_checks(check, S, Cocc, D, tolerance=2e-8):
    nocc = Cocc.shape[1]
    check.close("occupied_metric", Cocc.T @ S @ Cocc, np.eye(nocc), tolerance)
    check.close("density_definition", D, 2 * Cocc @ Cocc.T, tolerance)
    check.close("density_symmetry", D, D.T, tolerance)
    check.close("metric_idempotency", D @ S @ D, 2 * D, tolerance)
    check.close("electron_count", np.trace(D @ S), 2.0 * nocc, tolerance)


def oracle(eri, D):
    """Independent indexed contractions, not a call back into vibe-qc."""
    J = np.einsum("mnls,ls->mn", eri, D, optimize=False)
    K = np.einsum("mlns,ls->mn", eri, D, optimize=False)
    return J, K, J - 0.5 * K


def synthetic_checks(check):
    rng = np.random.default_rng(20261101)
    cases = []
    for n in (1, 2, 3, 5, 8):
        a = rng.normal(size=(n, n))
        S = a @ a.T + np.eye(n)
        se, su = np.linalg.eigh(S)
        invsqrt = (su / np.sqrt(se)) @ su.T
        q, _ = np.linalg.qr(rng.normal(size=(n, n)))
        nocc = max(1, n // 2)
        Cocc = (invsqrt @ q)[:, :nocc]
        D = 2 * Cocc @ Cocc.T
        density_checks(check, S, Cocc, D, 2e-10)
        u, _ = np.linalg.qr(rng.normal(size=(nocc, nocc)))
        check.close("occupied_rotation", 2 * (Cocc @ u) @ (Cocc @ u).T, D)
        b = rng.normal(size=(n, n, n + 1))
        b = 0.5 * (b + b.transpose(1, 0, 2))
        eri = np.ascontiguousarray(np.einsum("mnp,lsp->mnls", b, b))
        h = rng.normal(size=(n, n)); h = 0.5 * (h + h.T)
        X = rng.normal(size=(n, n)); X = 0.5 * (X + X.T)
        J, K, G = oracle(eri, D)
        native_g = np.asarray(core.build_fock_g(eri, D))
        check.close("native_coulomb", core.build_coulomb(eri, D), J)
        check.close("native_exchange", core.build_exchange(eri, D), K)
        check.close("native_fock_contraction", native_g, G)
        check.close("native_fock_symmetry", native_g, native_g.T)
        gx = np.asarray(core.build_fock_g(eri, X))
        check.close("interaction_reciprocity", np.sum(X * native_g), np.sum(D * gx))
        F = h + native_g

        def energy(p):
            return float(np.sum(p * h) + 0.5 * np.sum(p * core.build_fock_g(eri, p)))

        check.close("half_fock_energy", energy(D), 0.5 * np.sum(D * (h + F)))
        for t in (-0.3, 0.125, 0.7):
            expected = energy(D) + t * np.sum(X * F) + 0.5 * t * t * np.sum(X * gx)
            check.close("exact_energy_expansion", energy(D + t * X), expected)
        step = 2e-4
        check.close("energy_directional_derivative",
                    (energy(D + step * X) - energy(D - step * X)) / (2 * step),
                    np.sum(X * F), 2e-8)
        if n == 5:
            check.rejects("missing_density_factor_two", np.trace((D / 2) @ S), 2.0 * nocc)
            check.rejects("missing_overlap_metric", D @ D, 2 * D)
            check.rejects("exchange_factor_one", J - K, G)
            check.rejects("exchange_wrong_sign", J + 0.5 * K, G)
            wrong_layout = np.ascontiguousarray(eri.transpose(0, 2, 1, 3))
            check.rejects("exchange_index_permutation", core.build_fock_g(wrong_layout, D), G)
            check.rejects("missing_energy_half", np.sum(D * (h + F)), energy(D))
        cases.append({"ao": n, "occupied": nocc})
    return cases


def molecule_checks(check):
    # Positions are exactly specified in bohr, also in the reference input.
    atoms = [(8, [0.0, 0.0, 0.0]), (1, [0.0, 1.4, 1.1]), (1, [0.0, -1.4, 1.1])]
    mol = vibeqc.Molecule([vibeqc.Atom(z, r) for z, r in atoms], 0, 1)
    basis = vibeqc.BasisSet(mol, "sto-3g")
    opts = vibeqc.RHFOptions()
    opts.max_iter = 150
    opts.conv_tol_energy = 1e-12
    opts.conv_tol_grad = 1e-10
    opts.damping = 0.0
    opts.density_fit = False
    opts.initial_guess = vibeqc.InitialGuess.HCORE
    opts.scf_mode = vibeqc.SCFMode.CONVENTIONAL
    opts.stability_check = False
    result = vibeqc.run_rhf(mol, basis, opts)
    if not result.converged:
        raise AssertionError("H2O/STO-3G RHF did not converge")
    S = np.asarray(core.compute_overlap(basis))
    h = np.asarray(core.compute_kinetic(basis)) + np.asarray(core.compute_nuclear(basis, mol))
    eri = np.asarray(core.compute_eri(basis))
    D = np.asarray(result.density)
    C = np.asarray(result.mo_coeffs)
    eps = np.asarray(result.mo_energies)
    F = np.asarray(result.fock)
    density_checks(check, S, C[:, :5], D)
    _, _, G = oracle(eri, D)
    check.close("scf_fock", F, h + G, 2e-8)
    check.close("roothaan_residual", F @ C, (S @ C) * eps, 2e-8)
    check.close("stationarity_residual", F @ D @ S, S @ D @ F, 2e-8)
    energy = float(np.sum(D * h) + 0.5 * np.sum(D * G))
    check.close("scf_electronic_energy", result.e_electronic, energy, 2e-9)
    check.close("scf_total_energy", result.energy, energy + mol.nuclear_repulsion(), 2e-9)
    return {"name": "H2O", "basis": "STO-3G", "charge": 0, "multiplicity": 1,
            "geometry_bohr": atoms, "scf_mode": "CONVENTIONAL", "density_fit": False,
            "energy_hartree": float(result.energy), "iterations": int(result.n_iter),
            "converged": bool(result.converged)}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--molecule", action="store_true")
    parser.add_argument("--output", type=Path, help="JSON report; may include local runtime paths")
    args = parser.parse_args()
    checks = Checks()
    cases = synthetic_checks(checks)
    molecule = molecule_checks(checks) if args.molecule else None
    native = Path(core.__file__).resolve()
    report = {"schema": "vibeqc.formal-rhf-checks/1", "verdict": "PASS",
              "scope": "numerical implementation checks; not a compiled-program proof",
              "version": vibeqc.__version__, "python": sys.executable,
              "python_module": str(Path(vibeqc.__file__).resolve()),
              "native_module": str(native),
              "native_sha256": hashlib.sha256(native.read_bytes()).hexdigest(),
              "synthetic_cases": cases, "molecule": molecule,
              "max_absolute_residuals": checks.residuals,
              "negative_controls_detected": checks.negative_controls}
    text = json.dumps(report, indent=2, sort_keys=True)
    if args.output:
        args.output.write_text(text + "\n")
    print(text)


if __name__ == "__main__":
    main()