"""Part 2: (a) beta-binomial atmosphere fraction, (b) Bayesian cosmic-shoreline regression and
extrapolation to temperate M-dwarf planets, (c) subjective habitability probability p_M.

Model (b):   Delta = log10 S - [a + p log10(v_esc/v_Earth) + q log10(L*/Lsun)]
             P(atmosphere) = 1 / (1 + exp(Delta / w))            (logistic link; probit as a variant)
             likelihood_i  = pi_i P_i + (1 - pi_i)(1 - P_i)      (pi_i = soft label from the census)
The functional form follows Berta-Thompson et al. 2025 (arXiv:2507.02136); the data set, soft
labels, priors, sampler and predictions are ours.

What the fit can and cannot do (revision after referee report):
  * q (luminosity slope) is bounded by the data from below only; its central value and therefore
    the habitable-zone extrapolation follow the prior (see q_prior_variants and
    data/shoreline_q_prior_sensitivity.csv).
  * a single plane cannot distinguish TRAPPIST-1 c from e, so predictions for individual
    habitable-zone planets are structural, not measurements.
  * logistic tails are floored (common.F_ATM_FLOOR) before they are used in Part 3.
"""
import os

import numpy as np
import pandas as pd
from scipy import stats
from scipy.special import expit, ndtr

from common import DATA_DIR, SEED, V_ESC_EARTH, F_ATM_FLOOR
from curated_outcomes import SOLAR_SYSTEM

PARAMS = ["a", "p", "q", "lnw"]
PRIOR_WEAK = dict(a=(-3.0, 6.0), p=(4.0, 2.0, 0.0, 15.0), q=(0.6, 1.0, -2.0, 5.0),
                  lnw=(np.log(0.03), np.log(1.5)), flat=False)
PRIOR_FLAT = dict(a=(-3.0, 6.0), p=(4.0, 2.0, 0.0, 15.0), q=(0.6, 1.0, -2.0, 5.0),
                  lnw=(np.log(0.03), np.log(1.5)), flat=True)

INFORMATIVE = ("bare_rock_favoured", "thick_atm_ruled_out", "atm_detected")
AMBIGUOUS_ECLIPSE = ["GJ 357 b", "GJ 3473 b", "LHS 1478 b"]     # emission data that fit neither/both readings


def q_prior_variants():
    """Alternative priors on the luminosity slope q, used only to show that q is prior-dominated."""
    out = {}
    for label, (mu, sd) in (("N(0.6,1) [baseline]", (0.6, 1.0)), ("N(0,1)", (0.0, 1.0)), ("N(2,1)", (2.0, 1.0)),
                            ("N(0.6,0.3)", (0.6, 0.3)), ("N(0.6,3)", (0.6, 3.0))):
        pr = dict(PRIOR_WEAK)
        pr["q"] = (mu, sd, -2.0, 5.0)
        out[label] = pr
    return out


def apply_floor(f_atm, eps=F_ATM_FLOOR):
    """Mix the fitted retention probability with a constant 'lapse' term: eps + (1 - eps) f."""
    return eps + (1.0 - eps) * np.asarray(f_atm, dtype=float)


# ----------------------------------------------------------------------------------------------
# (a) fraction of observed M-dwarf rocky planets with a substantial atmosphere
# ----------------------------------------------------------------------------------------------
def beta_summary(k, n, a0, b0):
    post = stats.beta(a0 + k, b0 + n - k)
    return dict(k=k, n=n, prior=f"Beta({a0},{b0})", mean=post.mean(), median=post.median(),
                lo68=post.ppf(0.16), hi68=post.ppf(0.84), upper95=post.ppf(0.95))


def soft_fraction_posterior(pi, grid=None):
    """Posterior of the population fraction f with soft labels: prod_i [pi_i f + (1-pi_i)(1-f)], flat prior."""
    if grid is None:
        grid = np.linspace(0, 1, 2001)
    logl = np.zeros_like(grid)
    for x in pi:
        logl += np.log(x * grid + (1 - x) * (1 - grid) + 1e-300)
    post = np.exp(logl - logl.max())
    post /= np.trapezoid(post, grid)
    cdf = np.concatenate([[0], np.cumsum(0.5 * (post[1:] + post[:-1]) * np.diff(grid))])
    q = lambda u: float(np.interp(u, cdf, grid))  # noqa: E731
    return dict(mean=float(np.trapezoid(grid * post, grid)), median=q(0.5), lo68=q(0.16), hi68=q(0.84),
                upper95=q(0.95))


