"""
VMS + NFW + MOND analysis for SPARC rotation curves
Fixed units, safe file loading, correct statistics (Nov 2025 version)
"""

import numpy as np
import pandas as pd
from scipy.optimize import curve_fit
import glob

# =============================================================================
# Constants & units (SPARC standard system)
# =============================================================================
G_KPC = 4.302e-6          # G in kpc (km/s)? / M?   (do NOT apply any further conversion!)
A0 = 1.2e-10              # Standard MOND acceleration in m/s? (fixed for all galaxies)

# =============================================================================
# Data loading Ð handles both 5-column (no bulge) and 6-column SPARC files
# =============================================================================
def load_sparc_galaxy(galaxy_name: str, data_dir: str = './sparc_data/') -> pd.DataFrame:
    pattern = f"{data_dir}*{galaxy_name}*.dat"
    files = glob.glob(pattern)
    if not files:
        raise FileNotFoundError(f"No data file found for {galaxy_name}")
    filepath = files[0]

    raw = pd.read_csv(filepath, delim_whitespace=True, comment='#', header=None)
    if raw.shape[1] == 5:
        raw.columns = ['Rad', 'Vobs', 'errV', 'Vgas', 'Vdisk']
        raw['Vbul'] = 0.0
    elif raw.shape[1] == 6:
        raw.columns = ['Rad', 'Vobs', 'errV', 'Vgas', 'Vdisk', 'Vbul']
    else:
        raise ValueError(f"Unexpected column count in {filepath}")

    data = raw.dropna().reset_index(drop=True)
    print(f"Loaded {galaxy_name}: {len(data)} points")
    return data

# =============================================================================
# Velocity models
# =============================================================================
def vms_velocity(r: np.ndarray, rho0: float, rs: float) -> np.ndarray:
    """VMS (Zhao 1996, n=0) spherical profile Ð Eq. (4) in the paper"""
    v_squared = 4 * np.pi * G_KPC * rho0 * r / (r + rs)
    return np.sqrt(np.abs(v_squared))          # abs just in case of tiny negative floats

def nfw_mass(r: np.ndarray, rho_s: float, r_s: float) -> np.ndarray:
    x = r / r_s
    return 4 * np.pi * rho_s * r_s**3 * (np.log(1 + x) - x / (1 + x))

def nfw_velocity(r: np.ndarray, rho_s: float, r_s: float) -> np.ndarray:
    r_safe = np.maximum(r, 1e-10)  # Guard against r=0 division
    M = nfw_mass(r_safe, rho_s, r_s)
    return np.sqrt(G_KPC * M / r_safe)

def mond_velocity(r: np.ndarray, Vbar: np.ndarray, a0: float = A0) -> np.ndarray:
    """Flat-space interpolating function MOND Ð v? = Vbar? a? r"""
    r_m = r * 3.085677581e19          # kpc ? metres
    Vbar_ms = Vbar * 1e3
    v_ms = (Vbar_ms**2 * a0 * r_m)**0.25
    return v_ms / 1e3

# =============================================================================
# Statistics
# =============================================================================
def calculate_statistics(v_obs, v_model, v_err, n_params: int):
    residuals = v_obs - v_model
    chi2 = np.sum((residuals / v_err)**2)
    dof = len(v_obs) - n_params
    chi2_red = np.nan if dof <= 0 else chi2 / dof
    rms = np.sqrt(np.mean(residuals**2))
    return {'chi2': chi2, 'chi2_red': chi2_red, 'rms': rms, 'dof': dof}

# =============================================================================
# Fitting functions (one free parameter each)
# =============================================================================
def fit_vms(r: np.ndarray, v_obs: np.ndarray, v_err: np.ndarray, rs: float, v_flat: float, v_bar: np.ndarray):
    rho0_guess = v_flat**2 / (4 * np.pi * G_KPC * rs)
    def model(r, rho0):
        v_dm = vms_velocity(r, rho0, rs)
        return np.sqrt(v_dm**2 + v_bar**2)
    popt, pcov = curve_fit(model, r, v_obs, p0=[rho0_guess], sigma=v_err, bounds=(0, np.inf), absolute_sigma=True)
    return popt[0], np.sqrt(pcov[0][0])

def fit_nfw(r: np.ndarray, v_obs: np.ndarray, v_err: np.ndarray, r_s: float, v_flat: float, v_bar: np.ndarray):
    # Rough asymptotic guess: v_flat? Å 4¹ G ?_s r_s? / ln(2) or similar; this is heuristic
    rho_s_guess = v_flat**2 / (4 * np.pi * G_KPC * r_s**2 * np.log(2))
    def model(r, rho_s):
        v_dm = nfw_velocity(r, rho_s, r_s)
        return np.sqrt(v_dm**2 + v_bar**2)
    popt, pcov = curve_fit(model, r, v_obs, p0=[rho_s_guess], sigma=v_err, bounds=(0, np.inf), absolute_sigma=True)
    return popt[0], np.sqrt(pcov[0][0])

# =============================================================================
# Example / single-galaxy test
# =============================================================================
if __name__ == '__main__':
    galaxy = 'UGC02259'          # change as needed
    df = load_sparc_galaxy(galaxy)
    r = df['Rad'].values
    v_obs = df['Vobs'].values
    v_err = np.maximum(df['errV'].values, 0.5)  # Tiny floor to avoid zero errors
    # Baryonic velocity (quadratic sum Ð standard in SPARC papers)
    v_bar = np.sqrt(df['Vgas']**2 + df['Vdisk']**2 + df['Vbul']**2)

    # MOND prediction (0 free parameters)
    v_mond = mond_velocity(r, v_bar)
    mond_stats = calculate_statistics(v_obs, v_mond, v_err, n_params=0)

    # Very rough guesses for scale radii & flat velocity
    v_flat = np.median(v_obs[-5:]) if len(v_obs) >= 5 else np.max(v_obs)  # last few points Å flat part
    rs_guess = np.max(r) / 3.0          # reasonable starting guess

    # VMS fit (1 free parameter: ??)
    rho0, rho0_err = fit_vms(r, v_obs, v_err, rs_guess, v_flat, v_bar)
    v_dm_vms = vms_velocity(r, rho0, rs_guess)
    v_vms = np.sqrt(v_dm_vms**2 + v_bar**2)
    vms_stats = calculate_statistics(v_obs, v_vms, v_err, n_params=1)

    # NFW fit (1 free parameter: ?_s)
    rho_s, rho_s_err = fit_nfw(r, v_obs, v_err, rs_guess, v_flat, v_bar)
    v_dm_nfw = nfw_velocity(r, rho_s, rs_guess)
    v_nfw = np.sqrt(v_dm_nfw**2 + v_bar**2)
    nfw_stats = calculate_statistics(v_obs, v_nfw, v_err, n_params=1)

    print("\nResults for", galaxy)
    print(f"MOND ??_red = {mond_stats['chi2_red']:.3f}  (rms = {mond_stats['rms']:.2f} km/s)")
    print(f"VMS  ??_red = {vms_stats['chi2_red']:.3f}   ?? = {rho0:.2e} ± {rho0_err:.2e} M?/kpc?")
    print(f"NFW  ??_red = {nfw_stats['chi2_red']:.3f}   ?_s = {rho_s:.2e} ± {rho_s_err:.2e} M?/kpc?")