
#!/usr/bin/env python3
from __future__ import annotations
import argparse, json
from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from rbft_lib import load_config, ensure_dir, fit_phase_profile, inner_outer_median, phase_shift_to_reference

def classify(theta, chi2, B_y, spike_density, n_spikes):
    if n_spikes >= 2 and spike_density >= 0.08:
        return "phase_switching_unstable"
    if chi2 > 1.5:
        return "hard_residual"
    if np.isfinite(theta) and abs(theta) <= 30:
        return "near_aligned"
    if np.isfinite(B_y) and B_y <= -0.15:
        return "boundary_decaying"
    if np.isfinite(B_y) and B_y >= 0.15:
        return "radial_rise"
    if np.isfinite(theta) and abs(theta) >= 150:
        return "extreme_phase"
    if np.isfinite(theta) and abs(theta) >= 90:
        return "strong_phase_drift"
    return "phase_drifted"

def phase_bin(theta):
    if not np.isfinite(theta):
        return "unmeasured"
    a = abs(theta)
    sign = "+" if theta >= 0 else "-"
    if a <= 30:
        return f"{sign}0_30_aligned"
    if a <= 90:
        return f"{sign}30_90_mild"
    if a <= 150:
        return f"{sign}90_150_strong"
    return f"{sign}150_180_extreme"

def main():
    ap = argparse.ArgumentParser(description="Run the RBFT Galaxy Phase Tuner on clean SPARC data.")
    ap.add_argument("--root", default=".")
    args = ap.parse_args()
    root = Path(args.root)
    cfg = load_config(root)
    outdir = ensure_dir(root / "outputs" / "02_phase_tuner")
    figdir = ensure_dir(root / "figures")

    merged_path = root / "outputs" / "01_clean_sparc" / "sparc_merged_clean.csv"
    if not merged_path.exists():
        raise FileNotFoundError("Run 01_make_clean_sparc.py first.")
    df = pd.read_csv(merged_path)
    df = df[np.isfinite(df["y_lnPhi_ML050_unsigned"]) & np.isfinite(df["sigma_y_floor006"])].copy()

    ref_name = cfg.get("phase_reference_galaxy", "NGC7331")
    if ref_name not in set(df["Galaxy"]):
        ref_name = df["Galaxy"].value_counts().index[0]
    ref = df[df["Galaxy"] == ref_name].sort_values("R").copy()
    L_ref = float(np.log(ref["R"].max()/ref["R"].min())) if len(ref) >= 2 else np.nan

    results = []
    point_rows = []
    for gal, g in df.groupby("Galaxy"):
        g = g.sort_values("R").copy()
        if len(g) < 6:
            continue
        best, fits = fit_phase_profile(g)
        if best is None:
            continue
        inner_y, outer_y, B_y = inner_outer_median(g)
        delta, theta = phase_shift_to_reference(g, ref, L_ref)
        residual = (best["y"] - best["yhat"]) / best["sigma"]
        n_spikes = int(np.sum(np.abs(residual) > 3.0))
        spike_density = float(n_spikes / len(residual))
        cls = classify(theta, best["chi2_dof"], B_y, spike_density, n_spikes)

        results.append({
            "Galaxy": gal,
            "N_points": int(best["n"]),
            "phase_reference": ref_name,
            "L_ref": L_ref,
            "best_degree": int(best["degree"]),
            "chi2": float(best["chi2"]),
            "dof": int(best["dof"]),
            "chi2_dof": float(best["chi2_dof"]),
            "inside_tolerance_chi2_le_1": bool(best["chi2_dof"] <= 1.0),
            "soft_inside_chi2_le_1p5": bool(best["chi2_dof"] <= 1.5),
            "B_y_outer_minus_inner_lnPhi": float(B_y),
            "inner_y_median": float(inner_y),
            "outer_y_median": float(outer_y),
            "delta_lnR_to_ref": delta,
            "theta_deg": theta,
            "phase_abs_deg": abs(theta) if np.isfinite(theta) else np.nan,
            "phase_bin": phase_bin(theta),
            "D_spike": spike_density,
            "n_spikes": n_spikes,
            "phase_class": cls
        })

        # save point-level residuals
        for x, y, sig, yh, rr in zip(best["x"], best["y"], best["sigma"], best["yhat"], residual):
            point_rows.append({
                "Galaxy": gal, "lnR": x, "y_lnPhi": y, "sigma_y": sig,
                "y_model": yh, "residual_sigma": rr,
                "is_spike_3sigma": bool(abs(rr)>3)
            })

    res = pd.DataFrame(results).sort_values("Galaxy")
    pts = pd.DataFrame(point_rows)
    res.to_csv(outdir / "phase_tuner_results.csv", index=False)
    pts.to_csv(outdir / "phase_tuner_point_residuals.csv", index=False)

    summary = {
        "galaxies_analyzed": int(len(res)),
        "reference_galaxy": ref_name,
        "inside_tolerance_chi2_le_1": int(res["inside_tolerance_chi2_le_1"].sum()),
        "soft_inside_chi2_le_1p5": int(res["soft_inside_chi2_le_1p5"].sum()),
        "median_chi2_dof": float(res["chi2_dof"].median()),
        "median_abs_phase_deg": float(res["phase_abs_deg"].median()),
        "median_D_spike": float(res["D_spike"].median()),
        "phase_class_counts": res["phase_class"].value_counts().to_dict()
    }
    (outdir / "phase_tuner_summary.json").write_text(json.dumps(summary, indent=2))

    plt.figure(figsize=(8,5))
    res["theta_deg"].dropna().hist(bins=np.linspace(-180,180,37))
    for v in [-150,-90,-30,0,30,90,150]:
        plt.axvline(v, linewidth=0.8, linestyle="--")
    plt.xlabel("phase angle theta [deg]")
    plt.ylabel("galaxies")
    plt.title("RBFT phase-angle distribution")
    plt.tight_layout()
    plt.savefig(figdir / "phase_angle_distribution.png", dpi=160)
    plt.close()

    print(json.dumps(summary, indent=2))

if __name__ == "__main__":
    main()
    import sys as _sys, os as _os
    _sys.stdout.flush()
    _sys.stderr.flush()
    _os._exit(0)