def host_of(planet):
    return planet.rsplit(" ", 1)[0]


def beta_binomial_table(census):
    c = census
    m_rocky = (c.host_class == "M") & (c.group == "rocky")
    rows = []

    def add(label, mask, k):
        n = int(mask.sum())
        for a0, b0 in ((1, 1), (0.5, 0.5)):
            r = beta_summary(k, n, a0, b0)
            r["sample"] = label
            r["members"] = ", ".join(c.planet[mask])
            rows.append(r)

    strict = m_rocky & c.outcome.isin(INFORMATIVE)
    add("A. strict: M-dwarf rocky, informative verdict", strict, int((c.outcome[strict] == "atm_detected").sum()))
    cool = strict & ~c.lava_world & (c.S_earth < 25)
    add("A2. strict, S < 25 S_Earth only", cool, int((c.outcome[cool] == "atm_detected").sum()))
    emis = strict & (c.planet != "TRAPPIST-1 d")
    add("A3. strict, thermal-emission verdicts only (drops TRAPPIST-1 d, transmission only)", emis, 0)
    # one verdict per host star: planets of one star share an XUV history, so they are not independent
    hosts = sorted({host_of(p) for p in c.planet[strict]})
    r_host = [dict(beta_summary(0, len(hosts), a0, b0), sample="A4. strict, counted per host star",
                   members=", ".join(hosts)) for a0, b0 in ((1, 1), (0.5, 0.5))]
    rows.extend(r_host)
    emis_all = m_rocky & (c.outcome.isin(INFORMATIVE) | c.planet.isin(AMBIGUOUS_ECLIPSE))
    add("B. pessimistic: + ambiguous eclipses counted as airless", emis_all, 0)
    opt = m_rocky & (c.outcome.isin(INFORMATIVE) | c.planet.isin(["GJ 3473 b", "LHS 1478 b", "L 98-59 b"]))
    add("C. optimistic: + LHS 1478 b, GJ 3473 b, L 98-59 b counted as atmospheres", opt, 3)
    opt2 = (c.host_class == "M") & (c.group != "sub_neptune") & (
        c.outcome.isin(INFORMATIVE) | c.planet.isin(["GJ 3473 b", "LHS 1478 b", "L 98-59 b", "L 98-59 d", "LHS 1140 b"]))
    add("D. optimistic + volatile-rich candidates (L 98-59 d, LHS 1140 b) as atmospheres", opt2, 5)
    df = pd.DataFrame(rows)

    soft = []
    allm = (c.host_class == "M") & (c.group == "rocky") & (c.outcome != "unpublished")
    for col in ("pi_base", "pi_pess", "pi_opt"):
        r = soft_fraction_posterior(c.loc[allm, col].values)
        r.update(sample=f"E. soft labels ({col}), all {int(allm.sum())} M-dwarf rocky planets with published data",
                 k=np.nan, n=int(allm.sum()), prior="flat", members=", ".join(c.planet[allm]))
        soft.append(r)
    df = pd.concat([df, pd.DataFrame(soft)], ignore_index=True)
    df = df[["sample", "prior", "k", "n", "mean", "median", "lo68", "hi68", "upper95", "members"]]
    df.to_csv(os.path.join(DATA_DIR, "atmosphere_fraction_posteriors.csv"), index=False)
    return df


