"""Constraint, principal-mode and nonsymmetric 3D checks of the new candidate.

No scalar wave speed is assigned at the kinetic-conformal point: constraints
are derived there directly. A static 3D benchmark is not a spacetime simulation.
"""
from __future__ import annotations

import argparse
import json

import numpy as np
import sympy as sp
from scipy.fft import fftn, ifftn
from scipy.interpolate import PchipInterpolator
from scipy.sparse.linalg import LinearOperator, cg

from dark_constrained_metric import BASE, OUT, ConstrainedResponse, fit, digest
from dark_gravity_3d import Grid, EnergyLaw, controls, solve, resize, rms


class ConstrainedEnergy(EnergyLaw):
    def __init__(self, detuning=1e-6):
        self.newton = False
        self.law = ConstrainedResponse(detuning)
        x = np.r_[0.,np.geomspace(1e-11,1e4,30001)]
        mu,_ = self.law.evaluate(x)
        first = x*mu
        assert np.all(np.diff(first)>0)
        self.first = PchipInterpolator(x,first,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)
        mu = np.where(g==0,self.law.detuning/(1+self.law.detuning),mu)
        return energy,mu,radial


def symbolic():
    lam, zdot, lapb, a, k, omega = sp.symbols("lambda_H zdot lapb alpha_direction k omega",real=True)
    kinetic = 3*(1-3*lam)*zdot**2+2*(3*lam-1)*zdot*lapb+(1-lam)*lapb**2
    exact_kinetic = sp.simplify(kinetic.subs(lam,sp.Rational(1,3)))
    assert exact_kinetic == sp.Rational(2,3)*lapb**2
    matrix = sp.Matrix([[2*k**2,2*k**2,0],[2*k**2,a*k**2,0],[0,0,sp.Rational(2,3)*k**4]])
    determinant = sp.factor(matrix.det())
    assert sp.simplify(determinant-sp.Rational(4,3)*(a-2)*k**8)==0
    assert not determinant.has(omega)
    diagonal = sp.Matrix([[1-lam,-lam,-lam],[-lam,1-lam,-lam],[-lam,-lam,1-lam]])
    eigenvalues = diagonal.subs(lam,sp.Rational(1,3)).eigenvals()
    assert eigenvalues == {sp.Integer(0):1,sp.Integer(1):2}
    hubble = sp.symbols("H",real=True)
    flrw = (3-9*lam)*hubble**2
    assert flrw.subs(lam,sp.Rational(1,3))==0
    return {"scalar_kinetic_at_lambda_H_one_third":str(exact_kinetic),
            "unreduced_scalar_constraint_determinant":str(determinant),
            "determinant_independent_of_frequency":True,
            "kinetic_metric_hessian_rank":5,"kinetic_metric_hessian_trace_nullity":1,
            "phase_space_count":"(14 - 2*3 spatial first-class constraints - 4 second-class constraints)/2 = 2",
            "propagating_scalar_modes_on_regular_branch":0,
            "tensor_speed_squared_over_c_squared":1.,
            "scalar_speed_squared_over_c_squared":None,
            "cosmology":{"homogeneous_trace_kinetic_coefficient":0.,
                "homogeneous_lapse_and_trace_equations":"F(0)=16*pi*G*rho=-16*pi*G*p (c=1)",
                "supports_standard_flat_dust_radiation_FLRW":False}}


def response_screen(law):
    x = np.r_[0.,np.geomspace(1e-12,1e10,60001)]
    mu,radial = law.evaluate(x)
    at,ar = law.lapse_coefficients(x)
    assert np.min(mu)>0 and np.min(radial)>0
    # Retain the unstable lapse direction of the old propagating formulation.
    assert np.min(ar)<-.4
    # The constrained determinant depends on 2-alpha, not on alpha's sign.
    determinant_t = -4*mu
    determinant_r = -4*radial
    assert np.max(determinant_t)<0 and np.max(determinant_r)<0
    coefficient_error = float(max(np.max(abs(2-at-2*mu)),np.max(abs(2-ar-2*radial))))
    assert coefficient_error<1e-14
    # A direction is a convex combination of the two eigenvalues. Endpoints
    # therefore bound every orientation, including tilted wave vectors.
    rows=[]
    for acceleration in [0.,1e-6,.1,1.,6.417863741157173,100.,1e8]:
        mt,mr = map(float,law.evaluate(acceleration))
        atp,arp = map(float,law.lapse_coefficients(acceleration))
        for cosine in [0.,.37,1.]:
            alpha=atp+(arp-atp)*cosine**2
            expected=-4*(mt+(mr-mt)*cosine**2)
            # Use the uneliminated 2x2 matrix, including alpha=0 crossings.
            measured=float(np.linalg.det(np.array([[2.,2.],[2.,alpha]])))
            assert abs(measured-expected)<1e-14
            rows.append({"x":acceleration,"cos_theta":cosine,"lapse_coefficient":alpha,
                         "constraint_determinant":measured})
    return {"minimum_mu":float(mu.min()),"minimum_radial_stiffness":float(radial.min()),
            "minimum_lapse_parallel_coefficient":float(ar.min()),
            "maximum_lapse_direction_coefficient":float(max(at.max(),ar.max())),
            "minimum_principal_ellipticity_margin":float(min(mu.min(),radial.min())),
            "coefficient_identity_error":coefficient_error,"direction_checks":rows,
            "scope":"Exact endpoint criterion given the response; coefficients sampled over the stated x range, with analytic positive IR and Newtonian tails"}


