"""Symbolic action/constraint checks and metric benchmarks for the solar candidate."""
from __future__ import annotations

import json
import mpmath as mp
import numpy as np
import sympy as sp
from scipy.integrate import quad
from scipy.optimize import minimize_scalar

from dark_cmc_model import BASE, OUT, CMCResponse, candidate, digest
from dark_constrained_dynamics import gmres_controls
from dark_joint_solar import geodesic_control, shapiro_control
from dark_transition_tweak import REFERENCE
import dark_mercury_check as mercury


def action_checks():
    a,b,c,d,e,f,q=sp.symbols("a b c d e f q",real=True)
    tensor=sp.Matrix([[a,d,e],[d,b,f],[e,f,c]])
    trace=sp.trace(tensor)
    norm=sp.trace(tensor*tensor)
    kinetic=norm-trace**2/3+sp.Rational(2,3)*q*trace+q**2/6
    covariant=norm-trace**2+sp.Rational(2,3)*(trace+q/2)**2
    assert sp.expand(kinetic-covariant)==0
    hessian=sp.hessian(kinetic,[a,b,c,d,e,f])
    assert hessian.rank()==5
    assert hessian*sp.Matrix([1,1,1,0,0,0])==sp.zeros(6,1)
    momentum=tensor-sp.eye(3)*trace/3+sp.eye(3)*q/3
    assert sp.simplify(sp.trace(momentum)-q)==0
    legendre=2*sp.trace(momentum*tensor)-kinetic
    ham=sp.trace(momentum*momentum)-sp.trace(momentum)**2/2
    assert sp.expand(legendre-ham)==0
    assert sp.expand(ham-(norm-trace**2/3-q**2/6))==0
    # Homogeneous equations are checked independently of the local q=0 limit.
    rho,p,kappa,F0,N,H,qdot=sp.symbols("rho p kappa F0 N H qdot",real=True)
    hamiltonian=q**2/6+F0-kappa*rho
    q_evolution=qdot-sp.Rational(3,2)*N*kappa*(rho+p)
    differentiated=q*qdot/3+3*kappa*N*H*(rho+p)
    assert sp.simplify(differentiated.subs({qdot:sp.Rational(3,2)*N*kappa*(rho+p),H:-q/6}))==0
    assert sp.expand((3*H**2-kappa*rho/2+F0/2).subs(H,-q/6)-hamiltonian/2)==0
    # The concrete q(t)=-2/t clock has N=2/(1+w), not N=1 for every fluid.
    t,w=sp.symbols("t w",positive=True)
    clock_q=-2/t
    clock_rho=clock_q**2/(6*kappa)
    lapse=2/(1+w)
    assert sp.simplify(sp.diff(clock_q,t)-sp.Rational(3,2)*lapse*kappa*clock_rho*(1+w))==0
    # Homogeneous TT perturbations: det exp(gamma)=1 fixes the trace.
    # A diagonal plus polarization checks the normalization independently.
    adot,Nc,polar_velocity=sp.symbols("H Nc polar_velocity",real=True)
    extrinsic=sp.diag(adot+polar_velocity/(2*Nc),adot-polar_velocity/(2*Nc),adot)
    shear=sp.trace(extrinsic*extrinsic)-sp.trace(extrinsic)**2/3
    assert sp.simplify(shear-polar_velocity**2/(2*Nc**2))==0
    assert sp.trace(extrinsic)==3*adot
    return {"covariant_square_matches_ADM_action":True,"kinetic_Hessian_rank":5,
            "primary_constraint":"P-sqrt(h)*q(t)=0","primary_constraint_independent_of_lapse":True,
            "primary_self_Poisson_bracket_zero":"ultralocal scalar constraint; antisymmetry cancels the coincident-point terms",
            "Legendre_transform_matches_GR_momentum_Hamiltonian":True,
            "homogeneous_Friedmann_and_clock_identities":True,
            "concrete_clock":"q(t)=-2/t; for a flat constant-w fluid and F(0)=0, N=2/(1+w)",
            "IR_linear_scope":"Formal infinitesimal perturbations about zero acceleration; observed cosmological amplitudes need a nonlinear study",
            "homogeneous_TT_quadratic_action":"(64*pi*G)^-1 integral N*A^3 [(dot gamma_ij/N)^2-(partial_k gamma_ij/A)^2]",
            "homogeneous_TT_kinetic_normalization_checked":True,
            "gravitational_modes_on_regular_local_branch":2,"propagating_scalar_modes":0,
            "tensor_speed_squared_over_c_squared":1.,"scalar_speed_squared_over_c_squared":None}


