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

def main():
    ap = argparse.ArgumentParser(description="Build locked SPARC-derived RBFT v32 lensing-base predictor.")
    ap.add_argument("--root", default=".")
    args = ap.parse_args()
    root = Path(args.root)
    cfg = load_config(root)
    outdir = ensure_dir(root / "outputs" / "04_lensing_base")
    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")

    # Locked v32 lensing-base formula from the internal v32 base result.
    a = 1.33945
    b1 = 0.06217
    b2 = -0.18632
    b3 = -0.03816
    b4 = 0.54086
    b5 = 0.39392
    df["Y_lens_base"] = (
        a
        + b1*np.log(df["f_gas"].replace(0, np.nan))
        + b2*np.log(df["SBeff"].replace(0, np.nan))
        + b3*np.log(df["RHI_over_Rdisk"].replace(0, np.nan))
        + b4*np.log(df["RHI_over_Reff"].replace(0, np.nan))
        + b5*np.log(df["Rlast_over_RHI"].replace(0, np.nan))
    )
    df["A_lens_base"] = np.exp(df["Y_lens_base"])
    q1, q2 = df["Y_lens_base"].quantile([1/3, 2/3])
    df["lensing_base_class"] = pd.cut(
        df["Y_lens_base"],
        bins=[-np.inf, q1, q2, np.inf],
        labels=["predicted_low", "predicted_mid", "predicted_high"]
    )
    df.to_csv(outdir / "v32_lensing_base_by_galaxy.csv", index=False)

    stats = df.groupby("lensing_base_class", observed=False).agg(
        N=("Galaxy","size"),
        median_Y_lens_base=("Y_lens_base","median"),
        median_A_lens_base=("A_lens_base","median"),
        median_f_gas=("f_gas","median"),
        median_RHI_over_Rdisk=("RHI_over_Rdisk","median"),
        median_RHI_over_Reff=("RHI_over_Reff","median"),
        median_SBeff=("SBeff","median")
    ).reset_index()
    stats.to_csv(outdir / "v32_lensing_base_class_stats.csv", index=False)

    plt.figure(figsize=(7,5))
    plt.hist(df["A_lens_base"].dropna(), bins=30)
    plt.xlabel("A_lens_base = exp(Y_lens_base)")
    plt.ylabel("galaxies")
    plt.title("RBFT v32 lensing-base amplification distribution")
    plt.tight_layout()
    plt.savefig(figdir / "v32_lensing_base_amplification_distribution.png", dpi=160)
    plt.close()

    print(stats.to_string(index=False))

if __name__ == "__main__":
    main()
