"""Step 2: build the main-sequence 100 pc parent sample and attach AllWISE photometry.

Outputs
  RAW/derived/parent_sample.csv.gz : one row per GCNS main-sequence star with d50 <= 100 pc
                              (the denominator of every fraction in this study), with the
                              matched AllWISE row (if any) and neighbour diagnostics.
  RAW/derived/neighbours.csv.gz    : the other AllWISE sources within 12 arcsec of matched stars.
  data/sample_counts.csv    : the cut-by-cut row counts.
"""
import glob
import os

import numpy as np
import pandas as pd

from common import RAW, DATA, F0, ZP_DN, mag2jy
from dustmap import sfd_ebv

GCNS_FILE = os.path.join(RAW, "gcns_table1c.csv")
AW_DIR = os.path.join(RAW, "allwise_chunks")
MAMAJEK = os.path.join(RAW, "EEM_dwarf_UBVIJHK_colors_Teff.txt")
DERIVED = os.path.join(RAW, "derived")   # large intermediate tables (regenerable)

MATCH_RADIUS = 3.0      # arcsec, Gaia position at epoch 2010.5 vs AllWISE
NEIGH_RADIUS = 12.0     # arcsec, AllWISE neighbours (W4 beam FWHM)
DM_ABOVE, DM_BELOW = -1.5, 4.5   # main-sequence band in M_G relative to the ridge (4.5 mag below keeps
                                 # hosts dimmed by a 90% swarm, 2.5 mag, inside the sample; only 228 real
                                 # objects lie between 3.2 and 4.5 mag below the ridge)
WDPROB_MAX = 0.5                 # GCNS white-dwarf probability
MG_GIANT = 4.5                   # evolved-star cut applies only brighter than this
MG_F0, MG_M0, MG_MAX = 2.51, 8.16, 16.6   # Mamajek M_G of F0V, M0V, L0V


def read_mamajek():
    """Parse E. Mamajek's dwarf table into a DataFrame (numeric columns only)."""
    lines = open(MAMAJEK).read().splitlines()
    i0 = [i for i, l in enumerate(lines) if l.startswith("#SpT")]
    hdr = lines[i0[0]].replace("#SpT", "SpT", 1).split()[:-1]
    rows = []
    for l in lines[i0[0] + 1:i0[1]]:
        p = l.split()
        if len(p) < len(hdr):
            continue
        rows.append(p[:len(hdr)])
    df = pd.DataFrame(rows, columns=hdr)
    for c in hdr[1:]:
        df[c] = pd.to_numeric(df[c].str.replace(r"[:*]", "", regex=True), errors="coerce")
    return df


def ms_ridge(grp):
    """Main-sequence ridge M_G(G-RP) from the Mamajek table (monotonic interpolation)."""
    m = read_mamajek().dropna(subset=["G-Rp", "M_G"]).sort_values("G-Rp")
    m = m.groupby("G-Rp", as_index=False)["M_G"].mean()
    return np.interp(grp, m["G-Rp"].values, m["M_G"].values)


def sep_arcsec(ra1, de1, ra2, de2):
    d1, d2 = np.radians(de1), np.radians(de2)
    dra = np.radians(ra1 - ra2)
    a = np.sin((d1 - d2) / 2) ** 2 + np.cos(d1) * np.cos(d2) * np.sin(dra / 2) ** 2
    return np.degrees(2 * np.arcsin(np.sqrt(a))) * 3600.0