def response_checks(law):
    x=np.r_[0.,np.geomspace(1e-16,1e11,120001)]
    mu,radial=law.evaluate(x)
    at,ar=law.lapse_coefficients(x)
    assert mu[0]==1 and radial[0]==1 and at[0]==0 and ar[0]==0
    assert np.all(mu>0) and np.all(radial>0)
    assert np.min(ar)<-.4
    im,ir=int(np.argmin(mu)),int(np.argmin(radial))
    opt=minimize_scalar(lambda z:float(law.evaluate(np.exp(z))[1]),
                        bounds=(np.log(x[ir-1]),np.log(x[ir+1])),method="bounded")
    assert opt.success and opt.fun>0
    points=np.unique(np.r_[np.geomspace(1e-13,100,401),
                            law.join_interval,1e-4])
    step=points*1e-5
    force_fd=((points+step)*law.evaluate(points+step)[0]-(points-step)*law.evaluate(points-step)[0])/(2*step)
    derivative_error=float(np.max(abs(force_fd-law.evaluate(points)[1])))
    assert derivative_error<2e-7
    # The Hermite force has a quadratic derivative: test its exact extremum
    # as well as its endpoints, not just a convenient sampling grid.
    polynomial=law.force_join.c[:,0]
    width=np.diff(law.join_interval)[0]
    vertex=-polynomial[1]/(3*polynomial[0])
    extrema=law.join_interval[0]+np.array([0.,width,np.clip(vertex,0.,width)])
    join_min=float(np.min(law.force_join(extrema,1)))
    assert join_min>0
    join_error=float(np.max(abs(law.force_join(law.join_interval,1)
                                -law.raw.evaluate(law.join_interval)[1])))
    assert join_error<1e-15
    # Exact formula for the additional energy integral, against a logarithmic
    # quadrature that resolves both of its widely separated scales.
    eps,ell=law.epsilon,law.cutoff
    integral=eps*ell**2/(ell**2+eps**2)*(np.pi*ell/2+eps*np.log(eps/ell))
    numeric,error=quad(lambda u:float(np.exp(2*u)*law.correction(np.exp(u))[0]),
                       np.log(eps)-25,np.log(ell)+25,epsabs=1e-22,epsrel=1e-10,
                       points=[np.log(eps),np.log(ell)])
    assert abs(numeric/integral-1)<1e-9
    lo,hi=1e-4,100.
    reference_integral=(REFERENCE.normalization*(REFERENCE.integral(hi)-REFERENCE.moment(hi)/hi
                         -REFERENCE.integral(lo)+REFERENCE.moment(lo)/lo)
                        -REFERENCE.c/hi+REFERENCE.normalization*REFERENCE.d/hi**2
                        +REFERENCE.slope*lo**3/3-lo**2/2)
    transition_integral=(27/16)*(-.45)*4.25**2
    join_energy,_=quad(lambda x:float(law.force_join(x)-x*law.raw.evaluate(x)[0]),
                       *law.join_interval,points=[1e-4],epsabs=1e-19,epsrel=1e-8)
    f0=4*(reference_integral+transition_integral+integral+join_energy)
    scale=candidate(law)["galaxy_scale_m_s2"]
    astar=scale/mercury.LIGHT_SPEED**2
    return {"minimum_mu_on_scan":float(mu[im]),"x_at_minimum_mu_on_scan":float(x[im]),
            "minimum_radial_stiffness":float(opt.fun),"x_at_minimum_radial_stiffness":float(np.exp(opt.x)),
            "maximum_lapse_direction_coefficient":float(max(at.max(),ar.max())),
            "minimum_lapse_parallel_coefficient":float(ar.min()),
            "force_derivative_finite_difference_error":derivative_error,
            "force_join_minimum_derivative":join_min,
            "force_join_endpoint_derivative_error":join_error,
            "force_join_energy_correction":float(join_energy),
            "IR_energy_integral":float(integral),"IR_energy_quadrature_relative_error":float(abs(numeric/integral-1)),
            "F_zero_over_a_star_squared_if_F_infinity_zero":float(f0),
            "Lambda_inverse_m2_if_F_infinity_zero":float(-.5*f0*astar**2),
            "Lambda_scope":"An additive vacuum term remains free; this convention is not a cosmological-constant prediction",
            "all_orientation_criterion":"2-alpha_direction=2*[mu*(1-cos(theta)^2)+radial*cos(theta)^2]>0",
            "IR_core_interpolation_absolute_error_bound":1.1e-10,
            "IR_core_interpolation_scope":"Physical mu(0)=1 linearly joined to the first Rust table row at x=1e-14"}


