"""Step 10: assemble every automated vetting quantity for candidates and controls.

For each target (union candidates with D < log10(0.5), plus 300 random controls):
  * SDSS DR16 cone: spectroscopic neighbours 20"-150" away at the HI redshift (HI confusion); redshift of
    the optical counterpart itself (wrong-counterpart test); bright photometric neighbours.
  * aperture-photometry correction of the catalogue magnitude (ALFALFA-only galaxies).
  * 2MASS XSC K-band and AllWISE W1-W4 (CDS XMatch).
  * mid-infrared test: largest starlight fraction f that could be re-radiated as a blackbody of temperature
    T (100-600 K, evaluated per temperature) without exceeding the observed W3 and W4 fluxes.
Output: data/vet_table.csv
"""
import os
import numpy as np
import pandas as pd
from common import RAW, DATA, C_KMS, bb_frac
from s04_xmatch import xmatch

T50, T75, T90 = np.log10(0.5), np.log10(0.25), np.log10(0.1)
MSUN_I = 4.53
SAFETY_MIR = 2.5      # AllWISE profile-fit fluxes may miss extended emission: allow 1 mag more
T_WARM = np.array([100., 150., 200., 250., 300., 350., 400., 500., 600.])
LAM = {"W3": 11.56, "W4": 22.09, "i": 0.7499}
F0 = {"W3": 31.674, "W4": 8.363}      # Jy, Vega zero points
LSUN = 3.828e33


def nuLnu_from_mag(mag_vega, band, dist_mpc):
    fnu = F0[band] * 10 ** (-0.4 * mag_vega) * 1e-23               # erg/s/cm2/Hz
    nu = 2.998e14 / LAM[band]
    return 4 * np.pi * (dist_mpc * 3.0857e24) ** 2 * fnu * nu / LSUN


def nuLnu_i(logLi):
    """nu L_nu (Lsun) at SDSS i from the i-band luminosity in solar i-band units (AB)."""
    return 10 ** logLi * 0.70   # the Sun: nu L_nu(i) = 0.70 Lsun (M_i,AB = 4.53)


