#!/usr/bin/env python3
"""Pillar 2: direct age census of living Sun-like main-sequence stars in
volume-limited samples.

  (a) Geneva-Copenhagen Survey, Casagrande et al. (2011): d < 40 pc
  (b) Gaia DR3 FLAME: d < 100 pc and d < 200 pc

Individual isochrone ages of unevolved Sun-like stars are poorly constrained
(median 1-sigma errors of 50-70% in the GCS), so the raw histogram of
catalogue ages is dominated by the prior.  We therefore (i) forward-model the
errors with an EM deconvolution on a COARSE age grid (five bins, which is all
the data can support), (ii) repeat for precise-age subsamples, and (iii)
weight each star by its vertical orbit to turn the local volume census into
a census of the vertical column.  We also quantify the catalogue-to-catalogue
systematic by comparing FLAME and GCS ages star by star.
"""
import numpy as np
import pandas as pd
import common as C

COARSE = np.array([0.0, 2.0, C.T_SUN, C.T_SUN + 2.0, C.T_SUN + 5.0, 14.0])


def coarse_to_fine(phi_c, coarse=COARSE, edges=C.AGE_EDGES):
    lo, hi = edges[:-1], edges[1:]
    phi = np.zeros(len(lo))
    for k in range(len(coarse) - 1):
        ov = np.clip(np.minimum(hi, coarse[k + 1]) - np.maximum(lo, coarse[k]),
                     0, None)
        phi += phi_c[k] * ov / (coarse[k + 1] - coarse[k])
    return phi / phi.sum()


def coarse_likelihood(p16, p50, p84, lo=0.0, hi=14.0, coarse=COARSE, nsub=12):
    """Likelihood of each star integrated over each coarse age bin (uniform
    true-age density inside a bin).  The per-star likelihood is a normal
    truncated to the age range [lo, hi] of the isochrone grid, fitted to the
    reported 16/50/84th percentiles (common.truncnorm_fit)."""
    mu, sg = C.truncnorm_fit(p16, p50, p84, lo, hi)
    L = np.zeros((len(mu), len(coarse) - 1))
    for k in range(len(coarse) - 1):
        a, b = max(coarse[k], lo), min(coarse[k + 1], hi)
        t = a + (np.arange(nsub) + 0.5) / nsub * (b - a)
        dens = np.exp(-0.5 * ((t[None, :] - mu[:, None]) / sg[:, None]) ** 2)
        L[:, k] = dens.mean(axis=1) * (b - a) / (coarse[k + 1] - coarse[k])
    width = np.diff(coarse)[None, :]
    return L / np.maximum((L * width).sum(axis=1, keepdims=True), 1e-300)


def census(p16, p50, p84, w=None, n_boot=200, seed=5, label="", lo=0.0,
           hi=14.0):
    """Raw and deconvolved head-start statistics with bootstrap errors."""
    p16, p50, p84 = (np.asarray(v, float) for v in (p16, p50, p84))
    n = len(p50)
    w = np.ones(n) if w is None else np.asarray(w, float)
    L = coarse_likelihood(p16, p50, p84, lo, hi)

    def fit(idx):
        # maximum-likelihood bin masses m_k: p(data_i) = sum_k m_k L_ik
        m = C.ml_mixture(L[idx], w[idx])
        st = C.headstart_stats(coarse_to_fine(m))
        st["young_lt2"] = float(m[0])
        return st

    def raw(idx):
        st = C.headstart_stats(C.weighted_hist(p50[idx], w[idx]))
        st["young_lt2"] = float(np.sum(w[idx] * (p50[idx] < 2)) / w[idx].sum())
        return st

    rng = np.random.default_rng(seed)
    out = dict(sample=label, n=n, n_eff=float(w.sum() ** 2 / (w ** 2).sum()))
    full = np.arange(n)
    for tag, f in (("raw", raw), ("deconv", fit)):
        st = f(full)
        bs = [f(rng.integers(0, n, n)) for _ in range(n_boot)]
        for k in ("f_gt_0", "f_gt_1", "f_gt_2", "f_gt_3", "f_gt_5",
                  "mean_headstart_older", "young_lt2"):
            out[f"{tag}_{k}"] = st[k]
            lo_, hi_ = np.percentile([b[k] for b in bs], [16, 84])
            out[f"{tag}_{k}_lo"], out[f"{tag}_{k}_hi"] = lo_, hi_
    return out


