"""Step 3: fit the Tully-Fisher relations of the ALFALFA-SDSS gold sample, compute residuals.

Relations (forward fits, luminosity at fixed velocity, robust quadratic in log V):
   i-band luminosity  (SDSS cmodel; Durbala+2020 extinction corrections)
   W1 luminosity      (unWISE; expressed as the McGaugh stellar mass, M/L = 0.45)
   baryonic mass      (M*_W1 + 1.33 M_HI)
Output : data/alfalfa_resid.csv (gold + silver with residuals), data/alfalfa_tf_fits.csv
"""
import os
import numpy as np
import pandas as pd
from common import DATA, robust_polyfit, mad_sigma, running_median_fit

X0 = 2.2
BANDS = {"i": "logLi", "W1": "logLW1", "bar": "logMbar"}
VBINS = [80, 100, 120, 150, 200, 250, 700]


def fit_all(g):
    fits = {}
    for k, col in BANDS.items():
        p, sig = robust_polyfit(g.logV.values - X0, g[col].values, deg=2)
        fits[k] = (p, sig)
    return fits


def add_resid(d, fits):
    d = d.copy()
    for k, col in BANDS.items():
        p, _ = fits[k]
        d["r_" + k] = d[col] - np.polyval(p, d.logV - X0)
        d["reo_" + k] = d[col] - np.polyval(p, d.logVmin - X0)   # edge-on assumption (max possible residual)
    d["D"] = np.maximum(d.r_i, d.r_W1)          # least-underluminous of the two bands
    d["Deo"] = np.maximum(d.reo_i, d.reo_W1)
    return d


def add_running_median_resid(d, g):
    """Alternative residuals about a 15-bin running median of the gold sample (fit-form check)."""
    d = d.copy()
    for k in ("i", "W1"):
        f = running_median_fit(g.logV.values, g[BANDS[k]].values)
        d["rm_" + k] = d[BANDS[k]] - f(d.logV.values)
    d["D_rm"] = np.maximum(d.rm_i, d.rm_W1)
    return d


def main():
    d = pd.read_csv(os.path.join(DATA, "alfalfa_all.csv"), low_memory=False)
    g = d[d.gold]
    fits = fit_all(g)
    rows = []
    for k in BANDS:
        p, sig = fits[k]
        rows.append(dict(relation=k, c2=p[0], c1=p[1], c0=p[2], x0=X0, robust_sigma_dex=sig, N=len(g)))
        print(f"  TF[{k:3s}]: y = {p[2]:.3f} + {p[1]:.3f} x + {p[0]:.3f} x^2 (x = logV-{X0});  robust sigma = {sig:.3f} dex")
    pd.DataFrame(rows).to_csv(os.path.join(DATA, "alfalfa_tf_fits.csv"), index=False)

    s = add_resid(d[d.gold | d.silver], fits)
    s = add_running_median_resid(s, g)
    s.to_csv(os.path.join(DATA, "alfalfa_resid.csv"), index=False)

    g = s[s.gold]
    rows = []
    print("  scatter and median residual by rotation velocity (gold):")
    for lo, hi in zip(VBINS[:-1], VBINS[1:]):
        b = g[(g.V >= lo) & (g.V < hi)]
        row = dict(Vlo=lo, Vhi=hi, N=len(b))
        for k in BANDS:
            row[f"med_{k}"] = b["r_" + k].median()
            row[f"sig_{k}"] = mad_sigma(b["r_" + k])
        row["sig_D"] = mad_sigma(b.D)
        row["frac_D_lt_-0.3"] = (b.D < -0.301).mean()
        row["frac_D_lt_-0.6"] = (b.D < -0.602).mean()
        row["frac_D_lt_-1.0"] = (b.D < -1.0).mean()
        row["frac_ri_gt_0.6"] = (b.r_i > 0.602).mean()
        row["median_fgas"] = b.fgas.median()
        rows.append(row)
        print(f"    V {lo:3d}-{hi:3d}: N={len(b):4d}  med(i,W1,bar)=({row['med_i']:+.3f},{row['med_W1']:+.3f},{row['med_bar']:+.3f})"
              f"  sig=({row['sig_i']:.3f},{row['sig_W1']:.3f},{row['sig_bar']:.3f})  fgas={row['median_fgas']:.2f}")
    pd.DataFrame(rows).to_csv(os.path.join(DATA, "alfalfa_scatter_by_V.csv"), index=False)
    print(f"  corr(r_i, r_W1) = {np.corrcoef(g.r_i, g.r_W1)[0, 1]:.3f}")


if __name__ == "__main__":
    main()