def schwarzschild_checks():
    r,m=sp.symbols("r m",positive=True)
    psi=1+m/(2*r)
    lapse=(1-m/(2*r))/(1+m/(2*r))
    curvature=-8/psi**5*(sp.diff(psi,r,2)+2*sp.diff(psi,r)/r)
    laplace=sp.diff(r**2*psi**2*sp.diff(lapse,r),r)/(psi**6*r**2)
    assert sp.simplify(curvature)==0 and sp.simplify(laplace)==0
    u=sp.symbols("u")
    time_series=sp.series(((1-u/2)/(1+u/2))**2,u,0,3).removeO()
    space_series=sp.series((1+u/2)**4,u,0,2).removeO()
    assert time_series==1-2*u+2*u**2 and space_series==1+2*u
    # Exact timelike geodesic quadratures at fixed areal turning radii.
    # This is an additional geometry benchmark, not an ephemeris estimate.
    mp.mp.dps=70
    mass=mp.mpf(str(mercury.GM))/mp.mpf(str(mercury.LIGHT_SPEED))**2
    axis=mp.mpf(str(mercury.SEMIMAJOR))
    eccentricity=mp.mpf(str(mercury.ECCENTRICITY))
    p=axis*(1-eccentricity**2)/mass
    energy=mp.sqrt(((p-2)**2-4*eccentricity**2)/(p*(p-3-eccentricity**2)))
    angular=mass*p/mp.sqrt(p-3-eccentricity**2)
    def radius(chi):return mass*p/(1+eccentricity*mp.cos(chi))
    def dphi(chi):return mp.sqrt(p/(p-6-2*eccentricity*mp.cos(chi)))
    def dt(chi):
        rr=radius(chi)
        return energy*rr**2/(angular*(1-2*mass/rr))*dphi(chi)/mercury.LIGHT_SPEED
    advance=mp.quad(lambda chi:dphi(chi)-1,[0,mp.pi,2*mp.pi])
    period=mp.quad(dt,[0,mp.pi,2*mp.pi])
    rate=advance/period*mp.mpf(str(mercury.RAD_S_TO_ARCSEC_CENTURY))
    radial_errors=[]
    for chi in [mp.mpf('.3'),mp.mpf('1.1'),mp.mpf('2.5')]:
        rr=radius(chi)
        drdchi=mass*p*eccentricity*mp.sin(chi)/(1+eccentricity*mp.cos(chi))**2
        dtaudchi=rr**2/angular*dphi(chi)
        kinetic=(drdchi/dtaudchi)**2
        expected=energy**2-(1-2*mass/rr)*(1+angular**2/rr**2)
        relative=abs(kinetic/expected-1)
        assert relative<mp.mpf('1e-50')
        radial_errors.append(float(relative))
    return {"isotropic_spatial_curvature_zero":True,"isotropic_spatial_lapse_Laplacian_zero":True,
            "N_squared_series":str(time_series),"spatial_metric_series":str(space_series),
            "beta_PPN":1.,"gamma_PPN":1.,"alpha1_PPN":0.,"alpha2_PPN":0.,
            "exact_geodesic":{"scope":"Schwarzschild high-field branch, with the approximate orbit used as areal turning radii; not the same convention as a precision ephemeris",
                "precision_decimal_digits":70,"advance_radians_per_radial_cycle":str(advance),
                "coordinate_radial_period_seconds":str(period),"arcsec_per_Julian_century":str(rate),
                "radial_first_integral_relative_errors":radial_errors}}


def main():
    OUT.mkdir(parents=True,exist_ok=True)
    law=CMCResponse()
    report={"source_sha256":digest(__file__),"candidate":candidate(law),"action":action_checks(),
            "response":response_checks(law),"metric":schwarzschild_checks(),
            "first_PN_geodesic_controls":geodesic_control(),"Shapiro_control":shapiro_control(1.),
            "independent_block_solver_controls":gmres_controls(),
            "input_sha256":{name:digest(BASE/name) for name in ["dark_cmc_model.py","dark_constrained_dynamics.py",
                 "dark_joint_solar.py","dark_mercury_check.py","CMC_METRIC_DERIVATION.md"]}}
    (OUT/"theory_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()