# --------------------------------------------------------------------- GCS
def load_gcs():
    g = pd.read_csv(C.RAW + "gcs_casagrande2011.csv").rename(
        columns={"[Fe/H]": "feh", "[M/H]": "mh", "[a/Fe]": "afe"})
    g["d"] = 1000.0 / g.plx.where(g.plx > 0)
    g["z"] = g.d * np.sin(np.radians(g.GLAT)) + C.Z_SUN       # pc from plane
    return g


def gcs_census(pot, rows, dmax=40.0, mlo=0.9, mhi=1.2):
    g = load_gcs()
    base = g[(g.d < dmax) & g.ageMP.notna() & (g.M_MP > mlo) & (g.M_MP < mhi)
             & g.WVel.notna()].reset_index(drop=True)
    E = pot.energy(base.z.values, base.WVel.values + C.W_SUN)
    wcol = 1.0 / pot.sphere_fraction(E, dmax)
    tag = f"GCS d<{dmax:.0f} pc, {mlo}-{mhi} Msun"
    for iso, (a, lo, hi) in {"Padova": ("ageMP", "age16P", "age84P"),
                             "BaSTI": ("ageMB", "age16B", "age84B")}.items():
        p50, p16, p84 = base[a].values, base[lo].values, base[hi].values
        rows.append(census(p16, p50, p84, None, label=f"{tag}, {iso}, volume"))
        rows.append(census(p16, p50, p84, wcol, label=f"{tag}, {iso}, column"))
        rel = 0.5 * (p84 - p16) / np.maximum(p50, 0.1)
        m = rel < 0.25
        rows.append(census(p16[m], p50[m], p84[m], None,
                           label=f"{tag}, {iso}, volume, sigma<25% only"))
        for floor in (-0.5, -0.3):
            m = base.feh.values > floor
            rows.append(census(p16[m], p50[m], p84[m], wcol[m],
                               label=f"{tag}, {iso}, column, [Fe/H]>{floor}"))
    return base


# ------------------------------------------------------------- Gaia FLAME
def load_flame():
    g = pd.read_csv(C.RAW + "gaia_flame_200pc.csv")
    g["d"] = 1000.0 / g.parallax
    return g


def flame_census(pot, rows):
    g = load_flame()
    sun = g[(g.mass_flame > 0.8) & (g.mass_flame < 1.2)
            & (g.evolstage_flame < 420)]          # main sequence + subgiants
    info = dict(n_sunlike_200pc=len(sun),
                n_with_age=int(sun.age_flame.notna().sum()),
                n_with_rv=int(sun.radial_velocity.notna().sum()))
    s = sun[sun.age_flame.notna() & sun.radial_velocity.notna()
            & (sun.ruwe < 1.4)].reset_index(drop=True)
    U, V, W, z = C.galactic_uvw(s.l.values, s.b.values, s.parallax.values,
                                s.pmra.values, s.pmdec.values,
                                s.radial_velocity.values, s.ra.values,
                                s.dec.values)
    s["W"], s["z"] = W + C.W_SUN, z + C.Z_SUN
    info["n_used_200pc"] = len(s)
    for dmax in (100.0, 200.0):
        t = s[s.d < dmax]
        E = pot.energy(t.z.values, t.W.values)
        wcol = 1.0 / pot.sphere_fraction(E, dmax)
        p50 = t.age_flame.values
        p16, p84 = t.age_flame_lower.values, t.age_flame_upper.values
        tag = f"Gaia FLAME d<{dmax:.0f} pc, 0.8-1.2 Msun"
        nb = 20 if dmax > 150 else 40
        kw = dict(n_boot=nb, lo=0.2, hi=13.5)       # FLAME age grid limits
        rows.append(census(p16, p50, p84, None, label=f"{tag}, volume", **kw))
        rows.append(census(p16, p50, p84, wcol, label=f"{tag}, column", **kw))
        m = (t.mh_gspspec.values > -0.5)
        rows.append(census(p16[m], p50[m], p84[m], wcol[m],
                           label=f"{tag}, column, [M/H]spec>-0.5", **kw))
    return info


