#!/usr/bin/env python3
"""Reproduce the dimensionless stellar examples in blackhole.org Volume 12.

Run: python3 verify-stellar-teaching-models.py > stellar-teaching-results.json
Python 3.9+ standard library only; no downloaded data, file writes, or randomness.
The JSON records the actual Python/platform environment and convergence results.

Physics: cold, isotropic, nonrotating GR perfect fluid, p=K*rho0**2,
epsilon=rho0+p, with rho0 the baryon rest-mass density and G=c=1 in the
dimensionless equations. q=K*rho0; rbar=r*sqrt(G/K);
mbar=m*G**1.5/sqrt(K). The two TOV integrations use independent radial
and enthalpy coordinates and explicit classical RK4. Neither computes radial
eigenmodes nor certifies stability, a realistic EOS, or a rotating-star model.
The Lane-Emden n=3 problem independently gives the ideal white-dwarf constant.

Equations are derived on the Mass-Radius Sequences and Chandrasekhar Mass
pages. References: Oppenheimer & Volkoff, Phys.Rev.55:374 (1939), Eqs.9-10;
Chandrasekhar, ApJ74:81 (1931), Eqs.2-5. All code and computed values are
original to blackhole.org; no source table or third-party curve is copied.
Failures of declared numerical checks raise an exception and exit nonzero.
"""

import json
import math
import platform
import sys


def rk4(fun, x, y, step):
    """One classical fourth-order Runge-Kutta step."""
    k1 = fun(x, y)
    k2 = fun(x + step / 2, [v + step * k / 2 for v, k in zip(y, k1)])
    k3 = fun(x + step / 2, [v + step * k / 2 for v, k in zip(y, k2)])
    k4 = fun(x + step, [v + step * k for v, k in zip(y, k3)])
    return [v + step * (a + 2*b + 2*c + d) / 6
            for v, a, b, c, d in zip(y, k1, k2, k3, k4)]


def surface_step(fun, x, y, step, component):
    """Bisect the last RK4 step to locate a zero, rather than truncate a grid."""
    low, high = 0.0, step
    for _ in range(50):
        mid = (low + high) / 2
        trial = rk4(fun, x, y, mid)
        if trial[component] > 0:
            low = mid
        else:
            high = mid
    delta = (low + high) / 2
    return x + delta, rk4(fun, x, y, delta)


def state_from_h(h):
    # Negative-H RK stages occur only in a trial step bracketing the surface.
    # They use q=p=epsilon=0; accepted output is root-located at H=0, so no
    # exterior atmosphere is added to the stellar mass or radius. Refinement
    # and the independent H-coordinate solve test this event treatment.
    q = max(math.expm1(h) / 2, 0.0)
    return q, q*q, q + q*q


def tov_radius(qc, resolution, center_factor=1.0):
    """Integrate m(r), H(r), and baryon mass on a uniform radial grid."""
    step = 1 / resolution
    r = step * center_factor
    ec, pc = qc + qc*qc, qc*qc
    hc = math.log1p(2*qc)
    h2 = -2*math.pi*(ec/3 + pc)
    q2 = (1 + 2*qc)*h2/2
    e2 = (1 + 2*qc)*q2
    h = hc + h2*r*r
    y = [4*math.pi*(ec*r**3/3 + e2*r**5/5), h,
         4*math.pi*(qc*r**3/3 + (q2 + 4*math.pi*qc*ec/3)*r**5/5)]
    maximum_twom_over_r = 0.0

    def rhs(radius, values):
        mass, enthalpy, _ = values
        q, pressure, energy = state_from_h(enthalpy)
        factor = 1 - 2*mass/radius
        if factor <= 0:
            raise ArithmeticError("Radial integration entered trapped region")
        return [4*math.pi*radius**2*energy,
                -(mass + 4*math.pi*radius**3*pressure)/(radius**2*factor),
                4*math.pi*radius**2*q/math.sqrt(factor)]

    for _ in range(4*resolution):
        trial = rk4(rhs, r, y, step)
        if trial[1] <= 0:
            radius, final = surface_step(rhs, r, y, step, 1)
            return {"q_c": qc, "M": final[0], "R": radius, "M_b": final[2],
                    "compactness": final[0]/radius, "surface_H": final[1],
                    "max_2m_over_r": max(maximum_twom_over_r, 2*final[0]/radius)}
        r += step
        y = trial
        maximum_twom_over_r = max(maximum_twom_over_r, 2*y[0]/r)
    raise ArithmeticError("No TOV surface found before rbar=4")


