"""Nonsymmetric 3D AQUAL and principal-constraint checks for the final candidate."""
import json
import numpy as np
from scipy.interpolate import PchipInterpolator

from dark_cmc_model import BASE, OUT, CMCResponse, candidate, digest
from dark_gravity_3d import Grid, EnergyLaw, controls, solve, resize, rms
from dark_constrained_dynamics import constraint_solve


class CMCEnergy(EnergyLaw):
    def __init__(self):
        self.newton=False
        self.law=CMCResponse()
        x=np.r_[0.,np.geomspace(1e-15,1e4,48001)]
        mu,_=self.law.evaluate(x)
        force=x*mu
        assert np.all(np.diff(force)>0)
        self.first=PchipInterpolator(x,force,extrapolate=False)
        self.second=self.first.derivative()
        self.primitive=self.first.antiderivative()
        self.offset=float(self.primitive(0.))
        probe=np.geomspace(.001,100,18001)
        self.response_error=float(np.max(abs(self.first(probe)/probe-self.law.evaluate(probe)[0])))
        assert self.response_error<2e-7

    def evaluate(self,g):
        energy,mu,radial=super().evaluate(g)
        return energy,np.where(g==0,1.,mu),radial


def main():
    OUT.mkdir(parents=True,exist_ok=True)
    energy=CMCEnergy()
    record=candidate(energy.law)
    direction=np.array([2.,3.,5.])/np.sqrt(38.)
    external=direction*2.32e-10/record["galaxy_scale_m_s2"]
    check=controls(energy,external)
    homogeneous=[]
    # Include the most weakly elliptic IR direction, not just zero field.
    for strength in [0.,1.0663914801702836e-6,6.417863741157173]:
        grid=Grid(9,direction*strength)
        coeff=grid.energy_gradient(np.zeros((9,9,9)),energy)[2]
        homogeneous.append({"homogeneous_acceleration":strength,**constraint_solve(grid,coeff)})
    previous=None
    records=[]
    for n in [17,25,33,49,65]:
        grid=Grid(n,external)
        initial=None if previous is None else resize(previous,n)
        potential,item=solve(grid,energy,initial)
        if previous is not None:
            restricted=resize(potential,len(previous))
            pg=Grid(len(previous),external)
            item["refinement_gradient_relative_L2"]=rms(pg.gradient(restricted-previous))/rms(pg.gradient(restricted))
        if n in [17,33,49]:
            item["principal_constraint_block_solve"]=constraint_solve(grid,grid.energy_gradient(potential,energy)[2])
        if n==49:
            old_path=BASE/"outputs/dark_gravity_3d/Am0p45_T4p25/field_n49.npz"
            old=np.load(old_path)["potential"]
            item["change_from_previous_transition_gradient_relative_L2"]=rms(grid.gradient(potential-old))/rms(grid.gradient(old))
            item["previous_field_sha256"]=digest(old_path)
        path=OUT/f"field_3d_n{n}.npz"
        np.savez_compressed(path,potential=potential,source=grid.source,external=external,
                            acceleration=-(external[:,None,None,None]+grid.gradient(potential)))
        item["field_sha256"]=digest(path)
        records.append(item)
        previous=potential
        (OUT/"3d_checkpoint.json").write_text(json.dumps(records,indent=2,allow_nan=False)+"\n")
        print(json.dumps({k:v for k,v in item.items() if k!="history"}),flush=True)
    report={"source_sha256":digest(__file__),"candidate":record,"controls":check,
            "homogeneous_constraint_controls":homogeneous,"fields":records,
            "input_sha256":{name:digest(BASE/name) for name in ["dark_cmc_model.py","dark_gravity_3d.py",
                "dark_constrained_dynamics.py","dark_constrained_metric.py"]},
            "scope":"Synthetic periodic density contrast, tilted external acceleration, static AQUAL and weak-background principal metric constraints; no full spacetime evolution"}
    (OUT/"3d_summary.json").write_text(json.dumps(report,indent=2,allow_nan=False)+"\n")


if __name__=="__main__":
    main()