def flame_vs_gcs():
    """Star-by-star comparison of FLAME and GCS (Padova) ages for the stars
    in common (via the Hipparcos cross-match)."""
    h = pd.read_csv(C.RAW + "gaia_flame_hip.csv")
    g = load_gcs()
    m = g.merge(h, left_on="HIP", right_on="hip", how="inner")
    m = m[m.ageMP.notna() & m.age_flame.notna()
          & (m.M_MP > 0.8) & (m.M_MP < 1.2)]
    rows = [dict(subset="all Sun-like in common", n=len(m),
                 median_diff=float((m.age_flame - m.ageMP).median()),
                 rms_diff=float((m.age_flame - m.ageMP).std()),
                 f_older_flame=float((m.age_flame > C.T_SUN).mean()),
                 f_older_gcs=float((m.ageMP > C.T_SUN).mean()))]
    for lo, hi in [(-2.5, -0.5), (-0.5, -0.3), (-0.3, -0.1), (-0.1, 0.1),
                   (0.1, 0.6)]:
        t = m[(m.feh > lo) & (m.feh <= hi)]
        rows.append(dict(subset=f"{lo} < [Fe/H] <= {hi}", n=len(t),
                         median_diff=float((t.age_flame - t.ageMP).median()),
                         rms_diff=float((t.age_flame - t.ageMP).std()),
                         f_older_flame=float((t.age_flame > C.T_SUN).mean()),
                         f_older_gcs=float((t.ageMP > C.T_SUN).mean())))
    return pd.DataFrame(rows)


# ------------------------------------------------------------ solar twins
def solar_twins():
    """Sun's age rank among spectroscopic solar twins (precise ages, but a
    selection in Teff/logg/[Fe/H] that truncates old ages)."""
    rows = []
    sp = pd.read_csv(C.RAW + "spina2018_solar_twins.csv")
    rows.append(dict(sample="Spina+ 2018 solar twins (HARPS)", n=len(sp),
                     median_age=float(sp.Agemp.median()),
                     median_err=float(((sp.AgeUL1s - sp.AgeLL1s) / 2).median()),
                     f_older=float((sp.Agemp > C.T_SUN).mean())))
    ni = pd.read_csv(C.RAW + "nissen2020_solar_type.csv")
    rows.append(dict(sample="Nissen+ 2020 solar-type (HARPS)", n=len(ni),
                     median_age=float(ni.Age.median()),
                     median_err=float(ni.e_Age.median()),
                     f_older=float((ni.Age > C.T_SUN).mean())))
    ta = pd.read_csv(C.RAW + "taniguchi2026_solar_twins.csv")
    for col in ("Age(GMAG)", "Age(KsMAG)", "Age(logg)"):
        a = ta[col].dropna()
        err = ((ta["B_" + col] - ta["b_" + col]) / 2).median()
        rows.append(dict(sample=f"Taniguchi+ 2026 Gaia GSP-Spec twins, {col}",
                         n=len(a), median_age=float(a.median()),
                         median_err=float(err),
                         f_older=float((a > C.T_SUN).mean())))
    for r in rows:
        r["f_older_binom_err"] = float(np.sqrt(r["f_older"] * (1 - r["f_older"])
                                               / r["n"]))
    return pd.DataFrame(rows)


def main():
    pot = C.get_potential("fiducial")
    rows = []
    gcs_census(pot, rows)
    info = flame_census(pot, rows)
    df = pd.DataFrame(rows)
    df.to_csv(C.DATA + "/ms_census.csv", index=False)
    cmp_ = flame_vs_gcs()
    cmp_.to_csv(C.DATA + "/flame_vs_gcs.csv", index=False)
    tw = solar_twins()
    tw.to_csv(C.DATA + "/solar_twins.csv", index=False)
    pd.set_option("display.width", 250, "display.max_columns", 30,
                  "display.max_colwidth", 70)
    print("[ms] FLAME sample:", info)
    print(df[["sample", "n", "n_eff", "raw_f_gt_0", "deconv_f_gt_0",
              "deconv_f_gt_0_lo", "deconv_f_gt_0_hi", "deconv_f_gt_2",
              "deconv_f_gt_5", "deconv_young_lt2"]].round(3)
          .to_string(index=False))
    print(cmp_.round(3).to_string(index=False))
    print(tw.round(3).to_string(index=False))
    return df, info


if __name__ == "__main__":
    main()