"""Step 6: Cosmicflows-4 Tully-Fisher catalogue (Kourkchi et al. 2020) as a second, independent sample.

CF4 has by-eye inclinations, re-measured SDSS/WISE aperture photometry and vetted HI widths.
We do NOT use the CF4 Tully-Fisher distances (they assume the relation we are testing); distances are
redshift distances from V_cmb with the CF4 cosmological correction, H0 = 75 (only the zero point).
Outputs: data/cf4_resid.csv, data/cf4_tf_fits.csv, data/cf4_scatter_by_V.csv
"""
import os
import numpy as np
import pandas as pd
from scipy.spatial import cKDTree
from common import (RAW, DATA, C_KMS, cmb_velocity, robust_polyfit, mad_sigma, radec_to_xyz, chord_for_arcsec,
                    running_median_fit)

H0_CF4 = 75.0
X0 = 2.2
MSUN_I, MSUN_W1 = 4.53, 5.94   # AB absolute magnitudes of the Sun (SDSS i ; WISE W1 = 3.24 Vega + 2.699)


def build():
    c = pd.read_csv(os.path.join(RAW, "cf4_t1.csv"))
    c = c.rename(columns={"i*mag": "i_star", "W1*mag": "W1_star", "_RA": "ra", "_DE": "dec", "DM": "in_table4",
                          "(b/a)s": "ba_s", "(b/a)W": "ba_w"})
    c["in_table4"] = c.in_table4.astype(str).str.strip().eq("*")
    n_all = len(c)
    c = c[c.ra.notna()].copy()
    if len(c) < n_all:      # VizieR gives no position for a few rows; they cannot be cross-matched or vetted
        print(f"  CF4 rows without VizieR coordinates, dropped: {n_all - len(c)} of {n_all}")
    c["Vcmb"] = cmb_velocity(c.ra.values, c.dec.values, c.HRV.values)
    z = c.Vcmb / C_KMS
    q0, j0 = -0.595, 1.0
    f = 1 + 0.5 * (1 - q0) * z - (1.0 / 6) * (1 - q0 - 3 * q0 ** 2 + j0) * z ** 2
    c["Dist"] = f * c.Vcmb / H0_CF4
    c["DMz"] = 5 * np.log10(c.Dist.where(c.Dist > 0)) + 25
    c["logV"] = c.logWmxi - np.log10(2.0)
    c["V"] = 10 ** c.logV
    ok_i = c.i_star.between(5, 25) & (c.QSflag >= 3)
    ok_w = c.W1_star.between(5, 25) & (c.QWflag >= 3)
    c["logLi"] = np.where(ok_i, -0.4 * (c.i_star - c.DMz - MSUN_I), np.nan)
    c["logLW1"] = np.where(ok_w, -0.4 * (c.W1_star - c.DMz - MSUN_W1), np.nan)
    flux = 10 ** (-0.4 * (c.m21 - 17.40))
    c["logMHI"] = np.log10(2.356e5 * c.Dist ** 2 * flux)
    # cross-identify with a.100 optical counterparts
    a = pd.read_csv(os.path.join(DATA, "alfalfa_all.csv"), low_memory=False,
                    usecols=["AGC", "RAO", "DEO", "gold", "silver", "SNR"])
    a = a[a.RAO.notna()]
    tr = cKDTree(radec_to_xyz(a.RAO, a.DEO))
    dist, idx = tr.query(radec_to_xyz(c.ra, c.dec))
    hit = dist < chord_for_arcsec(20.0)
    c["AGC"] = np.where(hit, a.AGC.values[idx], -1)
    c["in_a100"] = hit
    c["in_alfalfa_gold"] = hit & a.gold.values[idx]
    c["sample"] = (c.Vcmb >= 3000) & (c.V >= 80) & (c.logLi.notna() | c.logLW1.notna())
    return c


