"""Independent checks of the pieces added or changed in the revision.  Run by run_all.py (stage 9) and
writes data/verification.csv; every row has a pass/fail flag.

 1. Earth-twin lifetimes: the fixed-step Heun integrator of thermal.py against scipy's adaptive LSODA,
    written here from the equations (not through thermal.integrate).
 2. Thermostat: a hot, old-epoch planet under the flat scaling, integrated with LSODA with the limit
    imposed as "no further heating at T_cap", against thermal.integrate.
 3. Mixed-criterion ranks: the weighted-sample rank against the same population built by replication.
 4. First-pass regression: with decay-model strength 1 and no thermostat the revised code must give the
    first-pass ranks (F2h, degassing criterion: 0.47 for 4 steps, 0.29 for 6).
 5. Fixed-cap ranks of the sister study (delta-function lifetimes).
"""
import os

import numpy as np
import pandas as pd
from scipy.integrate import solve_ivp

import analysis as an
import population as pop
import rank
import thermal
from common import DATA_DIR, AGE_SUN


def _lsoda_lifetime(model, fac, mass=1.0, eta=1.0, T0=None, t_max=40.0, f_crit=0.1, cap=None):
    """Lifetime (ridge law) and end of melting from an adaptive integration of C dT/dt = H - Q."""
    p = {"fac": np.atleast_2d(fac).astype(float), "mass": np.array([mass]), "eta_fac": np.array([eta])}
    model.prepare(p)

    def rhs(t, y):
        d = float(model.dTdt(t, np.array([y[0]]), p)[0])
        if cap is not None and y[0] >= cap and d > 0.0:
            return [0.0]
        return [d]

    T0 = model.T_init if T0 is None else T0
    if cap is not None:
        T0 = min(T0, cap)
    sol = solve_ivp(rhs, (0.0, t_max), [T0], method="LSODA", rtol=1e-8, atol=1e-6, dense_output=True, max_step=0.02)
    t = np.linspace(0.0, t_max, 40001)
    T = sol.sol(t)[0]
    D = model.degassing_rel(T, {k: (np.repeat(v, len(t), axis=0) if k == "fac" else np.repeat(v, len(t))) for k, v in
                                (("fac", p["fac"]), ("mass", p["mass"]), ("eta_fac", p["eta_fac"]))})
    life = t[D >= f_crit].max() if (D >= f_crit).any() else 0.0
    melt = t[D > 0].max() if (D > 0).any() else 0.0
    return float(life), float(melt), float(T.max())


def run(emp, n=600, seed=11):
    rows = []

    def add(check, value, reference, tol, note=""):
        rows.append({"check": check, "value": value, "reference": reference, "tolerance": tol,
                     "pass": bool(abs(value - reference) <= tol), "note": note})

    # 1. Earth twin
    for name in ("F1_classical", "F2h_hybrid", "F2_flat"):
        m = thermal.make_model(name)
        s = m.earth_summary()
        life, melt, _ = _lsoda_lifetime(m, np.ones(4), t_max=40.0)
        add(f"Earth twin lifetime, ridge law, {name}: LSODA vs Heun", life, s["lifetime_Gyr"], 0.03, "Gyr")
        add(f"Earth twin end of melting, {name}: LSODA vs Heun", melt, s["melt_end_Gyr"], 0.03, "Gyr")

    # 2. thermostat on a hot old-epoch planet (full decay model, 12 Gyr node, 1.5 Earth masses)
    h = an.base_hyper(emp, ism_scale=1.0)
    z = {k: np.zeros(1) for k in ("z_th", "z_k", "z_vol", "z_eta", "z_T", "u_m")}
    fac = pop.hpe_factors(12.0, z, h, vary=())
    for name in ("F2_flat", "F2h_hybrid"):
        for cap in (2200.0, None):
            m = thermal.make_model(name, T_cap=cap)
            p = {"fac": fac.copy(), "mass": np.array([1.5]), "eta_fac": np.ones(1)}
            res = thermal.integrate(m, np.array([m.T_init]), p, 40.0, dt=0.02, dt_early=0.005, dt_late=0.05)
            life, melt, tpk = _lsoda_lifetime(m, fac, mass=1.5, t_max=40.0, cap=cap)
            add(f"hot old-epoch planet, {name}, thermostat {cap}: lifetime LSODA vs Heun", life, float(res["lifetime"][0]), 0.06, "Gyr")
            add(f"hot old-epoch planet, {name}, thermostat {cap}: peak temperature LSODA vs Heun", tpk, float(res["T_peak"][0]), 3.0, "K")

    # 3. weighted mixture vs replication
    zz = pop.draw_normals(n, seed)
    r = pop.run_population("F1_classical", an.base_hyper(emp), zz)
    mix_w = np.concatenate([r["lifetime"], r["melt_end"]], axis=1)
    w = np.concatenate([np.full(n, 0.9), np.full(n, 0.1)])
    mix_r = np.concatenate([np.tile(r["lifetime"], (1, 9)), r["melt_end"]], axis=1)
    for n_ in (4, 6):
        a = rank.rank_from_lifetimes(mix_w, pop.NODE_AGES, n_, weights=w)
        b = rank.rank_from_lifetimes(mix_r, pop.NODE_AGES, n_)
        add(f"10% melting-criterion mixture, n = {n_}: weighted vs replicated samples", a, b, 1e-6 + 1e-6 * b, "rank")

    # 4. first-pass regression (F2h, degassing): 0.47 (n = 4) and 0.29 (n = 6) in the first-pass report
    r = pop.run_population("F2h_hybrid", an.base_hyper(emp, ism_scale=1.0, T_cap=None), pop.draw_normals(2000, seed + 1))
    add("first-pass configuration, F2h, degassing, n = 4", rank.rank_from_lifetimes(r["lifetime"], pop.NODE_AGES, 4), 0.468, 0.03,
        "rank; first-pass value 0.468 with 5000 planets")
    add("first-pass configuration, F2h, degassing, n = 6", rank.rank_from_lifetimes(r["lifetime"], pop.NODE_AGES, 6), 0.294, 0.03,
        "rank; first-pass value 0.294 with 5000 planets")

    # 5. fixed caps of the sister study
    for n_, cap, ref in ((4, 10.0, 0.52451), (6, 20.0, 0.01299)):
        a = rank.rank_from_lifetimes(np.full((2, 50), cap), [0.0, 12.0], n_, b_step=1)
        add(f"delta-function lifetime {cap:g} Gyr, n = {n_}: against mdwarf_early", a, ref, 1e-4, "rank")

    out = pd.DataFrame(rows)
    out.to_csv(os.path.join(DATA_DIR, "verification.csv"), index=False)
    return out


if __name__ == "__main__":
    import json
    emp = json.load(open(os.path.join(DATA_DIR, "results_summary.json")))["abundance_hyper"]
    with pd.option_context("display.width", 250, "display.max_colwidth", 90):
        print(run(emp).round(5).to_string())