
from __future__ import annotations

import json
from pathlib import Path
import numpy as np
import pandas as pd

G_SI = 6.67430e-11
MSUN_KG = 1.98847e30
KM_S = 1000.0

def load_config(root: Path) -> dict:
    with open(root / "config.json", "r") as f:
        return json.load(f)

def ensure_dir(path: Path) -> Path:
    path.mkdir(parents=True, exist_ok=True)
    return path

def wrap180(angle_deg):
    return ((angle_deg + 180.0) % 360.0) - 180.0

def read_sparc_table1(path: Path) -> pd.DataFrame:
    cols1 = [
        "Galaxy", "T", "D", "e_D", "f_D", "Inc", "e_Inc",
        "L36", "e_L36", "Reff", "SBeff", "Rdisk", "SBdisk0",
        "MHI", "RHI", "Vflat", "e_Vflat", "Q", "Ref"
    ]
    t1 = pd.read_csv(path, sep=r"\s+", skiprows=98, names=cols1, engine="python")
    t1["Galaxy"] = t1["Galaxy"].astype(str).str.strip()
    return t1

def read_sparc_table2(path: Path) -> pd.DataFrame:
    cols2 = [
        "Galaxy", "D", "R", "Vobs", "e_Vobs",
        "Vgas", "Vdisk", "Vbul", "SBdisk", "SBbul"
    ]
    t2 = pd.read_csv(path, sep=r"\s+", skiprows=25, names=cols2, engine="python")
    t2["Galaxy"] = t2["Galaxy"].astype(str).str.strip()
    return t2

def add_baryonic_phase_columns(df: pd.DataFrame, ml_star=0.5, sigma_floor=0.06) -> pd.DataFrame:
    out = df.copy()
    # SPARC Vgas can be signed. The baseline detector uses unsigned gas support.
    out["Vbar2_ML050_unsigned"] = (
        out["Vgas"].astype(float)**2
        + ml_star * out["Vdisk"].astype(float)**2
        + ml_star * out["Vbul"].astype(float)**2
    )
    valid = (out["Vobs"] > 0) & (out["Vbar2_ML050_unsigned"] > 0) & (out["R"] > 0)
    out["Phi_ML050_unsigned"] = np.nan
    out.loc[valid, "Phi_ML050_unsigned"] = out.loc[valid, "Vobs"]**2 / out.loc[valid, "Vbar2_ML050_unsigned"]
    out["y_lnPhi_ML050_unsigned"] = np.log(out["Phi_ML050_unsigned"])
    out["sigma_y_floor006"] = np.sqrt((2*out["e_Vobs"]/out["Vobs"])**2 + sigma_floor**2)
    out.loc[~np.isfinite(out["sigma_y_floor006"]), "sigma_y_floor006"] = np.nan
    out["lnR"] = np.log(out["R"])
    if "Rdisk" in out.columns:
        out["R_over_Rdisk"] = out["R"] / out["Rdisk"]
        out["lnR_over_Rdisk"] = np.log(out["R_over_Rdisk"])
    if "RHI" in out.columns:
        out["R_over_RHI"] = out["R"] / out["RHI"].replace(0, np.nan)
        out["lnR_over_RHI"] = np.log(out["R_over_RHI"])
    return out

def weighted_polyfit_predict(x, y, sigma, degree):
    mask = np.isfinite(x) & np.isfinite(y) & np.isfinite(sigma) & (sigma > 0)
    x = np.asarray(x)[mask]
    y = np.asarray(y)[mask]
    sigma = np.asarray(sigma)[mask]
    n = len(x)
    if n < degree + 2:
        return None
    w = 1.0 / sigma
    coeff = np.polyfit(x, y, degree, w=w)
    yhat = np.polyval(coeff, x)
    dof = max(n - (degree + 1), 1)
    chi2 = float(np.sum(((y-yhat)/sigma)**2))
    return {
        "coeff": coeff,
        "x": x,
        "y": y,
        "sigma": sigma,
        "yhat": yhat,
        "chi2": chi2,
        "dof": dof,
        "chi2_dof": chi2/dof,
        "n": n
    }

def fit_phase_profile(g: pd.DataFrame):
    x = g["lnR"].to_numpy(float)
    y = g["y_lnPhi_ML050_unsigned"].to_numpy(float)
    sig = g["sigma_y_floor006"].to_numpy(float)
    fits = []
    for deg in [1,2,3]:
        res = weighted_polyfit_predict(x, y, sig, deg)
        if res is not None:
            res["degree"] = deg
            fits.append(res)
    if not fits:
        return None, []
    simple = [f for f in fits if f["chi2_dof"] <= 1.0]
    if simple:
        best = sorted(simple, key=lambda z: (z["degree"], z["chi2_dof"]))[0]
    else:
        best = sorted(fits, key=lambda z: z["chi2_dof"])[0]
    return best, fits

