"""Independent finite variations of the nonlinear constrained Hamiltonian.

Off-shell, periodic, conformally curved test fields with canonical scalar
matter. This verifies the full two scalar secondary equations, including
nonlinear lapse acceleration terms; it is not a spacetime evolution.
"""
import json
import numpy as np
from dark_cmc_model import BASE, OUT, digest
from dark_cmc_3d import CMCEnergy
from dark_gravity_3d import Grid, BOX


def checks(n):
    grid=Grid(n,np.zeros(3))
    coordinate=2*np.pi*np.arange(n)/n
    x,y,z=coordinate[:,None,None],coordinate[None,:,None],coordinate[None,None,:]
    lapse=np.exp(.4*np.sin(x)+.3*np.cos(y)+.2*np.sin(z)+.1*np.sin(x+y-z))
    conformal=.07*np.cos(x-y)+.04*np.sin(y+z)+0*x
    test=np.cos(x)+.3*np.sin(y)+.2*np.cos(z)+.4*np.cos(x+y-z)
    phi=.13*np.sin(x+y)+.07*np.cos(y+z)+0*x
    momentum=.3+.04*np.cos(x-z)+0*y
    potential=.08*phi**2+.013
    norm=.07*(1+.2*np.cos(x+y+z))
    q=-.4
    kappa=1.7
    a_star=.075  # Dimensionless benchmark; samples the fitted transition.
    energy=CMCEnergy()

    def hamiltonian(N,Z,tf_norm):
        root=np.exp(3*Z)
        inverse=np.exp(-2*Z)
        gradz=grid.gradient(Z)
        curvature=inverse*(-4*grid.divergence(gradz)-2*np.sum(gradz**2,axis=0))
        acceleration=grid.gradient(np.log(N))
        a2=inverse*np.sum(acceleration**2,axis=0)
        a=np.sqrt(a2)
        U,mu,_=energy.evaluate(a/a_star)
        F=2*a2-4*a_star**2*U
        chi=2*(1-mu)
        kinetic_m=momentum**2/(2*root**2)
        gradient_m=.5*inverse*np.sum(grid.gradient(phi)**2,axis=0)
        rho=kinetic_m+gradient_m+potential
        stress=3*kinetic_m-gradient_m-3*potential
        H=float(np.mean(N*root*(tf_norm-q*q/6-curvature-F+kappa*rho)))
        divchi=grid.divergence(np.exp(Z)*chi*acceleration)/root
        laplaceN=grid.divergence(np.exp(Z)*grid.gradient(N))/root
        first=tf_norm-q*q/6-curvature-F+2*divchi+2*chi*a2+kappa*rho
        # The explicit -2 qdot/N is checked by the homogeneous symbolic
        # identity. At this fixed time qdot=0 in the functional variation.
        second=3*tf_norm+q*q/2+curvature+3*F-2*chi*a2-4*laplaceN/N+kappa*stress
        expected_N=float(np.mean(root*first*test))
        expected_C=-.5*float(np.mean(N*root*second*test))
        return H,expected_N,expected_C,float((a/a_star).min()),float((a/a_star).max())

    baseline=hamiltonian(lapse,conformal,norm)
    records=[]
    for step in [1e-3,3e-4,1e-4,3e-5,1e-5]:
        plus=hamiltonian(lapse+step*test,conformal,norm)[0]
        minus=hamiltonian(lapse-step*test,conformal,norm)[0]
        numerical_N=(plus-minus)/(2*step)
        # Integrate the canonical conformal flow exactly along its path:
        # h -> exp(f)h, P_TF -> exp(-f)P_TF, s -> exp(-3f)s.
        plus=hamiltonian(lapse,conformal+step*test/2,norm*np.exp(-3*step*test))[0]
        minus=hamiltonian(lapse,conformal-step*test/2,norm*np.exp(3*step*test))[0]
        numerical_C=(plus-minus)/(2*step)
        records.append({"step":step,"lapse_variation":numerical_N,"conformal_variation":numerical_C,
            "lapse_absolute_error":abs(numerical_N-baseline[1]),
            "conformal_absolute_error":abs(numerical_C-baseline[2])})
    assert records[-1]["lapse_absolute_error"]<2e-8,records
    assert records[-1]["conformal_absolute_error"]<2e-8,records
    return {"n":n,"acceleration_range_in_units_a_star":list(baseline[3:]),
            "analytic_lapse_variation":baseline[1],"analytic_conformal_variation":baseline[2],
            "finite_variations":records}


def main():
    report={"source_sha256":digest(__file__),"checks":[checks(n) for n in [17,25,33]],
            "scope":__doc__.strip(),"input_sha256":{name:digest(BASE/name) for name in
                ["dark_cmc_model.py","dark_cmc_3d.py","dark_gravity_3d.py","CMC_METRIC_DERIVATION.md"]}}
    OUT.mkdir(parents=True,exist_ok=True)
    (OUT/"variation_summary.json").write_text(json.dumps(report,indent=2,allow_nan=False)+"\n")
    print(json.dumps(report,indent=2),flush=True)


if __name__=="__main__":
    main()