def main():
    c = build()
    s = c[c["sample"]].copy()
    rows = []
    fits = {}
    for k, col in (("i", "logLi"), ("W1", "logLW1")):
        m = s[col].notna()
        p, sig = robust_polyfit(s.logV[m].values - X0, s[col][m].values, deg=2)
        fits[k] = p
        s["r_" + k] = s[col] - np.polyval(p, s.logV - X0)
        s["rm_" + k] = s[col] - running_median_fit(s.logV[m].values, s[col][m].values)(s.logV.values)
        rows.append(dict(relation=k, c2=p[0], c1=p[1], c0=p[2], x0=X0, robust_sigma_dex=sig, N=int(m.sum())))
        print(f"  CF4 TF[{k:2s}]: N={m.sum():5d}  y = {p[2]:.3f} + {p[1]:.3f} x + {p[0]:.3f} x^2 ; robust sigma {sig:.3f} dex")
    pd.DataFrame(rows).to_csv(os.path.join(DATA, "cf4_tf_fits.csv"), index=False)
    # stellar mass proxy on the W1 scale (0.45 Msun/Lsun,W1 in Vega-based solar units ~ here a constant offset)
    both = s.logLi.notna() & s.logLW1.notna()
    off = (s.logLW1[both] - s.logLi[both]).median()
    s["logLs"] = s.logLW1.fillna(s.logLi + off)
    s["D"] = np.fmax(s.r_i, s.r_W1)      # least-underluminous band; fmax ignores a missing band
    s["D_rm"] = np.fmax(s.rm_i, s.rm_W1)  # same statistic about the 15-bin running-median relation
    s["nband"] = s.r_i.notna().astype(int) + s.r_W1.notna().astype(int)
    # baryonic relation: calibrate the stellar-mass zero point so that the median gas fraction at
    # V=150-200 km/s matches the ALFALFA-SDSS value (McGaugh W1 masses); only residuals matter here
    s["logMs"] = s.logLs + np.log10(0.45)       # 0.45 Msun/Lsun in W1 (McGaugh & Schombert 2015)
    s["logMbar"] = np.log10(10 ** s.logMs + 1.33 * 10 ** s.logMHI)
    p, sig = robust_polyfit(s.logV.values - X0, s.logMbar.values, deg=2)
    s["r_bar"] = s.logMbar - np.polyval(p, s.logV - X0)
    s["fgas"] = 1.33 * 10 ** s.logMHI / 10 ** s.logMbar
    print(f"  CF4 BTFR robust sigma {sig:.3f} dex; i->W1 luminosity offset {off:+.3f} dex")
    # mean surface brightness inside the half-light ellipse (mag/arcsec^2) and its residual at fixed V;
    # positive dmu = lower surface brightness than typical.  Uniform dimming by f raises dmu by -2.5 log(1-f).
    s["dmu"] = np.nan
    for mag, re, ba, out in (("W1cmag", "ReW1", "ba_w", "mu_W1"), ("icmag", "Rei", "ba_s", "mu_i")):   # i overrides W1 in dmu
        ok = s[mag].between(5, 25) & (s[re] > 0.02) & s[ba].between(0.05, 1.0)
        mu = (s[mag] + 2.5 * np.log10(2 * np.pi * (s[re].where(ok) * 60.0) ** 2 * s[ba])).where(ok)
        s[out] = mu                      # absolute value, used for the surface-brightness selection term in s07
        pm, sg = robust_polyfit(s.logV[ok].values - X0, mu[ok].values, deg=1)
        s.loc[ok, "dmu"] = (mu - np.polyval(pm, s.logV - X0))[ok]
        print(f"  CF4 surface brightness-velocity relation ({mag}): N={ok.sum()}, robust scatter {sg:.2f} mag; "
              f"mean surface brightness percentiles 50/98/99.5: {np.nanpercentile(mu, [50, 98, 99.5]).round(2)}")
    s.to_csv(os.path.join(DATA, "cf4_resid.csv"), index=False)

    print(f"  CF4 rows {len(c)}; sample (Vcmb>=3000, V>=80, photometry Q>=3): {len(s)}; "
          f"of which not in ALFALFA gold: {(~s.in_alfalfa_gold).sum()}, not in a.100 at all: {(~s.in_a100).sum()}")
    rows = []
    for lo, hi in [(80, 100), (100, 120), (120, 150), (150, 200), (200, 250), (250, 700)]:
        b = s[(s.V >= lo) & (s.V < hi)]
        rows.append(dict(Vlo=lo, Vhi=hi, N=len(b), med_D=b.D.median(), sig_D=mad_sigma(b.D),
                         frac_lt_03=(b.D < -0.301).mean(), frac_lt_06=(b.D < -0.602).mean(),
                         frac_lt_10=(b.D < -1.0).mean(), frac_gt_06=(b.D > 0.602).mean()))
        print(f"    V {lo:3d}-{hi:3d}: N={len(b):4d} med={b.D.median():+.3f} sig={mad_sigma(b.D):.3f} "
              f"f(<-.3)={rows[-1]['frac_lt_03']:.3f} f(<-.6)={rows[-1]['frac_lt_06']:.4f} f(<-1)={rows[-1]['frac_lt_10']:.4f}"
              f" f(>.6)={rows[-1]['frac_gt_06']:.4f}")
    pd.DataFrame(rows).to_csv(os.path.join(DATA, "cf4_scatter_by_V.csv"), index=False)
    for t in (-0.301, -0.602, -1.0):
        cnd = s[s.D < t]
        print(f"    D<{t:+.3f}: {len(cnd)} (not in ALFALFA gold: {(~cnd.in_alfalfa_gold).sum()}; "
              f"in CF4 distance table: {cnd.in_table4.sum()}; dropped from it: {(~cnd.in_table4).sum()})")
    print(f"    overall fraction of sample dropped from CF4 distance table: {(~s.in_table4).mean():.3f}")


if __name__ == "__main__":
    main()