# ----------------------------------------------------------------------------------------------
# (b) shoreline regression
# ----------------------------------------------------------------------------------------------
def make_dataset(census, variant="base"):
    """Return dict of arrays for a named data variant."""
    c = census.copy()
    pi_col = {"pi_pess": "pi_pess", "pi_opt": "pi_opt"}.get(variant, "pi_base")
    keep = c.group.isin(["rocky", "volatile_rich?"]) & ~c.lava_world & (c.outcome != "unpublished")
    if variant == "plus_subneptunes":
        keep = keep | (c.group == "sub_neptune")
    if variant == "plus_lava":
        keep = c.group.isin(["rocky", "volatile_rich?"]) & (c.outcome != "unpublished")
    if variant == "rocky_only":
        keep = keep & (c.group == "rocky")
    if variant == "drop_trappist":
        keep = keep & ~c.planet.str.startswith("TRAPPIST-1")
    if variant == "drop_ltt1445":
        keep = keep & (c.planet != "LTT 1445 A b")
    if variant == "solar_system_only":
        keep = keep & False
    ex = c[keep]
    names = list(ex.planet)
    logS = list(np.log10(ex.S_earth))
    logv = list(np.log10(ex.v_esc_kms / V_ESC_EARTH))
    logL = list(ex.host_logL_Lsun)
    pi = list(ex[pi_col])
    kind = ["exo_M" if h == "M" else "exo_other" for h in ex.host_class]
    for b in SOLAR_SYSTEM:
        p_b = b["pi"]
        if variant == "ss_no_thin" and b["thin"]:
            p_b = 0.5
        names.append(b["name"]); logS.append(np.log10(b["S"])); logv.append(np.log10(b["v_esc"] / V_ESC_EARTH))
        logL.append(0.0); pi.append(p_b); kind.append("solar_system")
    return dict(names=np.array(names), logS=np.array(logS), logv=np.array(logv), logL=np.array(logL),
                pi=np.array(pi), kind=np.array(kind))


def prob_atm(theta, logS, logv, logL, link="logit"):
    """theta: (..., 4) array; other args broadcastable 1-D arrays -> (..., N) probabilities."""
    theta = np.atleast_2d(theta)
    a, p, q, lnw = (theta[:, i][:, None] for i in range(4))
    delta = logS[None, :] - (a + p * logv[None, :] + q * logL[None, :])
    z = -delta / np.exp(lnw)
    return expit(z) if link == "logit" else ndtr(z)


def log_prior(theta, prior):
    a, p, q, lnw = theta
    if not (prior["a"][0] <= a <= prior["a"][1]):
        return -np.inf
    if not (prior["lnw"][0] <= lnw <= prior["lnw"][1]):
        return -np.inf
    lp = 0.0
    for val, (mu, sd, lo, hi) in ((p, prior["p"]), (q, prior["q"])):
        if not (lo <= val <= hi):
            return -np.inf
        if not prior["flat"]:
            lp += -0.5 * ((val - mu) / sd) ** 2
    return lp


def log_post(theta, d, prior, link):
    lp = log_prior(theta, prior)
    if not np.isfinite(lp):
        return -np.inf
    P = prob_atm(theta, d["logS"], d["logv"], d["logL"], link)[0]
    like = d["pi"] * P + (1 - d["pi"]) * (1 - P)
    return lp + np.sum(np.log(like + 1e-300))


def run_mcmc(d, prior=PRIOR_WEAK, link="logit", n_chains=4, n_steps=30000, burn=6000, thin=8, seed=SEED):
    """Adaptive random-walk Metropolis (numpy only). Returns samples (N,4) and max split-R-hat."""
    rng = np.random.default_rng(seed)
    start = np.array([2.0, 4.0, 0.8, np.log(0.3)])
    # pilot run to learn the proposal covariance
    cov = np.diag([0.3, 0.6, 0.3, 0.4]) ** 2
    x = start.copy(); lx = log_post(x, d, prior, link)
    pilot = []
    for i in range(8000):
        y = rng.multivariate_normal(x, cov * 0.5)
        ly = log_post(y, d, prior, link)
        if np.log(rng.random()) < ly - lx:
            x, lx = y, ly
        pilot.append(x.copy())
    pilot = np.array(pilot[2000:])
    cov = np.cov(pilot.T) * (2.38 ** 2 / 4) + 1e-8 * np.eye(4)
    chains = []
    for c in range(n_chains):
        x = pilot[rng.integers(len(pilot))].copy(); lx = log_post(x, d, prior, link)
        out = np.empty((n_steps, 4)); acc = 0
        steps = rng.multivariate_normal(np.zeros(4), cov, size=n_steps)
        us = np.log(rng.random(n_steps))
        for i in range(n_steps):
            y = x + steps[i]
            ly = log_post(y, d, prior, link)
            if us[i] < ly - lx:
                x, lx = y, ly; acc += 1
            out[i] = x
        chains.append(out[burn::thin])
    chains = np.array(chains)                       # (C, N, 4)
    # split R-hat
    half = chains.shape[1] // 2
    sp = np.concatenate([chains[:, :half], chains[:, half:2 * half]], axis=0)
    W = sp.var(axis=1, ddof=1).mean(axis=0)
    B = half * sp.mean(axis=1).var(axis=0, ddof=1)
    rhat = np.sqrt(((half - 1) / half * W + B / half) / W)
    return chains.reshape(-1, 4), float(rhat.max())


