"""Part 1: merge the curated outcomes with NASA Exoplanet Archive parameters.

Output: data/rocky_atmosphere_census.csv (one row per planet) and
        data/temperate_targets.csv (prediction targets).
"""
import os
import re

import numpy as np
import pandas as pd

from common import DATA_DIR, V_ESC_EARTH, rocky_mass_from_radius, teq_from_insolation
from curated_outcomes import OUTCOMES, PREDICTION_TARGETS

LAVA_TEQ = 1100.0          # K; "lava world" threshold of Park Coy et al. 2026 (arXiv:2604.11911)
PASS_FACTOR = 2.6          # mid-point of the 2.1-3.1x XUV-fluence correction of Pass et al. 2025
PASS_MAX_MASS = 0.35       # Msun; "mid-to-late" (fully convective) M dwarfs


def host_class(teff, spectype=None):
    st = str(spectype).strip().upper() if spectype is not None else ""
    if st and st[0] in "FGKM":
        return st[0]
    if teff < 3900:
        return "M"
    if teff < 5300:
        return "K"
    if teff < 6000:
        return "G"
    return "F"


def derive(df):
    """Add derived physical columns to a NASA-archive frame."""
    df = df.copy()
    lum = 10.0 ** df["st_lum"]
    s_calc = lum / df["pl_orbsmax"] ** 2
    df["S_earth"] = np.where(np.isfinite(s_calc), s_calc, df["pl_insol"])
    df["S_source"] = np.where(np.isfinite(s_calc), "L*/a^2", "archive pl_insol")
    df["Teq_K"] = teq_from_insolation(df["S_earth"])
    # mass: use measured value unless it is an upper limit or itself from a mass-radius relation
    measured = (df["pl_bmasselim"].fillna(0) == 0) & df["pl_bmassprov"].isin(["Mass", "Msini"]) \
        & np.isfinite(df["pl_bmasse"])
    df["mass_used"] = np.where(measured, df["pl_bmasse"], rocky_mass_from_radius(df["pl_rade"]))
    df["mass_source"] = np.where(measured, df["pl_bmassprov"], "Earth-like M-R relation (R=M^0.27)")
    # RV-only planets: the archive mass is m sin i, a LOWER limit on the true mass (and so on v_esc);
    # for an isotropic orientation the median true mass is 1.15 x m sin i (v_esc +7%).
    df["mass_is_lower_limit"] = df["mass_source"] == "Msini"
    df["v_esc_kms"] = V_ESC_EARTH * np.sqrt(df["mass_used"] / df["pl_rade"])
    df["density_rel_earth"] = df["mass_used"] / df["pl_rade"] ** 3
    df["host_class"] = [host_class(t, s) for t, s in zip(df["st_teff"], df["st_spectype"])]
    # cumulative-XUV proxy, Zahnle & Catling (2017): I_xuv/I_xuv,Earth = S (L*/Lsun)^-0.6
    df["xuv_proxy_ZC17"] = df["S_earth"] * lum ** -0.6
    df["xuv_proxy_Pass25"] = df["xuv_proxy_ZC17"] * np.where(df["st_mass"] < PASS_MAX_MASS, PASS_FACTOR, 1.0)
    df["lava_world"] = df["Teq_K"] >= LAVA_TEQ
    return df


def build(nea_sel, meta):
    d = derive(nea_sel).set_index("pl_name")
    rows = []
    for o in OUTCOMES:
        p = d.loc[o["planet"]]
        ids = o["refs"].split(";")
        cites = []
        for i in ids:
            m = meta.get(i, {})
            au = m.get("first_author", "?").split()[-1]
            cites.append(f"{au}+ {m.get('published', '????')[:4]} arXiv:{i}")
        rows.append(dict(
            planet=o["planet"], host_spectype=re.sub(r"&plusmn;", "+/-", str(p["st_spectype"])),
            host_class=p["host_class"], host_teff_K=p["st_teff"], host_mass_Msun=p["st_mass"],
            host_logL_Lsun=round(p["st_lum"], 3), radius_Re=p["pl_rade"],
            mass_Me=round(p["mass_used"], 3), mass_source=p["mass_source"],
            mass_is_lower_limit=bool(p["mass_is_lower_limit"]),
            density_rel_earth=round(p["density_rel_earth"], 2),
            S_earth=round(p["S_earth"], 3), Teq_K=round(p["Teq_K"], 0),
            v_esc_kms=round(p["v_esc_kms"], 2),
            xuv_proxy_ZC17=round(p["xuv_proxy_ZC17"], 1), xuv_proxy_Pass25=round(p["xuv_proxy_Pass25"], 1),
            xuv_published=o["xuv"], lava_world=bool(p["lava_world"]), group=o["group"],
            modes=o["modes"], outcome=o["outcome"], key_result=o["key"],
            pi_base=o["pi"][0], pi_pess=o["pi"][1], pi_opt=o["pi"][2],
            references="; ".join(cites),
            urls="; ".join(f"https://arxiv.org/abs/{i}" for i in ids)))
    census = pd.DataFrame(rows)
    census.to_csv(os.path.join(DATA_DIR, "rocky_atmosphere_census.csv"), index=False)

    t = d.loc[[n for n in PREDICTION_TARGETS if n in d.index]].reset_index()
    targets = t[["pl_name", "st_spectype", "st_mass", "st_lum", "pl_rade", "mass_used", "mass_source",
                 "mass_is_lower_limit", "S_earth", "Teq_K", "v_esc_kms", "xuv_proxy_ZC17", "xuv_proxy_Pass25"]].copy()
    targets.to_csv(os.path.join(DATA_DIR, "temperate_targets.csv"), index=False)
    return census, targets


if __name__ == "__main__":
    import fetch_data
    sel, temp, meta = fetch_data.main()
    c, t = build(sel, meta)
    pd.set_option("display.width", 250); pd.set_option("display.max_rows", 200)
    print(c[["planet", "host_class", "host_mass_Msun", "host_logL_Lsun", "radius_Re", "mass_Me",
             "density_rel_earth", "S_earth", "Teq_K", "v_esc_kms", "xuv_proxy_ZC17", "lava_world",
             "group", "outcome", "pi_base"]].to_string())
    print(t.to_string())