def inner_outer_median(g: pd.DataFrame, value_col: str = "y_lnPhi_ML050_unsigned"):
    gg = g[[value_col, "R"]].dropna().sort_values("R")
    if len(gg) < 6:
        return np.nan, np.nan, np.nan
    k = max(len(gg)//3, 1)
    inner = float(gg.iloc[:k][value_col].median())
    outer = float(gg.iloc[-k:][value_col].median())
    return inner, outer, outer-inner

def phase_shift_to_reference(g: pd.DataFrame, ref: pd.DataFrame, L_ref: float, grid=None):
    if grid is None:
        grid = np.linspace(-0.5, 0.5, 201)
    # Use x = ln(R/Rdisk) where available.
    xg = g["lnR_over_Rdisk"].to_numpy(float) if "lnR_over_Rdisk" in g else g["lnR"].to_numpy(float)
    yg = g["y_lnPhi_ML050_unsigned"].to_numpy(float)
    xr = ref["lnR_over_Rdisk"].to_numpy(float) if "lnR_over_Rdisk" in ref else ref["lnR"].to_numpy(float)
    yr = ref["y_lnPhi_ML050_unsigned"].to_numpy(float)
    mg = np.isfinite(xg) & np.isfinite(yg)
    mr = np.isfinite(xr) & np.isfinite(yr)
    xg, yg = xg[mg], yg[mg]
    xr, yr = xr[mr], yr[mr]
    if len(xg) < 5 or len(xr) < 5 or not np.isfinite(L_ref) or L_ref <= 0:
        return np.nan, np.nan
    order = np.argsort(xr)
    xr, yr = xr[order], yr[order]
    scores = []
    for d in grid:
        xs = xg + d
        lo, hi = max(xs.min(), xr.min()), min(xs.max(), xr.max())
        m = (xs >= lo) & (xs <= hi)
        if m.sum() < 5:
            scores.append(np.inf)
            continue
        yr_i = np.interp(xs[m], xr, yr)
        # Compare shapes after removing offsets. This extracts phase shift more than normalization.
        a = yg[m] - np.nanmedian(yg[m])
        b = yr_i - np.nanmedian(yr_i)
        scores.append(float(np.nanmean((a-b)**2)))
    j = int(np.nanargmin(scores))
    delta = float(grid[j])
    theta = wrap180(360.0 * delta / L_ref)
    return delta, theta

def compute_structural_features(t1: pd.DataFrame, radial: pd.DataFrame | None = None, cfg: dict | None = None) -> pd.DataFrame:
    cfg = cfg or {}
    ml = cfg.get("stellar_ML_36", 0.5)
    he = cfg.get("helium_factor_HI", 1.33)
    out = t1.copy()
    out["Mstar_proxy"] = ml * out["L36"]
    out["Mgas_proxy"] = he * out["MHI"]
    out["Mb_proxy"] = out["Mstar_proxy"] + out["Mgas_proxy"]
    out["f_gas"] = out["Mgas_proxy"] / out["Mb_proxy"].replace(0, np.nan)
    out["RHI_over_Rdisk"] = out["RHI"] / out["Rdisk"].replace(0, np.nan)
    out["RHI_over_Reff"] = out["RHI"] / out["Reff"].replace(0, np.nan)
    out["Reff_over_Rdisk"] = out["Reff"] / out["Rdisk"].replace(0, np.nan)
    out["logSBeff"] = np.log(out["SBeff"].replace(0, np.nan))
    out["Psi_bar"] = (
        np.log(out["RHI_over_Rdisk"].replace(0, np.nan)) +
        np.log(out["RHI_over_Reff"].replace(0, np.nan)) +
        out["logSBeff"] -
        np.log(out["f_gas"].replace(0, np.nan))
    )
    if radial is not None:
        rlast = radial.groupby("Galaxy")["R"].max().rename("R_last")
        out = out.merge(rlast, on="Galaxy", how="left")
        out["Rlast_over_RHI"] = out["R_last"] / out["RHI"].replace(0, np.nan)
        out["Psi_stress"] = out["Psi_bar"] - np.log(out["Rlast_over_RHI"].replace(0, np.nan))
    return out

def v_rbft_kms_from_mb_msun(Mb_msun, aTheta):
    Mb_kg = np.asarray(Mb_msun, dtype=float) * MSUN_KG
    v_ms = (G_SI * Mb_kg * aTheta)**0.25
    return v_ms / KM_S