def main():
    t = pd.read_csv(os.path.join(DATA, "vet_targets.csv"))
    # refresh flags that depend on the union file (auto_clean definition lives there)
    u = pd.read_csv(os.path.join(DATA, "union_sample.csv")).set_index("uid")
    t["auto_clean"] = t.uid.map(u.auto_clean)
    for c in ("n_simbad_nb", "n_simbad_nb10", "simbad_nb_sep", "simbad_nb_name"):
        t[c] = t.uid.map(u[c])
    cosd = np.cos(np.radians(t.dec.values))

    # ---- SDSS DR16 cones
    rec = []
    for r in t.itertuples():
        o = dict(uid=r.uid, cone_ok=False, in_sdss=False, n_spec_same=0, sep_spec_same=np.nan,
                 self_dv=np.nan, n_phot_nb=0, n_gal_cone=0)
        p = os.path.join(RAW, "cones", f"sdss16_{r.uid}.csv")
        if os.path.exists(p):
            try:
                c = pd.read_csv(p)
                o["cone_ok"] = True
            except Exception:
                c = pd.DataFrame()
            if len(c):
                o["in_sdss"] = True
                sep = np.hypot((c.RA_ICRS - r.ra) * np.cos(np.radians(r.dec)), c.DE_ICRS - r.dec) * 3600
                vz = c.zsp * C_KMS
                has = c.zsp.notna()
                same = has & (np.abs(vz - r.cz) < r.W / 2 + 200)
                nb = same & (sep >= 20) & (sep <= 150)
                o["n_spec_same"] = int(nb.sum())
                o["sep_spec_same"] = float(sep[nb].min()) if nb.any() else np.nan
                me = has & (sep < 10)
                if me.any():
                    o["self_dv"] = float((vz[me] - r.cz).iloc[np.argmin(sep[me].values)])
                zt = r.cz / C_KMS
                ph = (~has) & (sep >= 20) & (sep <= 120) & (c.rmag < 18.0) & c.zph.notna() & \
                     (np.abs(c.zph - zt) < 2 * c.e_zph.fillna(0.03) + 0.02)
                o["n_phot_nb"] = int(ph.sum())
                o["n_gal_cone"] = int(len(c))
        rec.append(o)
    v = t.merge(pd.DataFrame(rec), on="uid", how="left")
    v["oc_z_mismatch"] = v.self_dv.abs() > 400

    # ---- aperture photometry (ALFALFA-only)
    ap_path = os.path.join(DATA, "aperture_phot.csv")
    if os.path.exists(ap_path):
        ap = pd.read_csv(ap_path)[["uid", "phot_ok", "detected", "ap_valid", "dmag_mean", "dmag_med", "r_ap", "q_img",
                                   "cen_off", "frac_masked", "i_ap_med", "i_ap_mean"]]
        v = v.merge(ap, on="uid", how="left")
        v["ap_valid"] = v.ap_valid.fillna(False).astype(bool)
        # residual with the image-based i-band flux; only where the photometered segment is the target itself
        v["D_ap_raw"] = v.r_i + 0.4 * v.dmag_med
        v["D_ap"] = v.D_ap_raw.where(v.ap_valid)
    # ---- 2MASS XSC and AllWISE
    v["AGC_key"] = v["AGC"]
    v["row"] = np.arange(len(v))
    # candidates and the 300 main controls only (the extra aperture controls need no infrared photometry)
    xin = v[v.role != "control_ap"].rename(columns={"AGC": "AGC_orig"})
    xin = xin.assign(AGC=xin.row)

    def cached_xmatch(cat, radius, fn):
        """CDS XMatch keyed by row number; a cache made for a different target list is detected and redone."""
        for force in (False, True):
            x = xmatch(xin, cat, radius, fn, force=force)
            pos = x[["AGC", "ra", "dec"]].drop_duplicates("AGC").merge(xin[["row", "ra", "dec"]], left_on="AGC",
                                                                      right_on="row", how="left", suffixes=("", "_t"))
            if len(pos) == 0 or (np.abs(pos.ra - pos.ra_t).max() < 1e-4 and np.abs(pos.dec - pos.dec_t).max() < 1e-4):
                return x
            print(f"    cached {fn} does not match the current target list; querying again")
        raise RuntimeError(f"{fn}: cross-match rows do not correspond to the targets")
    x = cached_xmatch("vizier:VII/233/xsc", 15, "xm_vet_2mass.csv").sort_values("angDist").drop_duplicates("AGC")
    w = cached_xmatch("vizier:II/328/allwise", 8, "xm_vet_allwise.csv").sort_values("angDist").drop_duplicates("AGC")
    v = v.merge(x.rename(columns={"AGC": "row", "K.ext": "Kext", "Sb/a": "ba_2mass"})[["row", "Kext", "ba_2mass"]],
                on="row", how="left")
    v = v.merge(w.rename(columns={"AGC": "row"})[["row", "W1mag", "W2mag", "W3mag", "e_W3mag", "W4mag", "e_W4mag", "ex"]],
                on="row", how="left")

    # ---- mid-infrared (warm waste heat) test
    # observed stellar light: i band where available, otherwise W1 converted with the sample-median i-W1 offset
    res_cf = pd.read_csv(os.path.join(DATA, "cf4_resid.csv"), usecols=["PGC", "logLi", "logLW1"])
    off = (res_cf.logLW1 - res_cf.logLi).median()
    al = pd.read_csv(os.path.join(DATA, "alfalfa_resid.csv"), low_memory=False, usecols=["AGC", "logLi"])
    d70 = np.where(v.src == "cf4", v.Dist * 75.0 / 70.0, v.Dist)
    logLi = np.where(v.src == "alfalfa", v.AGC.map(al.set_index("AGC").logLi), np.nan)
    cfm = res_cf.set_index("PGC")
    li_cf = v.PGC.map(cfm.logLi)
    lw_cf = v.PGC.map(cfm.logLW1) - off
    # CF4 luminosities were computed with H0 = 75; put them on the H0 = 70 scale (L ~ D^2)
    logLi = np.where(v.src == "cf4", np.where(li_cf.notna(), li_cf, lw_cf) + 2 * np.log10(75.0 / 70.0), logLi)
    v["logLi_70"] = logLi                      # i-band luminosity (solar units) on the H0=70 scale
    Lstar = nuLnu_i(v.logLi_70)                # conservative lower bound on the bolometric stellar luminosity
    L3 = nuLnu_from_mag(v.W3mag, "W3", d70) * SAFETY_MIR
    L4 = nuLnu_from_mag(v.W4mag, "W4", d70) * SAFETY_MIR
    v["nuLnu_W3_over_i"], v["nuLnu_W4_over_i"] = L3 / SAFETY_MIR / Lstar, L4 / SAFETY_MIR / Lstar
    fmax = np.zeros(len(v))
    has_mir = (v.W3mag.notna() | v.W4mag.notna()).values
    for T in T_WARM:
        Lw = np.fmin(L3 / bb_frac(LAM["W3"], T), L4 / bb_frac(LAM["W4"], T))   # max waste heat allowed at T
        fT = Lw / (Lw + Lstar)
        v[f"f_warm_max_T{T:.0f}"] = np.where(has_mir, fT, np.nan)             # per temperature
        fmax = np.fmax(fmax, fT)
    v["f_warm_max"] = np.where(has_mir, fmax, np.nan)                         # most permissive temperature

    # ---- automated verdict
    def verdict(r, thr):
        why = []
        if r.n_spec_nb > 0 or r.n_hi_nb > 0 or r.n_spec_same > 0 or r.n_simbad_nb > 0:
            why.append("confusion(spec/HI neighbour in beam)")
        if r.hi_opt_off > 40:
            why.append("HI-optical offset>40arcsec")
        if r.e_W > 20:
            why.append("line-width error>20km/s")
        if r.oc_z_mismatch:
            why.append("counterpart redshift mismatch")
        if thr > -5 and r.src == "alfalfa" and "D_ap" in v.columns and np.isfinite(r.D_ap) and r.D_ap > thr:
            why.append("catalogue photometry too faint (aperture)")
        return "; ".join(why)
    for name, thr in (("50", T50), ("75", T75), ("90", T90)):
        v["auto_reason_" + name] = [verdict(r, thr) if r.D < thr else "" for r in v.itertuples()]
    v["auto_reason_any"] = [verdict(r, -9.0) for r in v.itertuples()]   # flags only, no photometric rescue
    v = v.drop(columns=["row"])
    v.to_csv(os.path.join(DATA, "vet_table.csv"), index=False)

    c = v[v.role == "candidate"]
    k = v[v.role == "control"]
    print(f"  cones read: {v.cone_ok.sum()}/{len(v)}; inside SDSS imaging: {v.in_sdss.sum()}")
    print(f"  controls: DR16 spectroscopic neighbour in beam {(k.n_spec_same > 0).mean():.3f}; "
          f"counterpart z mismatch {k.oc_z_mismatch.mean():.3f}; any automated flag {(k.auto_reason_any != '').mean():.3f}")
    kc = k[k.auto_clean]
    print(f"  controls passing the catalogue-level flags ({len(kc)}): extra rejection by DR16 cone flags "
          f"{(kc.auto_reason_any != '').mean():.3f}")
    for name, thr in (("50", T50), ("75", T75), ("90", T90)):
        cc = c[c.D < thr]
        s = cc["auto_reason_" + name] == ""
        print(f"  D<{thr:+.3f}: {len(cc)} candidates; pass automated vetting: {s.sum()} "
              f"(CF4 {((cc.src == 'cf4') & s).sum()}, ALFALFA-only {((cc.src == 'alfalfa') & s).sum()})")


if __name__ == "__main__":
    main()