"""Part 3: Monte Carlo over planets.  Data-driven heat budgets + priors -> habitable lifetimes.

Per-planet random inputs
  EMPIRICAL (from abundances.py):
    [r/rock]  ~ N(a + b (age - 4.57), s)   zero point, RAW trend with stellar age and intrinsic scatter
                                            of the r-process/rock ratio (Eu data; zero point also Th)
    [K/rock]  trend and scatter (GALAH)
  INFERRED / ASSUMED:
    Th, U     f_X = f_r x (ISM free-decay factor of X)^ism_scale.  The one-zone free-decay model is a
              MODEL; ism_scale is its strength: 0 = data only (Th and U track Eu), 1 = full uniform-
              production model, baseline 0.5 because the only test available (Th/Eu at birth versus age
              in solar twins) finds 0.43 +/- 0.27 of the predicted gradient.  U was never measured.
    40K/K     ISM free-decay factor (same strength); planetary K depletion ~ lognormal(0, 0.25 dex)
    mass      log-uniform m_lo - m_hi Earth masses (baseline 0.5-2; 0.5-3 and 0.5-5 as variants)
    viscosity lognormal about the Earth-calibrated value, sigma = s_eta dex (default 0.5)
    primordial heat: initial potential temperature ~ N(T_init, 100 K), not above the thermostat T_cap

Hyper-parameters (drawn once per "hyper-draw" to propagate calibration/model uncertainty):
    a, b, s (within the spread between the abundance samples), ism_scale ~ U(0, 1), activation energy
    E ~ U(200, 400) kJ/mol, zero-pressure solidus ~ U(1359, 1423) K, continental heat-source fraction
    f_cc ~ U(0.2, 0.5), present convective heat loss Q0 ~ N(38, 3) TW, s_eta ~ U(0.25, 0.75),
    thermostat T_cap ~ U(2100, 2300) K.
    NOT drawn (treated as explicit scenarios in analysis.tail_stage): planet-mass range, degassing
    threshold, degassing law, share of planets on the melting criterion.
"""
import numpy as np

import thermal
from abundances import ism_factor
from common import AGE_SUN, T_UNIVERSE, SEED

TARGET_AGES = (5.0, 8.0, 10.0, 15.0, 20.0, 30.0, 45.0)
# formation-epoch nodes, expressed as the stellar age today [Gyr]; age 0 also stands for future births
NODE_AGES = np.array([0.0, AGE_SUN, 8.5, 12.0])
AGE_CLIP = 12.0   # the abundance trends and the ISM decay model are not extrapolated beyond 12 Gyr
T_MAX = {"F1_classical": 80.0, "F2h_hybrid": 70.0, "F2_flat": 40.0, "F3_stagnant": 20.0, "power": 70.0}
DT_EARLY = {"F1_classical": 0.005, "F2h_hybrid": 0.005, "F2_flat": 0.005, "F3_stagnant": 0.0025, "power": 0.005}

# ism_scale multiplies (in the log) the model ISM free-decay factors of Th, U and 40K: 1 = uniform-
# production one-zone model, 0 = data only (Th and U track Eu, 40K tracks K).  a, b, s describe the RAW
# europium-per-rock data (b is the measured age slope, no model factor).
ISM_SCALE_BASE = 0.5
DEFAULT_HYPER = {"ism_scale": ISM_SCALE_BASE, "a": 0.03, "b": 0.0055, "s": 0.05, "k_a": 0.02, "k_b": -0.007, "s_k": 0.06, "s_vol": 0.25,
                 "E": 300.0e3, "t_sol0": thermal.T_SOL0_MID, "f_cc": thermal.F_CC_EARTH,
                 "Q0": thermal.Q_CONV_EARTH, "s_eta": 0.5, "s_Ti": 100.0, "m_lo": 0.5, "m_hi": 2.0,
                 "T_cap": thermal.T_CAP_DEFAULT}