def main():
    counts = []
    g = pd.read_csv(GCNS_FILE)
    g = g.rename(columns={"GaiaEDR3": "sid", "E(BP/RP)": "bprp_excess", "2MASS": "tmass", "WISE": "gcns_wise"})
    counts.append(("GCNS rows (VizieR J/A+A/649/A6 table1c)", len(g)))
    g["dist_pc"] = g.Dist50 * 1000.0
    g = g[g.dist_pc <= 100.0].copy()
    counts.append(("median distance Dist50 <= 100 pc", len(g)))
    g = g[np.isfinite(g.Gmag) & np.isfinite(g.RPmag)].copy()
    counts.append(("G and RP photometry present", len(g)))
    g["g_rp"] = g.Gmag - g.RPmag
    g["M_G"] = g.Gmag - 5 * np.log10(g.dist_pc / 10.0)
    g["dM"] = g.M_G - ms_ridge(g.g_rp.values)
    # Main-sequence selection.  G-RP is used only to remove white dwarfs / subluminous objects
    # (more than DM_BELOW mag below the ridge at their colour) and evolved stars (more than
    # |DM_ABOVE| mag above the ridge AND M_G < MG_GIANT).  Stars that sit above the ridge at
    # M_G >= MG_GIANT are kept: they are pre-main-sequence stars, unresolved multiples, or
    # (mostly) close-pair components whose RP flux is contaminated, i.e. real dwarfs.
    below = (g.dM >= DM_BELOW) | (g.WDprob > WDPROB_MAX)
    giant = (g.dM <= DM_ABOVE) & (g.M_G < MG_GIANT)
    late = g.M_G > MG_MAX
    ms = ~below & ~giant & ~late
    g = g[ms].copy()
    counts.append((f"main sequence: dM_G < {DM_BELOW} below ridge, not (dM_G < {DM_ABOVE} and M_G < {MG_GIANT}), "
                   f"M_G <= {MG_MAX} (removed: {int(below.sum())} below band or GCNS WDprob > {WDPROB_MAX}, "
                   f"{int(giant.sum())} evolved, {int((late & ~below & ~giant).sum())} fainter than L0)", len(g)))
    # spectral class from absolute magnitude (Mamajek: F0V M_G=2.51, M0V M_G=8.16, L0V M_G=16.6)
    g["cls"] = np.where(g.M_G < MG_F0, "A", np.where(g.M_G < MG_M0, "FGK", "M"))
    for c in ("A", "FGK", "M"):
        counts.append((f"  class {c}", int((g.cls == c).sum())))

    # ------------------------------------------------------------ AllWISE
    files = sorted(glob.glob(os.path.join(AW_DIR, "aw_*.csv")))
    aw = pd.concat([pd.read_csv(f, dtype={"designation": str, "cc_flags": str, "ph_qual": str, "var_flg": str,
                                          "moon_lev": str}, low_memory=False)
                    for f in files], ignore_index=True)
    aw["sep"] = sep_arcsec(aw.ra_in.values, aw.dec_in.values, aw.ra.values, aw.dec.values)
    aw = aw[aw.sid.isin(g.sid)].copy()
    aw["cc_flags"] = aw.cc_flags.fillna("0000").str.zfill(4)
    aw["moon_lev"] = aw.moon_lev.fillna("0000").str.zfill(4)
    # primary = nearest AllWISE source within MATCH_RADIUS
    aw = aw.sort_values(["sid", "sep"])
    prim = aw.groupby("sid", as_index=False).first()
    prim = prim[prim.sep <= MATCH_RADIUS].copy()
    # neighbours = other AllWISE sources within NEIGH_RADIUS of the Gaia position
    key = aw.sid.astype(str) + aw.designation
    pkey = set(prim.sid.astype(str) + prim.designation)
    nb = aw[~key.isin(pkey)].copy()
    for b in ("w1", "w3", "w4"):
        nb[b + "f"] = np.clip(nb[b + "flux"].fillna(0), 0, None)
    agg = nb.groupby("sid").agg(n_aw_neigh=("sep", "size"), neigh_minsep=("sep", "min"),
                                neigh_w1dn=("w1f", "sum"), neigh_w3dn=("w3f", "sum"),
                                neigh_w4dn=("w4f", "sum")).reset_index()
    prim = prim.merge(agg, on="sid", how="left")
    for c in ("n_aw_neigh", "neigh_w1dn", "neigh_w3dn", "neigh_w4dn"):
        prim[c] = prim[c].fillna(0)
    # compact table of the separately catalogued neighbours of stars that have a primary match
    # (used in step 5 to test empirically inside which radius AllWISE merges a source into the star)
    os.makedirs(DERIVED, exist_ok=True)
    nb[nb.sid.isin(prim.sid)][["sid", "sep", "w1mpro", "w3mpro", "w4mpro", "w1snr", "w3snr", "w4snr", "cc_flags",
                               "ext_flg"]].to_csv(os.path.join(DERIVED, "neighbours.csv.gz"), index=False,
                                                  float_format="%.5g")

    # one AllWISE source can be the nearest match of several GCNS stars (close pairs
    # unresolved by WISE): keep it only for the brightest star in G, the others are
    # treated as having no usable WISE photometry (they stay in the denominator).
    prim = prim.merge(g[["sid", "Gmag"]], on="sid")
    prim["n_share"] = prim.groupby("designation").sid.transform("size")
    prim = prim.sort_values(["designation", "Gmag"])
    prim["share_rank"] = prim.groupby("designation").cumcount()
    n_shared_lost = int((prim.share_rank > 0).sum())
    prim = prim[prim.share_rank == 0].drop(columns=["Gmag", "share_rank"])

    # physical fluxes in Jy from the profile-fit DN (kept for non-detections and negatives)
    for b in ("W1", "W2", "W3", "W4"):
        k = 10 ** (-0.4 * ZP_DN[b]) * F0[b]
        prim[f"F{b}"] = prim[f"{b.lower()}flux"] * k
        prim[f"e{b}"] = prim[f"{b.lower()}sigflux"] * k
        # neighbour flux sums in Jy
    for b in ("W1", "W3", "W4"):
        prim[f"neighF{b}"] = prim[f"neigh_{b.lower()}dn"] * 10 ** (-0.4 * ZP_DN[b]) * F0[b]

    keep = ["sid", "designation", "sep", "glon", "glat", "FW1", "eW1", "FW2", "eW2", "FW3", "eW3", "FW4", "eW4",
            "w1mpro", "w1sigmpro", "w2mpro", "w2sigmpro", "w3mpro", "w3sigmpro", "w4mpro", "w4sigmpro",
            "w1snr", "w2snr", "w3snr", "w4snr", "w1rchi2", "w2rchi2", "w3rchi2", "w4rchi2",
            "nb", "na", "w1sat", "w2sat", "w3sat", "w4sat", "cc_flags", "ext_flg", "ph_qual", "moon_lev",
            "w3m", "w4m", "w3cov", "w4cov", "w3sky", "w3sigsk", "w3conf", "w4sky", "w4sigsk", "w4conf",
            "w3mag", "w3sigm", "w4mag", "w4sigm", "xscprox", "n_share",
            "n_aw_neigh", "neigh_minsep", "neighFW1", "neighFW3", "neighFW4"]
    out = g.merge(prim[keep], on="sid", how="left")
    out["has_aw"] = out.designation.notna()
    counts.append((f"  of which matched to an AllWISE source within {MATCH_RADIUS} arcsec "
                   f"({n_shared_lost} more share a WISE source with a brighter GCNS star)", int(out.has_aw.sum())))

    # Galactic coordinates and SFD reddening for every parent star (from the Gaia position)
    from astropy.coordinates import SkyCoord
    import astropy.units as u
    gal = SkyCoord(out.RA_ICRS.values * u.deg, out.DE_ICRS.values * u.deg, frame="icrs").galactic
    out["gl"], out["gb"] = gal.l.deg, gal.b.deg
    out["ebv"] = sfd_ebv(out.gl.values, out.gb.values)
    out["pm_asyr"] = np.hypot(out.pmRA.fillna(0), out.pmDE.fillna(0)) / 1000.0
    out = out.sort_values("sid").reset_index(drop=True)
    os.makedirs(DERIVED, exist_ok=True)
    out.to_csv(os.path.join(DERIVED, "parent_sample.csv.gz"), index=False, float_format="%.9g")
    pd.DataFrame(counts, columns=["cut", "n"]).to_csv(os.path.join(DATA, "sample_counts.csv"), index=False)
    for c, n in counts:
        print(f"{n:8d}  {c}")


if __name__ == "__main__":
    main()