def block_gmres(operator, source, precondition, rtol=2e-11, restart=40, max_cycles=40):
    """Restarted left-preconditioned GMRES with explicit finite contractions.

    This avoids a reproducible floating-status warning in this environment's
    SciPy/BLAS solution update. No warnings are suppressed. Twice-reorthogonalized
    Arnoldi and Givens rotations are followed by a true, unpreconditioned
    residual check. Exact small-system controls are included in the report.
    """
    solution=np.zeros_like(source)
    source_norm=float(np.linalg.norm(source))
    iterations=0
    if source_norm==0:
        return solution,iterations
    for _ in range(max_cycles):
        true_residual=source-operator(solution)
        if np.linalg.norm(true_residual)<=rtol*source_norm:
            return solution,iterations
        residual=precondition(true_residual)
        beta=float(np.linalg.norm(residual))
        assert beta>0 and np.isfinite(beta)
        basis=np.zeros((restart+1,len(source)))
        basis[0]=residual/beta
        hessenberg=np.zeros((restart+1,restart))
        cosines=np.zeros(restart)
        sines=np.zeros(restart)
        transformed=np.zeros(restart+1)
        transformed[0]=beta
        dimension=restart
        for j in range(restart):
            work=precondition(operator(basis[j]))
            for _ in range(2):
                projections=np.einsum("ij,j->i",basis[:j+1],work,optimize=False)
                hessenberg[:j+1,j]+=projections
                work-=np.einsum("i,ij->j",projections,basis[:j+1],optimize=False)
            next_norm=float(np.linalg.norm(work))
            hessenberg[j+1,j]=next_norm
            if next_norm>0:
                basis[j+1]=work/next_norm
            for i in range(j):
                upper,lower=hessenberg[i,j],hessenberg[i+1,j]
                hessenberg[i,j]=cosines[i]*upper+sines[i]*lower
                hessenberg[i+1,j]=-sines[i]*upper+cosines[i]*lower
            diagonal=float(np.hypot(hessenberg[j,j],hessenberg[j+1,j]))
            assert diagonal>0
            cosines[j]=hessenberg[j,j]/diagonal
            sines[j]=hessenberg[j+1,j]/diagonal
            hessenberg[j,j]=diagonal
            hessenberg[j+1,j]=0.
            transformed[j+1]=-sines[j]*transformed[j]
            transformed[j]*=cosines[j]
            iterations+=1
            if abs(transformed[j+1])<=rtol*beta or next_norm<1e-14:
                dimension=j+1
                break
        coefficients=np.zeros(dimension)
        for i in range(dimension-1,-1,-1):
            remainder=float(np.einsum("i,i->",hessenberg[i,i+1:dimension],coefficients[i+1:],optimize=False))
            coefficients[i]=(transformed[i]-remainder)/hessenberg[i,i]
        solution+=np.einsum("i,ij->j",coefficients,basis[:dimension],optimize=False)
        assert np.all(np.isfinite(solution))
    raise RuntimeError("Independent block GMRES did not reach its true residual tolerance")


def gmres_controls():
    rows=[]
    for matrix in [np.eye(5),np.array([[2.,2.],[2.,-.46]]),
                   np.array([[4.,1.,0.],[-2.,3.,.2],[1.,0.,-1.]])]:
        exact=np.arange(1.,len(matrix)+1)
        action=lambda v:np.einsum("ij,j->i",matrix,v,optimize=False)
        value,iterations=block_gmres(action,action(exact),lambda v:v,rtol=1e-12)
        error=float(np.max(abs(value-exact)))
        assert error<1e-11
        rows.append({"dimension":len(matrix),"iterations":iterations,"max_error":error})
    return rows