def hyper_from_fits(summary):
    """Adopted abundance hyper-parameters and their uncertainties from the fits of abundances.analyse().

    zero point a: mean and standard deviation of the four sample zero points (Th twins, Eu twins,
                  HARPS-GTO, Battistini & Bensby) at the solar age (the model decay factor is zero there);
    slope b:      RAW [Eu/rock] age slope, mean of the two broad (not metallicity-restricted) samples;
                  its uncertainty is taken as half the difference between the broad-sample mean and the
                  solar-twin slope, because the twins, restricted to |[Fe/H]| < 0.15, show a negative trend;
    scatter s:    mean and standard deviation of the four intrinsic scatters (error-model dependent for
                  the two broad samples; see report)."""
    fe = summary["fits_eu_raw"]
    zp = np.array([summary["fit_th_age"]["a"], fe["solar_twins"]["a"], fe["harps_gto"]["a"], fe["battistini_bensby"]["a"]])
    sc = np.array([summary["fit_th_age"]["s_int"], fe["solar_twins"]["s_int"], fe["harps_gto"]["s_int"], fe["battistini_bensby"]["s_int"]])
    b_broad = 0.5 * (fe["harps_gto"]["b"] + fe["battistini_bensby"]["b"])
    k3, k4 = summary["fits_k"]["galah_dr3"], summary["fits_k"]["galah_dr4"]
    return {"a": float(zp.mean()), "e_a": float(zp.std(ddof=1)),
            "b": float(b_broad), "e_b": float(0.5 * abs(b_broad - fe["solar_twins"]["b"])),
            "s": float(sc.mean()), "e_s": float(sc.std(ddof=1)),
            "k_a": float(0.5 * (k3["a"] + k4["a"])), "k_b": float(0.5 * (k3["b"] + k4["b"])),
            "s_k": float(0.5 * (k3["s_int"] + k4["s_int"]))}


def draw_hyper(rng, base, emp):
    h = dict(base)
    h["a"] = rng.normal(emp["a"], emp["e_a"])
    h["b"] = rng.normal(emp["b"], emp["e_b"])
    h["s"] = float(np.clip(rng.normal(emp["s"], emp["e_s"]), 0.01, None))
    h["ism_scale"] = rng.uniform(0.0, 1.0)
    h["T_cap"] = rng.uniform(2100.0, 2300.0)
    h["E"] = rng.uniform(200.0e3, 400.0e3)
    h["t_sol0"] = rng.uniform(thermal.T_SOL0_LOW, thermal.T_SOL0_HIGH)
    h["f_cc"] = rng.uniform(0.2, 0.5)
    h["Q0"] = rng.normal(thermal.Q_CONV_EARTH, 3.0)
    h["s_eta"] = rng.uniform(0.25, 0.75)
    return h


def draw_normals(n, seed=SEED):
    """Standard normal / uniform deviates shared by all formation-epoch nodes (common random numbers)."""
    rng = np.random.default_rng(seed)
    return {"z_th": rng.standard_normal(n), "z_k": rng.standard_normal(n), "z_vol": rng.standard_normal(n),
            "z_eta": rng.standard_normal(n), "z_T": rng.standard_normal(n), "u_m": rng.random(n)}


def hpe_factors(age, z, h, vary=("th", "k", "vol")):
    """(N, 4) initial concentration factors (238U, 235U, 232Th, 40K) relative to the Earth for stars of
    age `age` today (age 0 is also used for stars that form in the future).  The r-process-per-rock
    level (zero point, raw age slope, scatter) is data; the nuclide-specific free-decay factors are the
    one-zone model, applied with strength ism_scale."""
    age_c = float(np.clip(age, 0.0, AGE_CLIP))
    logf_r = h["a"] + h["b"] * (age_c - AGE_SUN) + (h["s"] * z["z_th"] if "th" in vary else 0.0)
    sc = h.get("ism_scale", ISM_SCALE_BASE)
    d_th = sc * ism_factor("Th232", age_c)
    d_u238 = sc * ism_factor("U238", age_c)
    d_u235 = sc * ism_factor("U235", age_c)
    logf_k = (h["k_a"] + h["k_b"] * (age_c - AGE_SUN) + sc * ism_factor("K40", age_c)
              + (h["s_k"] * z["z_k"] if "k" in vary else 0.0) + (h["s_vol"] * z["z_vol"] if "vol" in vary else 0.0))
    zero = np.zeros(len(z["z_th"]))
    logf = np.column_stack([zero + logf_r + d_u238, zero + logf_r + d_u235, zero + logf_r + d_th, zero + logf_k])
    return 10.0 ** logf


