"""ORAEVA cosmology: scalar-tensor background consistency test (2026-10-14).

Einstein-frame spatially flat FRW, canonical x=phi/M, universal conformal
matter coupling A=exp(alpha*(x-x0)), exponential scalar potential
V=V0*exp(lambda*(x-x0)). Radiation is ideal (trace free).

Evolution variable N=ln(a_E). Units: M=1, reference densities are
relative to 3 M^2 H_ref^2. Here the model's N=0 expansion is close to the
reference (not forced to be EXACTLY H_ref when the scalar is dynamic).

The radiation-dust-potential solution is integrated from N=-22 to N=10.
Observed (Jordan) redshift is 1+z_J=A(0)/(a_E A(N)); don't substitute
1/a_E when comparing physical observables.

Requires numpy, scipy. Usage: python oraeva_coupled_quintessence_check.py
"""
from math import exp, log, sqrt
import numpy as np
from scipy.integrate import solve_ivp
from scipy.optimize import brentq

ALPHA=0.002
BETA=2*ALPHA**2
NI,NF=-22.,10.
OM0,OR0=0.315,9.2e-5
OV0=1-OM0-OR0


def integrate(lam, rtol=5e-12, atol=5e-14, nstart=NI, nend=NF):
    # state: x, y=dx/dN, log(rho_m), log(rho_r), log(V)
    initial = (0., 0., log(OM0)-3*nstart, log(OR0)-4*nstart, log(OV0))
    def system(N, state):
        x,y,lm,lr,lv=state
        rm,rr,v=np.exp([lm,lr,lv])
        H2=(rm+rr+v)/(3.-y*y/2.)
        om,orr,ov=np.array([rm,rr,v])/(3.*H2)
        dlogh = -1.5*om -2.*orr-0.5*y*y
        yp=-(3.+dlogh)*y-3.*ALPHA*om-3.*lam*ov
        return np.array([y,yp,-3.+ALPHA*y,-4.,lam*y])
    sol=solve_ivp(system,(nstart,nend),initial,method='DOP853',dense_output=True,
                  rtol=rtol,atol=atol)
    assert sol.success,sol.message
    return sol,system


def compute(sol,system,N):
    x,y,lm,lr,lv=sol.sol(N)
    rm,rr,v=np.exp([lm,lr,lv]); H2=(rm+rr+v)/(3-y*y/2)
    om,orr,ov=np.array([rm,rr,v])/(3*H2)
    dlogh=-1.5*om-2*orr-0.5*y*y
    yp=system(N,sol.sol(N))[1]
    # d ln a_J/dN = 1+ALPHA*y
    jordan_q=-1.-(dlogh+ALPHA*yp/(1+ALPHA*y)-ALPHA*y)/(1+ALPHA*y)
    return {'N':N,'x':x,'y':y,'beta_eff':-ALPHA*y,
            'beta_ratio':-y/(2*ALPHA),'Om':om,'Or':orr,'Ov':ov,
            'Ok':y*y/6., 'qJ':jordan_q,'H2':H2}


def physical_redshift(sol,N):
    # physical scale factor is normalized a_J(0)=1.
    xnow=sol.sol(0)[0]
    x=sol.sol(N)[0]
    return exp(-N-ALPHA*(x-xnow))-1


def main():
    print('alpha=',ALPHA,'lambda=',2*ALPHA,'beta=',BETA)
    for title,lam in [('tuned exponential',2*ALPHA),('constant potential',0.)]:
        sol,system=integrate(lam)
        print('\n'+title)
        print('era / redshift(J) | beta_eff/beta | Omega_r | Omega_m | Omega_V | Omega_kin | Jordan q')
        for label,zJ in [('radiation',1.e6),('equality vicinity',3400.),
                         ('matter z~100',100.),('matter z~10',10.),
                         ('acceleration today',0.),('future a_J=10',-0.9)]:
            N=brentq(lambda N:physical_redshift(sol,N)-zJ,NI,NF,xtol=1.e-12)
            s=compute(sol,system,N)
            print(f'{label:23s} zJ={zJ:10g} ratio={s["beta_ratio"]:.9f} '
                  f'Or={s["Or"]:.7f} Om={s["Om"]:.7f} Ov={s["Ov"]:.7f} '
                  f'Ok={s["Ok"]:.9f} qJ={s["qJ"]:.7f}')
        print('Today x =',compute(sol,system,0.)['x'])
        # Two independent integrated conservation-law invariants.
        Ns=np.linspace(NI,NF,401)
        vals=sol.sol(Ns)
        x,y,lm,lr,lv=vals
        inv_m=lm+3*Ns-ALPHA*x
        inv_r=lr+4*Ns
        inv_v=lv-lam*x
        assert np.ptp(inv_m)<1e-8
        assert np.ptp(inv_r)<1e-8
        assert np.ptp(inv_v)<1e-8
        print('max conservation-invariant range:',*[f'{np.ptp(v):.2e}' for v in (inv_m,inv_r,inv_v)])

    # Exact radiation-free orbit when slope == 2 alpha.
    # All variables initialized consistently at N=0 on y=-2 alpha.
    def orbit_system(N,z):
        x,y,lm,lv=z
        rm,v=exp(lm),exp(lv)
        H2=(rm+v)/(3-y*y/2)
        om,ov=rm/(3*H2),v/(3*H2)
        E=-1.5*om-.5*y*y
        return [y, -(3+E)*y-3*ALPHA*om-3*(2*ALPHA)*ov,
                -3+ALPHA*y,2*ALPHA*y]
    initial=[0.,-2*ALPHA,log(.315),log(.685)]
    exact=solve_ivp(orbit_system,(0,8),initial,rtol=5e-12,atol=5e-14,
                    method='DOP853',dense_output=True)
    Ns=np.linspace(0,8,101)
    x,y,lm,lv=exact.sol(Ns)
    err=max(np.max(np.abs(x+2*ALPHA*Ns)),
            np.max(np.abs(y+2*ALPHA)),
            np.max(np.abs(lm-log(.315)+(3+BETA)*Ns)),
            np.max(np.abs(lv-log(.685)+2*BETA*Ns)))
    print('\nExact dust+potential orbit: max absolute numerical deviation =',err)
    assert err<1e-8
    print('PASS: exact scaling and matter, radiation, potential continuity checks')

if __name__=='__main__':
    main()