"""Part 4: the stellar side.  How long does a planet stay inside the habitable zone of a K dwarf?

Stellar evolution tracks: Geneva grid of Mowlavi et al. (2012, A&A 541, A41), Z = 0.014, 0.5-1.0 Msun
(VizieR J/A+A/541/A41).  Habitable-zone flux limits: Kopparapu et al. (2014, ApJL 787, L29), table 1,
1 Earth-mass coefficients: runaway greenhouse (inner edge) and maximum greenhouse (outer edge).
"""
import os

import numpy as np
import pandas as pd

from common import RAW_DIR, DATA_DIR

KOPP = {  # S_eff_sun, a, b, c, d   (T* = Teff - 5780 K)
    "runaway": (1.107, 1.332e-4, 1.58e-8, -8.308e-12, -1.931e-15),
    "max_greenhouse": (0.356, 6.171e-5, 1.698e-9, -3.198e-12, -5.575e-16),
}


def s_eff(limit, teff):
    s0, a, b, c, d = KOPP[limit]
    x = np.clip(np.asarray(teff, dtype=float), 2600.0, 7200.0) - 5780.0
    return s0 + a * x + b * x ** 2 + c * x ** 3 + d * x ** 4


def load_tracks():
    t = pd.read_csv(os.path.join(RAW_DIR, "geneva_tracks_lowmass.csv"))
    t = t.rename(columns={"Xc(H)": "Xc"})
    t["age_gyr"] = t["Age"] / 1e9
    t["L"] = 10.0 ** t["logL"]
    t["Teff"] = 10.0 ** t["logTe"]
    return t.sort_values(["Minit", "Age"]).reset_index(drop=True)


def hz_table(n_orbit=400):
    tracks = load_tracks()
    rows, curves = [], []
    for m, tr in tracks.groupby("Minit"):
        ms = tr[tr["Xc"] > 1e-3]                      # core hydrogen burning
        age, L, teff = ms["age_gyr"].values, ms["L"].values, ms["Teff"].values
        t_ms = age[-1]
        s_in, s_out = s_eff("runaway", teff), s_eff("max_greenhouse", teff)
        # orbits: log-uniform between the HZ edges at the start of the main sequence
        d_in0, d_out0 = np.sqrt(L[0] / s_in[0]), np.sqrt(L[0] / s_out[0])
        d = np.exp(np.linspace(np.log(d_in0), np.log(d_out0), n_orbit + 2)[1:-1])
        flux = L[None, :] / d[:, None] ** 2
        inside = (flux <= s_in[None, :]) & (flux >= s_out[None, :])
        # exit time through the inner edge = first time flux exceeds the runaway limit
        hot = flux > s_in[None, :]
        t_exit = np.where(hot.any(axis=1), age[np.argmax(hot, axis=1)], t_ms)
        # Earth-equivalent orbit: same flux relative to the runaway limit as Earth today (1/1.107) at 4.5 Gyr
        i45 = np.argmin(np.abs(age - 4.5))
        d_e = np.sqrt(L[i45] / (s_in[i45] / 1.107))
        hot_e = L / d_e ** 2 > s_in
        t_exit_e = age[np.argmax(hot_e)] if hot_e.any() else t_ms
        # luminosity growth rate around 4.5 Gyr
        j0, j1 = np.argmin(np.abs(age - 3.5)), np.argmin(np.abs(age - 5.5))
        growth = (L[j1] / L[j0] - 1.0) / (age[j1] - age[j0]) * 100.0
        rows.append({"mass_msun": m, "t_ms_gyr": t_ms, "L_zams": L[0], "L_end_ms_over_zams": L[-1] / L[0],
                     "lum_growth_pct_per_gyr_at_4.5": growth,
                     "hz_exit_median_gyr": float(np.median(t_exit)), "hz_exit_p10_gyr": float(np.percentile(t_exit, 10)),
                     "hz_exit_p90_gyr": float(np.percentile(t_exit, 90)),
                     "frac_zams_hz_orbits_still_in_hz_at_10gyr": float(np.mean(t_exit > 10.0)) if t_ms > 10 else 0.0,
                     "frac_still_in_hz_at_20gyr": float(np.mean(t_exit > 20.0)) if t_ms > 20 else 0.0,
                     "earth_equivalent_orbit_exit_gyr": float(t_exit_e),
                     "max_hz_residence_gyr": float(t_exit.max())})
        curves.append(pd.DataFrame({"mass_msun": m, "age_gyr": age, "L_over_Lzams": L / L[0], "Teff": teff,
                                    "frac_in_hz": inside.mean(axis=0)}))
    tab = pd.DataFrame(rows)
    tab.to_csv(os.path.join(DATA_DIR, "stellar_hz_lifetimes.csv"), index=False)
    cur = pd.concat(curves, ignore_index=True)
    cur.to_csv(os.path.join(DATA_DIR, "stellar_tracks_hz.csv"), index=False)
    return tab, cur


if __name__ == "__main__":
    tab, _ = hz_table()
    with pd.option_context("display.width", 250, "display.max_columns", 30):
        print(tab.round(2).to_string())