"""
RFA_komplett.py  –  Python translation of RFA_komplett.m (GNU Octave / Mathcad 6)
Daniel's cosmological model (Neue Kosmologie)

Physics summary:
  The model treats the universe as a self-gravitating sphere of baryonic matter
  (density rho_b).  The dark-energy-like term is encoded in the dimensionless
  ratio q = Omega_m / Omega_b.  The time-scale-factor relation uses an asinh
  form (cf. LambdaCDM but with a model-specific x(a) instead of Omega_Lambda).

Dependencies:
    pip install numpy matplotlib
"""

import math
import numpy as np
import matplotlib
matplotlib.use("Agg")          # headless – change to "TkAgg" for interactive
import matplotlib.pyplot as plt

# ─────────────────────────────────────────────────────────────────────────────
#  LOCAL HELPER FUNCTIONS
# ─────────────────────────────────────────────────────────────────────────────

def x_from_qR(q, R, M, G, c):
    """x(q, R) from the Grundgleichung."""
    return (2*q - 1) - 0.5 * (1 - q * (1 - M*q*G / (R*c**2))**2)


def R_q_formula(q, x, M, G, c):
    """
    R_q from the quadratic root of the energy equation.
    R_q = [2*c^2*M*q^2*G + 2*sqrt(disc)] / [2*(5*c^4*q - 3*c^4 - 2*c^4*x)]
    Returns NaN when the discriminant is negative or denominator is zero.
    """
    denom = 2 * c**4 * (5*q - 3 - 2*x)
    disc  = M**2 * q**3 * G**2 * c**4 * (-4*q + 3 + 2*x)
    if disc < 0 or abs(denom) < 1e-300:
        return float("nan")
    return (2*c**2*M*q**2*G + 2*math.sqrt(disc)) / denom


def t_direct(a, q, x, H0_si, Omega_b, q0):
    """
    Direct t(a) from the asinh formula:
      t = 2/(3*H0*sqrt(x*Omega_b)) * asinh(a^(3/2) * sqrt(x/q0))
    """
    xOb = x * Omega_b
    if xOb <= 0 or x <= 0:
        return float("nan")
    return (2.0 / (3.0 * H0_si * math.sqrt(xOb))) * math.asinh(a**1.5 * math.sqrt(x / q0))


def t_of_a(a, q0, R0, M, G, c, H0_si, Omega_b):
    """
    Full t(a) with iterative q correction (deep-past regime).
    Transcription of the while-loop algorithm from Mathcad - MatLab_S.2.pdf.
    """
    R = a * R0
    q = q0
    x = x_from_qR(q, R, M, G, c)

    t_val = t_direct(a, q, x, H0_si, Omega_b, q0)
    if math.isnan(t_val):
        return float("nan")

    R_int = c * math.sqrt(max(1 - 1/q**2, 0)) * t_val
    R_q   = R_q_formula(q, x, M, G, c)

    if math.isnan(R_q) or R_q < R_int:
        return t_val   # normal regime – no iteration needed

    # Outer loop: while R_q >= R_int
    for _ in range(500):
        if math.isnan(R_q) or R_q < R_int:
            break

        q -= 1e-6
        if q <= 1.001:
            q = 1.001
            break

        n = 0
        while n <= 4:
            v = math.sqrt(max(1 - 1/q**2, 0))

            # (a) Newton-Raphson for t
            dt = max(abs(t_val) * 1e-7, 1.0)

            Ri_a = c * v * t_val
            xa   = x_from_qR(q, Ri_a, M, G, c)
            Rqa  = R_q_formula(q, xa, M, G, c)
            if math.isnan(Rqa):
                break
            fa = Rqa - Ri_a

            Ri_b = c * v * (t_val + dt)
            xb   = x_from_qR(q, Ri_b, M, G, c)
            Rqb  = R_q_formula(q, xb, M, G, c)
            if math.isnan(Rqb):
                break
            fb = Rqb - Ri_b

            df_dt = (fb - fa) / dt
            if abs(df_dt) > 1e-300:
                t_val -= fa / df_dt

            # (b) R_int <- c * sqrt(1-1/q^2) * t
            v     = math.sqrt(max(1 - 1/q**2, 0))
            R_int = c * v * t_val

            # (c) x <- energy equation evaluated at R_int
            x = x_from_qR(q, R_int, M, G, c)

            # (d) Newton-Raphson for q
            dq = max(abs(q) * 1e-7, 1e-10)

            v1  = math.sqrt(max(1 - 1/q**2, 0))
            x1  = x_from_qR(q, R_int, M, G, c)
            Rq1 = R_q_formula(q, x1, M, G, c)
            if math.isnan(Rq1):
                break
            fq1 = Rq1 - c * v1 * t_val

            q2  = q + dq
            v2  = math.sqrt(max(1 - 1/q2**2, 0))
            x2  = x_from_qR(q2, R_int, M, G, c)
            Rq2 = R_q_formula(q2, x2, M, G, c)
            if math.isnan(Rq2):
                break
            fq2 = Rq2 - c * v2 * t_val

            df_dq = (fq2 - fq1) / dq
            if abs(df_dq) > 1e-300:
                q -= fq1 / df_dq
            if q <= 1.001:
                q = 1.001
                break

            # (e) R_q <- updated with new q and x
            x   = x_from_qR(q, R_int, M, G, c)
            R_q = R_q_formula(q, x, M, G, c)
            if math.isnan(R_q):
                break

            n += 1

        # Refresh for outer loop condition
        v     = math.sqrt(max(1 - 1/q**2, 0))
        R_int = c * v * t_val
        x     = x_from_qR(q, R_int, M, G, c)
        R_q   = R_q_formula(q, x, M, G, c)

    return t_val


