
#!/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 scipy import stats
from sklearn.linear_model import Ridge
from sklearn.model_selection import KFold, cross_val_predict
from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error
from rbft_lib import load_config, ensure_dir, v_rbft_kms_from_mb_msun, G_SI, MSUN_KG

def main():
    ap = argparse.ArgumentParser(description="Run real weak/strong lensing pilot checks from curated input tables.")
    ap.add_argument("--root", default=".")
    args = ap.parse_args()
    root = Path(args.root)
    cfg = load_config(root)
    outdir = ensure_dir(root / "outputs" / "05_real_lensing_pilots")
    figdir = ensure_dir(root / "figures")
    aTheta = cfg["aTheta_m_s2"]

    # Weak-lensing BTFR
    weak = pd.read_csv(root / "data" / "external" / "mistele2024_weak_lensing_btfr_input.csv")
    weak["Mb_Msun"] = 10**weak["log10_Mb_Msun"]
    weak["V_RBFT_locked_kms"] = v_rbft_kms_from_mb_msun(weak["Mb_Msun"], aTheta)
    weak["delta_log10_V"] = np.log10(weak["Vflat_kms"]) - np.log10(weak["V_RBFT_locked_kms"])
    weak["V_ratio_obs_to_RBFT"] = weak["Vflat_kms"] / weak["V_RBFT_locked_kms"]
    weak["aTheta_implied_m_s2"] = (weak["Vflat_kms"]*1000)**4 / (G_SI * weak["Mb_Msun"] * MSUN_KG)
    weak["aTheta_over_locked"] = weak["aTheta_implied_m_s2"] / aTheta
    weak.to_csv(outdir / "weak_lensing_btfr_recomputed.csv", index=False)

    rows = []
    for sample, g in weak.groupby("sample"):
        if len(g) >= 2:
            slope, intercept = np.polyfit(g["log10_Mb_Msun"], np.log10(g["Vflat_kms"]), 1)
        else:
            slope, intercept = np.nan, np.nan
        rows.append({
            "sample": sample,
            "N": len(g),
            "median_aTheta_over_locked": float(g["aTheta_over_locked"].median()),
            "rms_delta_log10_V": float(np.sqrt(np.mean(g["delta_log10_V"]**2))),
            "median_V_ratio_obs_to_RBFT": float(g["V_ratio_obs_to_RBFT"].median()),
            "slope_logV_vs_logMb": float(slope)
        })
    weak_sum = pd.DataFrame(rows).sort_values("sample")
    weak_sum.to_csv(outdir / "weak_lensing_btfr_summary.csv", index=False)

    plt.figure(figsize=(7,5))
    plt.scatter(weak["log10_Mb_Msun"], np.log10(weak["Vflat_kms"]), label="weak lensing rows")
    xs = np.linspace(weak["log10_Mb_Msun"].min(), weak["log10_Mb_Msun"].max(), 200)
    vpred = v_rbft_kms_from_mb_msun(10**xs, aTheta)
    plt.plot(xs, np.log10(vpred), label="locked RBFT", linewidth=2)
    plt.xlabel("log10 baryonic mass [Msun]")
    plt.ylabel("log10 Vflat [km/s]")
    plt.title("Weak-lensing BTFR vs locked RBFT")
    plt.legend()
    plt.tight_layout()
    plt.savefig(figdir / "weak_lensing_btfr_vs_locked_rbft.png", dpi=160)
    plt.close()

    # Strong-lensing pilot
    strong = pd.read_csv(root / "data" / "external" / "slacs_bells_strong_lensing_pilot_input.csv")
    corr_rows = []
    pairs = [
        ("logMstar_Salpeter", "logM_E"),
        ("logSigma_star_Re", "R_E_over_R_eff"),
        ("logSigma_star_Re", "Y_E_enclosed_ln"),
        ("logMstar_enclosed_deV", "logM_E")
    ]
    for x,y in pairs:
        sub = strong[[x,y]].replace([np.inf,-np.inf], np.nan).dropna()
        r,p = stats.spearmanr(sub[x], sub[y]) if len(sub) >= 5 else (np.nan, np.nan)
        corr_rows.append({"predictor": x, "target": y, "N": len(sub), "spearman_r": r, "spearman_p": p})
    corr = pd.DataFrame(corr_rows)
    corr.to_csv(outdir / "strong_lensing_pilot_correlations.csv", index=False)

    features = ["logMstar_Salpeter", "logSigma_star_Re"]
    target = "logM_E"
    sub = strong[[target]+features].dropna()
    X = sub[features].to_numpy(float)
    y = sub[target].to_numpy(float)
    X = (X-X.mean(axis=0))/(X.std(axis=0)+1e-12)
    pred = cross_val_predict(Ridge(alpha=10.0), X, y, cv=KFold(n_splits=5, shuffle=True, random_state=42))
    sr, sp = stats.spearmanr(pred, y)
    cv = pd.DataFrame([{
        "target": target, "features": ",".join(features), "N": len(sub),
        "cv_r2": r2_score(y, pred),
        "mae": mean_absolute_error(y, pred),
        "rmse": np.sqrt(mean_squared_error(y, pred)),
        "spearman_pred_true": sr,
        "spearman_p": sp
    }])
    cv.to_csv(outdir / "strong_lensing_pilot_cv.csv", index=False)

    plt.figure(figsize=(7,5))
    plt.scatter(strong["logMstar_Salpeter"], strong["logM_E"], alpha=0.85)
    plt.xlabel("log10 stellar mass Salpeter")
    plt.ylabel("log10 Einstein mass")
    plt.title("Strong-lensing pilot: Einstein mass vs stellar mass")
    plt.tight_layout()
    plt.savefig(figdir / "strong_lensing_logME_vs_logMstar.png", dpi=160)
    plt.close()

    print("Weak lensing summary:")
    print(weak_sum.to_string(index=False))
    print("\nStrong lensing correlations:")
    print(corr.to_string(index=False))
    print("\nStrong lensing CV:")
    print(cv.to_string(index=False))

if __name__ == "__main__":
    main()