def tov_enthalpy(qc, resolution):
    """Independently integrate x=r^2, z=m/r, M_b using H as coordinate."""
    hc = math.log1p(2*qc)
    ec, pc = qc + qc*qc, qc*qc
    delta = hc * 1e-5 / resolution
    h = hc - delta
    x = delta / (2*math.pi*(ec/3 + pc))
    z = 4*math.pi*ec*x/3
    y = [x, z, 4*math.pi*qc*x**1.5/3]
    step = -h / resolution

    def rhs(enthalpy, values):
        radius2, ratio, _ = values
        q, pressure, energy = state_from_h(enthalpy)
        factor = 1 - 2*ratio
        denom = ratio + 4*math.pi*radius2*pressure
        if radius2 <= 0 or factor <= 0 or denom <= 0:
            raise ArithmeticError("Invalid enthalpy-coordinate TOV state")
        return [-2*radius2*factor/denom,
                -(4*math.pi*radius2*energy-ratio)*factor/denom,
                -4*math.pi*radius2**1.5*q*math.sqrt(factor)/denom]

    for i in range(resolution):
        y = rk4(rhs, h + i*step, y, step)
    radius = math.sqrt(y[0])
    return {"q_c": qc, "M": y[1]*radius, "R": radius,
            "M_b": y[2], "compactness": y[1]}


def lane_emden(index, resolution):
    """Newtonian polytrope; evolve theta and v=xi^2 theta' to first zero."""
    step = 1 / resolution
    x = step
    y = [1-x*x/6+index*x**4/120, -x**3/3+index*x**5/30]

    def rhs(radius, values):
        theta, flux = values
        return [flux/radius**2, -radius**2*theta**index]

    for _ in range(10*resolution):
        trial = rk4(rhs, x, y, step)
        if trial[0] <= 0:
            root, final = surface_step(rhs, x, y, step, 0)
            return {"index": index, "xi_1": root, "minus_xi2_theta_prime": -final[1]}
        x += step
        y = trial
    raise ArithmeticError("No Lane-Emden surface before xi=10")


def relative(a, b):
    return abs(a-b)/max(abs(a), abs(b), 1e-30)


def find_maximum(resolution):
    # Bracket belongs to this teaching EOS only. Search log(qc) to avoid
    # using a discrete sampled mass maximum as the continuum extremum.
    # Golden-section search assumes a single peak in this bracket; this is
    # checked on an auxiliary grid below, not asserted for arbitrary EOSs.
    low, high = math.log(0.1), math.log(0.8)
    golden = (math.sqrt(5)-1)/2
    left = high - golden*(high-low)
    right = low + golden*(high-low)
    a = tov_enthalpy(math.exp(left), resolution)
    b = tov_enthalpy(math.exp(right), resolution)
    for _ in range(42):
        if a["M"] < b["M"]:
            low, left, a = left, right, b
            right = low + golden*(high-low)
            b = tov_enthalpy(math.exp(right), resolution)
        else:
            high, right, b = right, left, a
            left = high - golden*(high-low)
            a = tov_enthalpy(math.exp(left), resolution)
    return tov_enthalpy(math.exp((low+high)/2), resolution)