def summarize_params(samples):
    out = {}
    for i, k in enumerate(PARAMS):
        lo, med, hi = np.percentile(samples[:, i], [16, 50, 84])
        out[k] = (med, lo, hi)
    w = np.exp(samples[:, 3])
    out["w_dex"] = tuple(np.percentile(w, [50, 16, 84]))
    return out


# ----------------------------------------------------------------------------------------------
# population model: habitable-zone rocky planets as a function of stellar mass
# ----------------------------------------------------------------------------------------------
def fit_mass_luminosity(small_mdwarf_planets):
    """Quadratic fit of log L* on log M* for the M-dwarf hosts in the NASA archive sample."""
    h = small_mdwarf_planets.drop_duplicates("hostname")
    h = h[np.isfinite(h.st_mass) & np.isfinite(h.st_lum) & (h.st_mass > 0.07) & (h.st_mass < 0.75)]
    coef = np.polyfit(np.log10(h.st_mass), h.st_lum, 2)
    resid = h.st_lum - np.polyval(coef, np.log10(h.st_mass))
    return coef, len(h), float(np.std(resid))


def logL_of_mass(m, coef):
    return np.polyval(coef, np.log10(m))


HZ_S = np.exp(np.linspace(np.log(0.25), np.log(0.9), 7))      # instellation range of the habitable zone
HZ_R = np.exp(np.linspace(np.log(0.8), np.log(1.5), 8))       # rocky-planet radii, flat in ln R
MASS_GRID_M = np.exp(np.linspace(np.log(0.08), np.log(0.60), 30))


def f_atm_population(samples, coef, link="logit", mass_grid=MASS_GRID_M):
    """f_atm[sample, mass]: mean retention probability of HZ rocky planets for each stellar mass."""
    SS, RR = np.meshgrid(HZ_S, HZ_R)
    logS = np.log10(SS.ravel())
    logv = 1.35 * np.log10(RR.ravel())                         # v_esc/v_E = sqrt(M/R), M = R^3.7
    out = np.empty((len(samples), len(mass_grid)))
    for j, m in enumerate(mass_grid):
        logL = np.full(logS.shape, logL_of_mass(m, coef))
        out[:, j] = prob_atm(samples, logS, logv, logL, link).mean(axis=1)
    return out


def imf_weights_kroupa(mass_grid):
    """Kroupa (2001) number weights dN/dlnM on a log-spaced grid (normalised)."""
    alpha = np.where(mass_grid < 0.5, 1.3, 2.3)
    xi = mass_grid ** (-alpha) * np.where(mass_grid < 0.5, 1.0, 0.5)   # continuity at 0.5
    w = xi * mass_grid                                                  # dN/dlnM
    return w / w.sum()


def predict_targets(samples, targets, link="logit"):
    logS = np.log10(targets.S_earth.values)
    logv = np.log10(targets.v_esc_kms.values / V_ESC_EARTH)
    logL = targets.st_lum.values
    P = prob_atm(samples, logS, logv, logL, link)
    return P


F_WATER_RANGE = (0.05, 1.0)
F_OTHER_RANGE = (0.1, 1.0)


def draw_p_M_factors(n, rng):
    """Two ASSUMED factors beyond present-day atmosphere retention (report.md, section 3.4).

    They are not derived from data.  Each is log-uniform between a lower bound chosen by the author
    and 1; the literature cited in the report motivates that each could be well below 1 or equal
    to 1, but does not fix either bound.  Every Part 3 result is therefore reported both with and
    without these factors.
      f_water : planet keeps or regains a habitable surface water inventory after the host's
                pre-main-sequence super-luminous phase (Luger & Barnes 2015; Kimura & Ikoma 2022)
      f_other : tidally locked climate stays stable, flare/UV environment is survivable and usable,
                conditions persist for the multi-Gyr span the hard steps need
    """
    f_water = np.exp(rng.uniform(np.log(F_WATER_RANGE[0]), np.log(F_WATER_RANGE[1]), n))
    f_other = np.exp(rng.uniform(np.log(F_OTHER_RANGE[0]), np.log(F_OTHER_RANGE[1]), n))
    return f_water, f_other