def planet_params(age, z, h, model, vary=("th", "k", "vol", "mass", "eta", "Ti")):
    n = len(z["z_th"])
    mass = np.exp(np.log(h["m_lo"]) + z["u_m"] * np.log(h["m_hi"] / h["m_lo"])) if "mass" in vary else np.ones(n)
    eta = 10.0 ** (h["s_eta"] * z["z_eta"]) if "eta" in vary else np.ones(n)
    T0 = np.full(n, float(model.T_init)) + (h["s_Ti"] * z["z_T"] if "Ti" in vary else 0.0)
    T0 = np.clip(T0, model.t_sol0 + 50.0, 2600.0)       # integrate() additionally applies the thermostat
    return {"fac": hpe_factors(age, z, h, vary), "mass": mass, "eta_fac": eta}, T0


_MODEL_CACHE = {}


def make_model(name, h):
    """Calibrated thermal model for the hyper-parameters h (cached: calibration depends only on these)."""
    key = (name, h["E"], h["f_cc"], h["Q0"], h["t_sol0"], h.get("T_cap", thermal.T_CAP_DEFAULT))
    if key not in _MODEL_CACHE:
        if len(_MODEL_CACHE) > 64:
            _MODEL_CACHE.clear()
        _MODEL_CACHE[key] = thermal.make_model(name, E=h["E"], f_cc=h["f_cc"], Q0=h["Q0"], t_sol0=h["t_sol0"],
                                               T_cap=h.get("T_cap", thermal.T_CAP_DEFAULT))
    return _MODEL_CACHE[key]


def threshold_draw(n, lo, hi, seed):
    """Per-planet degassing thresholds, log-uniform between lo and hi (fraction of modern Earth)."""
    rng = np.random.default_rng(seed)
    return np.exp(rng.uniform(np.log(lo), np.log(hi), n))


def run_population(name, h, z, node_ages=NODE_AGES, f_crit=0.1, dt=0.02, vary=("th", "k", "vol", "mass", "eta", "Ti"),
                   model=None):
    """Lifetimes for every node.  Returns dict with arrays of shape (n_nodes, N): `lifetime`
    (degassing criterion, ridge law), `lifetime_arc` (degassing criterion, plate-speed law) and
    `melt_end` (melting criterion), plus peak temperatures and the calibrated model.  f_crit may be a
    scalar or an array of N per-planet thresholds."""
    if model is None:
        model = make_model(name, h)
    n = len(z["z_th"])
    facs, masses, etas, T0s = [], [], [], []
    for age in node_ages:
        p, T0 = planet_params(age, z, h, model, vary)
        facs.append(p["fac"]); masses.append(p["mass"]); etas.append(p["eta_fac"]); T0s.append(T0)
    p_all = {"fac": np.concatenate(facs), "mass": np.concatenate(masses), "eta_fac": np.concatenate(etas)}
    k = len(node_ages)
    fc = f_crit if np.isscalar(f_crit) else np.tile(np.asarray(f_crit, dtype=float), k)
    res = thermal.integrate(model, np.concatenate(T0s), p_all, T_MAX[name], dt=dt, f_crit=fc,
                            dt_early=DT_EARLY[name], t_early=1.0, dt_late=0.05)
    return {"lifetime": res["lifetime"].reshape(k, n), "melt_end": res["melt_end"].reshape(k, n),
            "lifetime_arc": res["lifetime_arc"].reshape(k, n),
            "T_peak": res["T_peak"].reshape(k, n), "at_cap": res["at_cap"].reshape(k, n),
            "mass": p_all["mass"].reshape(k, n), "logf_th": np.log10(p_all["fac"][:, 2]).reshape(k, n),
            "log_eta": np.log10(p_all["eta_fac"]).reshape(k, n), "model": model, "t_max": T_MAX[name]}


def survival(life, ages=TARGET_AGES):
    return np.array([(life > a).mean() for a in ages])