def constraint_solve(grid, coefficients):
    """Independent full block solve versus an eliminated elliptic equation.

    L=-Delta, H=-div(M grad), A=2L-2H. The two principal constraints are
    2L(zeta+n)=s_zeta, 2L*zeta+A*n=s_n. Eliminating zeta gives
    H*n=(s_zeta-s_n)/2. Mean modes are fixed by the periodic gauge choice.
    """
    n=grid.n
    shape=(n,n,n)
    size=n**3
    coordinate=np.arange(n)*2*np.pi/n
    x,y,z=np.meshgrid(coordinate,coordinate,coordinate,indexing="ij")
    exact_n=np.sin(x+2*y)+.31*np.cos(y-z)+.11*np.sin(2*z-x)
    exact_zeta=.43*np.cos(x-y+z)+.21*np.sin(2*x+z)
    exact_n-=exact_n.mean()
    exact_zeta-=exact_zeta.mean()
    def lap(v):
        return ifftn(grid.k2*fftn(v)).real
    def invlap(v):
        return ifftn(grid.inverse_k2*fftn(v)).real
    def hessian(v):
        return grid.hessian(v,coefficients)
    def block(v):
        zeta, lapse=v.reshape((2,)+shape)
        first=2*lap(zeta+lapse)
        second=2*lap(zeta+lapse)-2*hessian(lapse)
        return np.r_[first.ravel(),second.ravel()]
    exact=np.r_[exact_zeta.ravel(),exact_n.ravel()]
    source=block(exact)
    sz,sn=source.reshape((2,)+shape)
    mu,radial,_,_=coefficients
    mean_coefficient=float(np.mean((2*mu+radial)/3))
    positive_floor=float(min(mu.min(),radial.min()))
    assert positive_floor>0
    op=LinearOperator((size,size),matvec=lambda v:hessian(v.reshape(shape)).ravel(),dtype=float)
    pre=LinearOperator((size,size),matvec=lambda v:(invlap(v.reshape(shape))/mean_coefficient).ravel(),dtype=float)
    lapse,status=cg(op,((sz-sn)/2).ravel(),M=pre,rtol=2e-11,atol=0.,maxiter=900)
    assert status==0
    lapse=lapse.reshape(shape)
    zeta=.5*invlap(sz)-lapse
    eliminated=np.r_[zeta.ravel(),lapse.ravel()]
    def block_pre(v):
        first,second=v.reshape((2,)+shape)
        pn=invlap((first-second)/(2*mean_coefficient))
        pz=.5*invlap(first)-pn
        return np.r_[pz.ravel(),pn.ravel()]
    direct,iterations=block_gmres(block,source,block_pre)
    relative=rms(direct-exact)/rms(exact)
    disagreement=rms(direct-eliminated)/rms(exact)
    residual=rms(block(direct)-source)/rms(source)
    assert relative<2e-8 and disagreement<2e-8 and residual<1e-9
    # Independent self-adjointness and positive quadratic-form checks.
    hv=hessian(exact_n)
    hz=hessian(exact_zeta)
    symmetry=abs(float(np.mean(exact_zeta*hv-exact_n*hz)))
    quadratic=float(np.mean(exact_n*hv))
    lower_bound=positive_floor*float(np.mean(np.sum(grid.gradient(exact_n)**2,axis=0)))
    assert symmetry<1e-10 and quadratic>=lower_bound*(1-1e-12)
    return {"n":n,"direct_block_relative_solution_error":relative,
            "block_vs_eliminated_relative_difference":disagreement,
            "direct_block_relative_residual":residual,"gmres_iterations":iterations,
            "self_adjointness_error":symmetry,"positive_quadratic_form":quadratic,
            "pointwise_ellipticity_lower_bound":lower_bound,
            "minimum_mu":float(mu.min()),"minimum_radial":float(radial.min())}


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--stage",choices=["principal","3d"],default="principal")
    args=parser.parse_args()
    OUT.mkdir(parents=True,exist_ok=True)
    response=ConstrainedResponse()
    candidate=fit(response)
    common={"source_sha256":digest(__file__),"candidate":candidate,
            "input_sha256":{name:digest(BASE/name) for name in ["dark_constrained_metric.py","dark_gravity_3d.py",
                    "dark_joint_solar.py","dark_transition_tweak.py"]}}
    if args.stage=="principal":
        report={**common,"symbolic":symbolic(),"response":response_screen(response),
                "independent_GMRES_controls":gmres_controls()}
        energy=ConstrainedEnergy()
        checks=[]
        for strength in [0.,6.417863741157173]:
            vector=strength*np.array([2.,3.,5.])/np.sqrt(38.)
            grid=Grid(9,vector)
            coeff=grid.energy_gradient(np.zeros((9,9,9)),energy)[2]
            checks.append({"homogeneous_acceleration":strength,**constraint_solve(grid,coeff)})
        report["constant_background_block_controls"]=checks
        (OUT/"principal_summary.json").write_text(json.dumps(report,indent=2,allow_nan=False)+"\n")
        print(json.dumps(report,indent=2),flush=True)
        return
    energy=ConstrainedEnergy()
    external=np.array([2.,3.,5.])/np.sqrt(38.)*2.32e-10/candidate["galaxy_scale_m_s2"]
    check=controls(energy,external)
    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,record=solve(grid,energy,initial)
        if previous is not None:
            restricted=resize(potential,len(previous))
            pg=Grid(len(previous),external)
            record["refinement_gradient_relative_L2"]=rms(pg.gradient(restricted-previous))/rms(pg.gradient(restricted))
        if n in [17,33,49]:
            record["principal_constraint_block_solve"]=constraint_solve(grid,grid.energy_gradient(potential,energy)[2])
        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)))
        record["field_sha256"]=digest(path)
        records.append(record)
        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 record.items() if k!="history"}),flush=True)
    report={**common,"controls":check,"fields":records,
            "scope":"Synthetic periodic density contrast and tilted external acceleration; static AQUAL plus the local principal constrained-metric block, not full spacetime evolution"}
    (OUT/"3d_summary.json").write_text(json.dumps(report,indent=2,allow_nan=False)+"\n")


if __name__=="__main__":
    main()