def main():
    checks = []

    def check(name, value, bound):
        if not math.isfinite(value) or value > bound:
            raise AssertionError(f"{name}: {value:.8g} exceeds {bound:.8g}")
        checks.append({"check": name, "value": value, "upper_bound": bound})

    central = [0.01, 0.05, 0.10, 0.20, 0.32, 0.50, 1.00]
    convergence = []
    sequence = []
    for qc in central:
        radial = [tov_radius(qc, n) for n in (1024, 2048, 4096)]
        enthalpy = [tov_enthalpy(qc, n) for n in (1024, 2048, 4096)]
        for field in ("M", "R", "M_b"):
            check(f"radial refinement qc={qc} {field}",
                  relative(radial[-1][field], radial[-2][field]), 2e-6)
            check(f"enthalpy refinement qc={qc} {field}",
                  relative(enthalpy[-1][field], enthalpy[-2][field]), 2e-6)
            check(f"independent coordinates qc={qc} {field}",
                  relative(radial[-1][field], enthalpy[-1][field]), 2e-6)
        check(f"surface enthalpy qc={qc}", abs(radial[-1]["surface_H"]), 1e-12)
        if not 0 < radial[-1]["max_2m_over_r"] < 1:
            raise AssertionError("Nonpositive compactness or trapped surface")
        check(f"causal sound speed qc={qc}", 2*qc/(1+2*qc), 1.0)
        convergence.append({"q_c": qc, "resolutions": [1024, 2048, 4096],
                            "radial": radial, "enthalpy": enthalpy})
        sequence.append(radial[-1])

    centered = tov_radius(0.2, 4096, center_factor=0.5)
    check("center-start refinement", max(relative(centered[k], sequence[3][k])
          for k in ("M", "R", "M_b")), 1e-7)

    newtonian = []
    for qc in (1e-3, 1e-4, 1e-5):
        model = tov_enthalpy(qc, 4096)
        newtonian.append({"q_c": qc, "R": model["R"], "M_over_qc": model["M"]/qc,
                          "R_error": relative(model["R"], math.sqrt(math.pi/2)),
                          "M_over_qc_error": relative(model["M"]/qc, math.sqrt(2*math.pi))})
    check("Newtonian n=1 radius limit", newtonian[-1]["R_error"], 5e-5)
    check("Newtonian n=1 mass limit", newtonian[-1]["M_over_qc_error"], 1e-4)
    if any(newtonian[i+1]["R_error"] >= newtonian[i]["R_error"] for i in (0, 1)):
        raise AssertionError("No approach to analytic Newtonian limit")

    # Cold first law dM/dMb=mu_infinity/mb=sqrt(1-2M/R), since the
    # surface chemical potential of this polytrope equals mb.
    first_law = []
    for delta in (1e-3, 5e-4, 2.5e-4):
        left = tov_enthalpy(0.1-delta, 8192)
        right = tov_enthalpy(0.1+delta, 8192)
        ratio = (right["M"]-left["M"])/(right["M_b"]-left["M_b"])
        central_model = tov_enthalpy(0.1, 8192)
        expected = math.sqrt(1-2*central_model["compactness"])
        first_law.append({"delta_qc": delta, "dM_dMb": ratio,
                          "surface_redshift_factor": expected,
                          "relative_error": relative(ratio, expected)})
    check("baryon first law", first_law[-1]["relative_error"], 2e-6)

    peak_grid = [tov_enthalpy(0.1 * 8**(i/32), 512) for i in range(33)]
    peak_index = max(range(len(peak_grid)), key=lambda i: peak_grid[i]["M"])
    if not 0 < peak_index < len(peak_grid)-1:
        raise AssertionError("Sampled peak is not inside the search bracket")
    if any(peak_grid[i+1]["M"] <= peak_grid[i]["M"] for i in range(peak_index)):
        raise AssertionError("Multiple sampled extrema before bracketed peak")
    if any(peak_grid[i+1]["M"] >= peak_grid[i]["M"] for i in range(peak_index, 32)):
        raise AssertionError("Multiple sampled extrema after bracketed peak")
    maxima = [find_maximum(n) for n in (1024, 2048, 4096)]
    for field in ("M", "R", "q_c"):
        check(f"extremum refinement {field}",
              relative(maxima[-1][field], maxima[-2][field]), 2e-5)
    for factor in (0.95, 1.05):
        if tov_enthalpy(maxima[-1]["q_c"]*factor, 4096)["M"] >= maxima[-1]["M"]:
            raise AssertionError("Mass extremum is not a local maximum")

    lanes = [lane_emden(3, n) for n in (128, 256, 512)]
    check("Lane-Emden n=3 radius convergence",
          relative(lanes[-1]["xi_1"], lanes[-2]["xi_1"]), 1e-7)
    check("Lane-Emden n=3 mass convergence",
          relative(lanes[-1]["minus_xi2_theta_prime"], lanes[-2]["minus_xi2_theta_prime"]), 1e-7)
    analytic = lane_emden(1, 512)
    check("Lane-Emden n=1 analytic radius", relative(analytic["xi_1"], math.pi), 1e-7)
    check("Lane-Emden n=1 analytic mass",
          relative(analytic["minus_xi2_theta_prime"], math.pi), 1e-7)
    # Explicit dimensional constants for the rounded textbook coefficient,
    # not a precision determination of Newton's constant.
    hbar, speed, grav, mu, solar = 1.054571817e-34, 299792458, 6.67430e-11, 1.66053906660e-27, 1.98847e30
    mass_constant = (math.sqrt(3*math.pi)/2 * lanes[-1]["minus_xi2_theta_prime"]
                     * (hbar*speed/grav)**1.5 / (mu*mu*solar))
    check("Chandrasekhar rounded coefficient", abs(mass_constant-5.83), 0.01)
    result = {
        "schema": "blackhole-stellar-teaching-models-v1",
        "environment": {"python": sys.version, "platform": platform.platform(),
                        "dependencies": "Python standard library only"},
        "scope": "Teaching EOS and Newtonian Lane-Emden only; no radial-mode stability calculation",
        "units": {"q": "K*rho0", "rbar": "r*sqrt(G/K)", "mbar": "m*G**1.5/sqrt(K)",
                  "rho0": "baryon reference rest-mass density", "c": 1},
        "method": {
            "integrator": "classical fixed-step RK4 in radius and separately in enthalpy",
            "radius": "step=1/resolution; center start r=step with m,Mb through r^5 and H through r^2",
            "enthalpy": "start Hc-H=Hc*1e-5/resolution with leading regular center series",
            "surface": "H=0; radial final step bisected 50 times; enthalpy coordinate ends at zero",
            "error_scope": "reported refinement and cross-method differences, not rigorous error bounds",
            "peak": "golden-section search in log(qc), qc in [0.1,0.8]; sampled unimodality check"
        },
        "sequence": sequence, "convergence": convergence, "newtonian_limit": newtonian,
        "baryon_first_law": first_law, "mass_extrema": maxima,
        "mass_peak_bracket_samples": peak_grid,
        "lane_emden_n3": lanes, "lane_emden_n1_control": analytic,
        "chandrasekhar_coefficient_Msun_times_mu_e_squared": mass_constant,
        "checks": checks, "all_checks_passed": True
    }
    print(json.dumps(result, indent=2, sort_keys=True, allow_nan=False))


if __name__ == "__main__":
    main()