#!/usr/bin/env python3
"""Numerical UHF, Coulomb-metric fitting and canonical MP2 correspondence checks.
Run through the calculation scheduler; reports include local runtime paths.
"""
import argparse
import hashlib
import json
import sys
from pathlib import Path
import numpy as np
import vibeqc as v
from vibeqc import _vibeqc_core as core
from vibeqc.density_fitting import DensityFitting
sys.path.insert(0, str(Path(__file__).resolve().parent))
from check_rhf import Checks, oracle

parser=argparse.ArgumentParser(description=__doc__)
parser.add_argument('--cpp',action='store_true',help='also require resident C++ density-fitting bindings')
args=parser.parse_args()
check = Checks()
def options(cls):
    o=cls(); o.max_iter=150; o.conv_tol_energy=1e-12; o.conv_tol_grad=1e-10
    o.density_fit=False; o.initial_guess=v.InitialGuess.HCORE
    o.scf_mode=v.SCFMode.CONVENTIONAL; o.stability_check=False
    return o

def integrals(m,b):
    return (np.asarray(core.compute_overlap(b)),
            np.asarray(core.compute_kinetic(b))+np.asarray(core.compute_nuclear(b,m)),
            np.asarray(core.compute_eri(b)))

m=v.Molecule([v.Atom(8,[0.,0.,0.]),v.Atom(1,[0.,0.,1.8])],0,2)
b=v.BasisSet(m,'sto-3g'); u=v.run_uhf(m,b,options(v.UHFOptions))
assert u.converged
S,h,eri=integrals(m,b)
A=np.asarray(u.density_alpha); B=np.asarray(u.density_beta)
J,_,_=oracle(eri,A+B)
for spin,D,n in [('alpha',A,5),('beta',B,4)]:
    C=np.asarray(getattr(u,'mo_coeffs_'+spin))[:,:n]
    F=np.asarray(getattr(u,'fock_'+spin)); K=oracle(eri,D)[1]
    check.close(spin+'_density',D,C@C.T,2e-8)
    check.close(spin+'_metric',D@S@D,D,2e-8)
    check.close(spin+'_count',np.trace(D@S),float(n),2e-8)
    check.close(spin+'_fock',F,h+J-K,2e-8)
    check.close(spin+'_stationarity',F@D@S,S@D@F,2e-8)
EU=np.sum((A+B)*h)+.5*np.sum((A+B)*J)-.5*np.sum(A*oracle(eri,A)[1]+B*oracle(eri,B)[1])
check.close('uhf_energy',u.e_electronic,EU,2e-8)
check.rejects('uhf_wrong_exchange_half',h+J-.5*oracle(eri,A)[1],u.fock_alpha)

m2=v.Molecule([v.Atom(8,[0.,0.,0.]),v.Atom(1,[0.,1.4,1.1]),v.Atom(1,[0.,-1.4,1.1])],0,1)
b2=v.BasisSet(m2,'sto-3g'); r=v.run_rhf(m2,b2,options(v.RHFOptions)); assert r.converged
S,h,eri=integrals(m2,b2); C=np.asarray(r.mo_coeffs); e=np.asarray(r.mo_energies)
D=np.asarray(r.density); co=C[:,:5]; cv=C[:,5:]
mo=np.einsum('mnls,mi,na,lj,sb->iajb',eri,co,cv,co,cv,optimize=True)
den=e[:5,None,None,None]+e[None,None,:5,None]-e[None,5:,None,None]-e[None,None,None,5:]
assert np.all(den<0)
x=mo; y=mo.transpose(0,3,2,1)
os=float(np.sum(x*x/den)); ss=float(np.sum((x*x-x*y)/den))
o=v.MP2Options(); o.n_frozen_core=0; o.density_fit=False
p=v.run_mp2(m2,b2,r,o)
check.close('mp2_os',p.e_os,os,2e-10); check.close('mp2_ss',p.e_ss,ss,2e-10)
check.close('mp2_correlation',p.e_correlation,os+ss,2e-10)
check.rejects('mp2_wrong_denominator_sign',-(os+ss),p.e_correlation)
check.rejects('mp2_missing_exchange',2*os,p.e_correlation)
aux=v.BasisSet(m2,'def2-universal-jkfit')
V=np.asarray(core.compute_2c_eri(aux)); T=np.asarray(core.compute_3c_eri(b2,aux))
na=V.shape[0]; n=D.shape[0]; T=T.reshape(na,n*n)
L=np.linalg.cholesky(V); factors=np.linalg.solve(L,T); gram=factors.T@factors
check.close('df_metric_reconstruction',gram,T.T@np.linalg.solve(V,T),2e-8)
fitted=gram.reshape(n,n,n,n); jf,kf,_=oracle(fitted,D)
df=DensityFitting(b2,aux)
check.close('df_public_j',df.build_J(D),jf,2e-8)
check.close('df_public_k',df.build_K_density(D),kf,2e-8)
if args.cpp:
    native_df=core._DensityFittingStorage(b2,aux)
    check.close('df_cpp_j',native_df.build_J(D),jf,2e-8)
    check.close('df_cpp_k',native_df.build_K_density(D),kf,2e-8)
check.rejects('df_omitted_metric',T.T@T,gram)
report={'verdict':'PASS','cpp_density_fitting_checked':args.cpp,'version':v.__version__,'native_module':core.__file__,
        'native_sha256':hashlib.sha256(Path(core.__file__).read_bytes()).hexdigest(),
        'uhf_oh_energy':float(u.energy),'rhf_water_energy':float(r.energy),
        'mp2_water_os':os,'mp2_water_ss':ss,'mp2_water_total':float(r.energy+os+ss),
        'max_absolute_residuals':check.residuals,'negative_controls':check.negative_controls}
Path('extensions-checks.json').write_text(json.dumps(report,indent=2)+'\n')
print(json.dumps(report,indent=2))