def t_a_Rq(a, q0, R0, M, G, c):
    """
    t(a) = R_q(q0, x) / v   [Schritt 2, page 2 – direct, non-iterative]
    """
    R  = a * R0
    x  = x_from_qR(q0, R, M, G, c)
    Rq = R_q_formula(q0, x, M, G, c)
    if math.isnan(Rq):
        return float("nan")
    v = c * math.sqrt(1 - 1/q0**2)
    return Rq / v


def da_dt(a, x_a, t_a, q0, H0_si, Omega_b):
    """
    da/dt: time derivative of the scale factor.
    From a(t) = (sqrt(q0/x) * sinh(3/2 * H0 * t * sqrt(x*Omega_b)))^(2/3)
    """
    xOb  = x_a * Omega_b
    arg  = 1.5 * H0_si * t_a * math.sqrt(xOb)
    adot = ((2/3) * a**(-0.5) * math.sqrt(q0 / x_a)
            * math.cosh(arg) * (1.5 * H0_si * math.sqrt(xOb)))
    return adot


# ─────────────────────────────────────────────────────────────────────────────
#  MAIN
# ─────────────────────────────────────────────────────────────────────────────

def main():
    import os
    out_dir = os.path.dirname(os.path.abspath(__file__))

    # ── SECTION 1: Fundamental Constants ─────────────────────────────────────
    Mpc     = 3.085677581e22        # [m]  1 Megaparsec
    G       = 6.67430e-11           # [m^3/(kg*s^2)]
    c       = 299792458.0           # [m/s]
    H0_si   = 67.4e3 / Mpc         # [1/s]
    H0_kms  = 67.4                  # [km/s/Mpc] for display

    Omega_b   = 0.04960
    Omega_m0S = 0.31882             # starting estimate
    Omega_CDM = 0.315               # ΛCDM matter density (comparison)

    rho_c = 3 * H0_si**2 / (8 * math.pi * G)   # [kg/m^3] critical density
    rho_b = Omega_b * rho_c                     # [kg/m^3] baryon density

    # Preliminary
    q_prelim = Omega_m0S / Omega_b
    x_prelim = (1 - q_prelim * Omega_b) / Omega_b

    print("\n=== Section 1: Fundamental Constants ===")
    print(f"H0        = {H0_kms:.4f} km/s/Mpc  = {H0_si:.6e} 1/s")
    print(f"rho_c     = {rho_c:.6e} kg/m^3")
    print(f"rho_b     = {rho_b:.6e} kg/m^3")

    # ── SECTION 2: Step 1 output ──────────────────────────────────────────────
    q0       = 6.427606904079293
    Omega_m0 = q0 * Omega_b

    print("\n=== Section 2: Step 1 Output ===")
    print(f"q0        = {q0:.15f}")
    print(f"Omega_m0  = {Omega_m0:.15f}")
    print( "  (target:  0.318809302442333)")

    # ── SECTION 3 & 4: R0 self-consistency ───────────────────────────────────
    R0 = 1.285e26   # [m] initial estimate

    print("\n=== Section 3 & 4: R0 self-consistency ===")

    for k in range(1, 51):
        M_k    = (4/3) * math.pi * R0**3 * rho_b
        x_k    = x_from_qR(q0, R0, M_k, G, c)
        t1_k   = t_direct(1.0, q0, x_k, H0_si, Omega_b, q0)
        R0_new = c * math.sqrt(1 - 1/q0**2) * t1_k
        err    = abs(R0_new - R0) / R0
        R0     = R0_new
        if err < 1e-9:
            print(f"  R0 converged after {k} iterations")
            break

    M      = (4/3) * math.pi * R0**3 * rho_b
    MG_c2  = M * G / c**2
    Gyr    = 1e9 * 365.25 * 24 * 3600   # [s]
    t1_val = t_of_a(1.0, q0, R0, M, G, c, H0_si, Omega_b)

    print(f"R0 (self-consistent)  = {R0:.6e} m  (target: 1.285e26 m)")
    print(f"M                     = {M:.6e} kg")
    print(f"t(a=1)                = {t1_val/Gyr:.9f} Gyr  (target: 13.749542 Gyr)")

    # ── SCHRITT 2 ─────────────────────────────────────────────────────────────
    t0_s2 = t_a_Rq(1.0, q0, R0, M, G, c)

    print("\n=== Schritt 2: Unvermeidliche Vorberechnungen ===")
    print(f"t0  = {t0_s2:.3e} s   (target: 4.339e17 s)")
    print(f"R0  = {R0:.8e} m   (target: 1.28496778e26 m)")
    print(f"M   = {M:.3e} kg  (target: 3.761e51 kg)")

    print("\n=== Schritt 2: Bestimmung t(a), x(a) fuer a = 1, 0.9 .. 0.1 ===")
    print(f"{'a':>5}  {'t [s]':>14}  {'t [Gyr]':>10}  {'x(a)':>12}")
    a_range = [round(1.0 - 0.1*k, 10) for k in range(10)]
    for a_i in a_range:
        t_i = t_a_Rq(a_i, q0, R0, M, G, c)
        x_i = x_from_qR(q0, a_i * R0, M, G, c)
        print(f"{a_i:5.2f}  {t_i:14.6e}  {t_i/Gyr:10.6f}  {x_i:12.6f}")

    # ── SECTION 5: Output functions  (a = 2 .. 0.06) ─────────────────────────
    a_vec = np.flip(np.arange(0.06, 2.005, 0.005))
    N     = len(a_vec)

    R_vec        = np.zeros(N)
    x_vec        = np.zeros(N)
    OmLambda_vec = np.zeros(N)
    OmM_vec      = np.zeros(N)

    for i, a_i in enumerate(a_vec):
        R_vec[i]        = a_i * R0
        x_vec[i]        = x_from_qR(q0, R_vec[i], M, G, c)
        OmLambda_vec[i] = x_vec[i] * Omega_b
        OmM_vec[i]      = q0 * Omega_b   # constant

    idx1  = int(np.argmin(np.abs(a_vec - 1.0)))
    idx05 = int(np.argmin(np.abs(a_vec - 0.5)))

    x_001   = x_from_qR(q0, 0.01 * R0, M, G, c)
    OmL_001 = x_001 * Omega_b

    print("\n=== Section 5: Output checks ===")
    print(f"Omega_Lambda(a=1)    = {OmLambda_vec[idx1]:.10f}  (target: 0.6811906976)")
    print(f"Omega_Lambda(a=0.01) = {OmL_001:.3f}  (target: 27.387)")
    print( "Singularity (doc.)   = 0.0740  (numerical observation)")
    print(f"R(a=1)               = {R_vec[idx1]:.4e} m  (target: 1.285e26 m)")

    # ── SECTION 6: Hubble parameter functions ─────────────────────────────────
    Hg_vec   = np.zeros(N)
    H1_vec   = np.full(N, float("nan"))
    H2_vec   = np.full(N, float("nan"))
    HCDM_vec = np.zeros(N)
    t_vec    = np.full(N, float("nan"))

    for i, a_i in enumerate(a_vec):
        R_i = R_vec[i]
        x_i = x_vec[i]

        # Geometric Hubble
        Hg_vec[i]   = c / R_i * Mpc / 1000.0

        # ΛCDM comparison
        HCDM_vec[i] = H0_kms * math.sqrt(Omega_CDM / a_i**3 + (1 - Omega_CDM))

        if x_i > 0:
            t_i    = t_a_Rq(a_i, q0, R0, M, G, c)
            xOb    = x_i * Omega_b
            arg    = 1.5 * H0_si * t_i * math.sqrt(xOb)
            sqQX   = math.sqrt(q0 / x_i)
            adot_i = sqQX * math.cosh(arg) * H0_si * math.sqrt(xOb) / (sqQX * math.sinh(arg))**(1/3)
            H1_vec[i] = (adot_i / a_i * Mpc / 1000.0).real
            t_vec[i]  = t_i

            H2_vec[i] = (math.sqrt(8/3 * math.pi * G
                         * (q0*rho_b/a_i**3 + x_i*rho_b))
                         * Mpc / 1000.0)

    print("\n=== Section 6: Hubble parameters ===")
    print(f"H_g  (a=0.5) = {Hg_vec[idx05]:.3f} km/s/Mpc  (target: 143.982)")
    print(f"H1   (a=0.5) = {H1_vec[idx05]:.3f} km/s/Mpc  (target: 118.344)")
    print(f"H2   (a=0.5) = {H2_vec[idx05]:.3f} km/s/Mpc  (target: 120.502)")
    print(f"HCDM (a=0.5) = {HCDM_vec[idx05]:.3f} km/s/Mpc  (target: 120.663)")
    print(f"H1   (a=1.0) = {H1_vec[idx1]:.3f} km/s/Mpc  (target:  67.400)")
    print(f"H2   (a=1.0) = {H2_vec[idx1]:.3f} km/s/Mpc  (target:  67.400)")
    print(f"HCDM (a=1.0) = {HCDM_vec[idx1]:.3f} km/s/Mpc  (target:  67.400)")
    print(f"H_g  (a=1.0) = {Hg_vec[idx1]:.3f} km/s/Mpc  (target:  71.991)")

    # ── SECTION 7: Tabulated output ───────────────────────────────────────────
    print("\n=== Section 7: Table ===")
    print(f"{'a':>6}  {'R [m]':>11}  {'x(a)':>9}  {'OmLambda':>9}  "
          f"{'H_g':>8}  {'H1':>8}  {'H2':>8}  {'H_CDM':>8}")

    for ap in [2.0, 1.5, 1.0, 0.75, 0.5, 0.25, 0.1, 0.06]:
        ii = int(np.argmin(np.abs(a_vec - ap)))
        print(f"{a_vec[ii]:6.3f}  {R_vec[ii]:11.4e}  {x_vec[ii]:9.5f}  "
              f"{OmLambda_vec[ii]:9.5f}  {Hg_vec[ii]:8.3f}  "
              f"{H1_vec[ii]:8.3f}  {H2_vec[ii]:8.3f}  {HCDM_vec[ii]:8.3f}")

    # ── SECTION 8: Plots ──────────────────────────────────────────────────────

    # Figure 1: Hubble parameters
    fig1, ax1 = plt.subplots(figsize=(9, 5))
    ax1.plot(a_vec, Hg_vec,   "k-",  lw=1.5, label=r"$H_g$ (geometric)")
    ax1.plot(a_vec, H1_vec,   "b--", lw=1.5, label=r"$H_1$ (da/dt)")
    ax1.plot(a_vec, H2_vec,   "r-",  lw=1.5, label=r"$H_2$ (Friedmann)")
    ax1.plot(a_vec, HCDM_vec, "m-.", lw=1.5, label=r"$H_\mathrm{CDM}$ ($\Omega_\mathrm{CDM}=0.315$)")
    ax1.plot(1.0, H0_kms, "ko", ms=8, label=f"$H_0 = {H0_kms:.1f}$ km/s/Mpc")
    ax1.set_xlabel("Scale factor a")
    ax1.set_ylabel("H [km/s/Mpc]")
    ax1.set_title("Hubble parameter functions (Daniel's cosmological model)")
    ax1.set_xlim(0.06, 2.0)
    ax1.set_ylim(0, 500)
    ax1.grid(True)
    ax1.legend(loc="upper right")
    fig1.tight_layout()
    fig1.savefig(os.path.join(out_dir, "RFA_fig1_Hubble.pdf"), dpi=300)
    print(f"\nFigure saved: {os.path.join(out_dir, 'RFA_fig1_Hubble.pdf')}")

    # Figure 2: Density parameters
    fig2, ax2 = plt.subplots(figsize=(9, 5))
    ax2.plot(a_vec, OmLambda_vec, "r-",  lw=1.5, label=r"$\Omega_\Lambda(a)$")
    ax2.plot(a_vec, OmM_vec,      "b--", lw=1.5, label=r"$\Omega_m$ (const)")
    ax2.axhline(0.68119,  color="r", ls=":", lw=1.0)
    ax2.axhline(Omega_m0, color="b", ls=":", lw=1.0)
    ax2.set_xlabel("Scale factor a")
    ax2.set_ylabel(r"$\Omega$")
    ax2.set_title("Density parameters vs scale factor")
    ax2.set_xlim(0.06, 2.0)
    ax2.grid(True)
    ax2.legend(loc="upper right")
    fig2.tight_layout()
    fig2.savefig(os.path.join(out_dir, "RFA_fig2_Omega.pdf"), dpi=300)
    print(f"Figure saved: {os.path.join(out_dir, 'RFA_fig2_Omega.pdf')}")

    # Figure 3: Age of universe
    fig3, ax3 = plt.subplots(figsize=(9, 5))
    ax3.plot(a_vec, t_vec / Gyr, "b-", lw=1.5, label="t(a)")
    ax3.plot(1.0, t1_val / Gyr, "ro", ms=8,
             label=f"t(1) = {t1_val/Gyr:.3f} Gyr")
    ax3.set_xlabel("Scale factor a")
    ax3.set_ylabel("t [Gyr]")
    ax3.set_title("Age of universe t(a)")
    ax3.set_xlim(0.06, 2.0)
    ax3.grid(True)
    ax3.legend(loc="upper left")
    fig3.tight_layout()
    fig3.savefig(os.path.join(out_dir, "RFA_fig3_Age.pdf"), dpi=300)
    print(f"Figure saved: {os.path.join(out_dir, 'RFA_fig3_Age.pdf')}")

    # ── SECTION 9: Summary ────────────────────────────────────────────────────
    print("\n=== Summary ===")
    print(f"q0              = {q0:.15f}")
    print(f"Omega_m0        = {Omega_m0:.15f}")
    print(f"R0              = {R0:.8e} m")
    print(f"M               = {M:.6e} kg")
    print(f"t(a=1)          = {t1_val/Gyr:.6f} Gyr  (target: 13.749542)")
    print(f"H1(a=1)         = {H1_vec[idx1]:.3f} km/s/Mpc  (target: 67.4)")
    print(f"H2(a=1)         = {H2_vec[idx1]:.3f} km/s/Mpc  (target: 67.4)")
    print(f"HCDM(a=0.5)     = {HCDM_vec[idx05]:.3f} km/s/Mpc  (target: 120.663)")
    print(f"Omega_Lambda(1) = {OmLambda_vec[idx1]:.10f}  (target: 0.6811906976)")
    print( "Singularity     = 0.0740  (numerical; no closed-form formula in document)")
    print("Done.")


if __name__ == "__main__":
    main()