
#!/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, compute_structural_features

def safe_log(s):
    return np.log(pd.Series(s).replace(0, np.nan))

def corr_row(df, x, y):
    sub = df[[x,y]].replace([np.inf,-np.inf], np.nan).dropna()
    if len(sub) < 8:
        return {"predictor": x, "target": y, "N": len(sub), "spearman_r": np.nan, "spearman_p": np.nan}
    r, p = stats.spearmanr(sub[x], sub[y])
    return {"predictor": x, "target": y, "N": len(sub), "spearman_r": float(r), "spearman_p": float(p)}

def cv_model(df, target, features, label):
    sub = df[[target]+features].replace([np.inf,-np.inf], np.nan).dropna()
    if len(sub) < 20:
        return {"target": target, "feature_set": label, "N": len(sub), "cv_r2": np.nan}
    X = sub[features].to_numpy(float)
    y = sub[target].to_numpy(float)
    X = (X - X.mean(axis=0)) / (X.std(axis=0)+1e-12)
    model = Ridge(alpha=30.0)
    cv = KFold(n_splits=5, shuffle=True, random_state=42)
    pred = cross_val_predict(model, X, y, cv=cv)
    sr, sp = stats.spearmanr(pred, y)
    return {
        "target": target, "feature_set": label, "N": len(sub),
        "cv_r2": float(r2_score(y, pred)),
        "mae": float(mean_absolute_error(y, pred)),
        "rmse": float(np.sqrt(mean_squared_error(y, pred))),
        "spearman_pred_true": float(sr),
        "spearman_p": float(sp)
    }

def main():
    ap = argparse.ArgumentParser(description="Test baryonic structural prediction of phase state.")
    ap.add_argument("--root", default=".")
    args = ap.parse_args()
    root = Path(args.root)
    cfg = load_config(root)
    outdir = ensure_dir(root / "outputs" / "03_phase_correlations")
    figdir = ensure_dir(root / "figures")

    t1 = pd.read_csv(root / "outputs" / "01_clean_sparc" / "sparc_table1_clean.csv")
    merged = pd.read_csv(root / "outputs" / "01_clean_sparc" / "sparc_merged_clean.csv")
    phase = pd.read_csv(root / "outputs" / "02_phase_tuner" / "phase_tuner_results.csv")

    struct = compute_structural_features(t1, merged, cfg)
    df = phase.merge(struct, on="Galaxy", how="left")
    df["ln_f_gas"] = safe_log(df["f_gas"])
    df["ln_RHI_over_Rdisk"] = safe_log(df["RHI_over_Rdisk"])
    df["ln_RHI_over_Reff"] = safe_log(df["RHI_over_Reff"])
    df["ln_Rlast_over_RHI"] = safe_log(df["Rlast_over_RHI"])

    predictors = [
        "Psi_bar", "Psi_stress", "f_gas", "ln_f_gas", "logSBeff",
        "RHI_over_Rdisk", "RHI_over_Reff", "Reff_over_Rdisk",
        "ln_Rlast_over_RHI", "D_spike"
    ]
    targets = ["B_y_outer_minus_inner_lnPhi", "phase_abs_deg", "chi2_dof"]
    rows = []
    for t in targets:
        for p in predictors:
            if p in df:
                rows.append(corr_row(df, p, t))
    cor = pd.DataFrame(rows).sort_values(["target","spearman_p"], na_position="last")
    cor.to_csv(outdir / "phase_structure_correlations.csv", index=False)

    features = ["ln_f_gas", "logSBeff", "ln_RHI_over_Rdisk", "ln_RHI_over_Reff", "ln_Rlast_over_RHI"]
    cvrows = [
        cv_model(df, "B_y_outer_minus_inner_lnPhi", features, "baryonic_structure"),
        cv_model(df, "phase_abs_deg", features, "baryonic_structure"),
        cv_model(df, "chi2_dof", features, "baryonic_structure"),
    ]
    cv = pd.DataFrame(cvrows)
    cv.to_csv(outdir / "phase_structure_cv_models.csv", index=False)

    # Main figure
    plt.figure(figsize=(7,5))
    sub = df[["Psi_bar","B_y_outer_minus_inner_lnPhi"]].replace([np.inf,-np.inf], np.nan).dropna()
    plt.scatter(sub["Psi_bar"], sub["B_y_outer_minus_inner_lnPhi"], s=18, alpha=0.75)
    if len(sub) > 2:
        m,b = np.polyfit(sub["Psi_bar"], sub["B_y_outer_minus_inner_lnPhi"], 1)
        xs = np.linspace(sub["Psi_bar"].min(), sub["Psi_bar"].max(), 200)
        plt.plot(xs, m*xs+b, linewidth=1.5)
    plt.xlabel("Psi_bar")
    plt.ylabel("B_y = y_outer - y_inner")
    plt.title("Baryonic stress proxy vs radial phase drift")
    plt.tight_layout()
    plt.savefig(figdir / "phase_structure_Psi_bar_vs_By.png", dpi=160)
    plt.close()

    print("Top correlations:")
    print(cor.head(12).to_string(index=False))
    print("\nCV models:")
    print(cv.to_string(index=False))

if __name__ == "__main__":
    main()
