import numpy as np
from scipy.io import savemat
import os
import warnings
from scipy.io import loadmat
from scipy.interpolate import PchipInterpolator
import matplotlib.pyplot as plt
from scipy.stats import norm, skew
from matplotlib.patches import Rectangle
from matplotlib.colors import Normalize
from matplotlib.cm import get_cmap


# ===== CONSTANTS =====
Compliance = 377  # Free Parameter : Vacuum Compliance = 377
G_kpc = 4.300917270437e-06  
solar_mass_kg = 1.9885e30

# ===== HELPER FUNCTIONS =====
def read_rotation_data(filename):
    galaxy_names = []
    observed_radius = []
    observed_speed = []
    e_Vobs = []
    observed_Vgas = []
    e_Vgas = []

    with open(filename, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#') or line.startswith('ID') or line.startswith('='):
                continue
            parts = line.split()
            if len(parts) < 6:
                continue

            galaxy_name = parts[0]
            radius = float(parts[1])
            Vobs = float(parts[2])
            eVobs = float(parts[3])
            Vgas = float(parts[4])
            eVgas = 0.0  

            # Find existing galaxy or create new entry
            try:
                idx = galaxy_names.index(galaxy_name)
            except ValueError:
                idx = len(galaxy_names)
                galaxy_names.append(galaxy_name)
                observed_radius.append([])
                observed_speed.append([])
                e_Vobs.append([])
                observed_Vgas.append([])
                e_Vgas.append([])

            observed_radius[idx].append(radius)
            observed_speed[idx].append(Vobs)
            e_Vobs[idx].append(eVobs)
            observed_Vgas[idx].append(Vgas)
            e_Vgas[idx].append(eVgas)

    # Convert to numpy arrays for vectorized math
    observed_radius = [np.array(r) for r in observed_radius]
    observed_speed = [np.array(v) for v in observed_speed]
    e_Vobs = [np.array(e) for e in e_Vobs]
    observed_Vgas = [np.array(v) for v in observed_Vgas]
    e_Vgas = [np.array(e) for e in e_Vgas]

    return galaxy_names, observed_radius, observed_speed, e_Vobs, observed_Vgas, e_Vgas


def read_mass_data(filename):
    gs_galaxy_names = []
    M2L, e_M2L, M2Lb = [], [], []
    L36, e_L36, Lbulge = [], [], []
    Rdisk, Reff, Distance, e_D = [], [], [], []
    Mbar, e_Mbar = [], []

    with open(filename, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#') or line.startswith('ID') or line.startswith('='):
                continue
            parts = line.split()
            if len(parts) < 21:
                continue

  
            gs_galaxy_names.append(parts[0])
            M2L.append(float(parts[1]))
            e_M2L.append(float(parts[2]))
            M2Lb.append(float(parts[3]))
            Distance.append(float(parts[5]))    
            e_D.append(float(parts[6]))       
            L36.append(float(parts[9]))         
            e_L36.append(float(parts[10]))      
            Lbulge.append(float(parts[11]))     
            Reff.append(float(parts[12]))       
            Rdisk.append(float(parts[16]))      
            Mbar.append(float(parts[19]))       
            e_Mbar.append(float(parts[20]))     

    # Convert to numpy arrays
    return (gs_galaxy_names, 
            np.array(M2L), np.array(e_M2L), np.array(M2Lb),
            np.array(L36), np.array(e_L36), np.array(Lbulge),
            np.array(Rdisk), np.array(Reff), np.array(Distance), np.array(e_D),
            np.array(Mbar), np.array(e_Mbar))


def compute_enclosed_disk_mass(r, M_disk, e_M_disk, R_disk):
    if R_disk > 0 and M_disk > 0:
        x = r / R_disk
        profile_factor = 1 - (1 + x) * np.exp(-x)
        enc_disk = M_disk * profile_factor
        e_enc_disk = e_M_disk * profile_factor
    else:
        enc_disk = 0.0
        e_enc_disk = 0.0
    return enc_disk, e_enc_disk


def compute_enclosed_bulge_mass(r, M_bulge, e_M_bulge, R_eff):
    if R_eff > 0 and M_bulge > 0:
        a = R_eff / 1.8153
        profile_factor = (r / (r + a))**2
        enc_bulge = M_bulge * profile_factor
        e_enc_bulge = e_M_bulge * profile_factor
    else:
        enc_bulge = 0.0
        e_enc_bulge = 0.0
    return enc_bulge, e_enc_bulge


# ===== MAIN EXECUTION =====
if __name__ == "__main__":
    # 1. Load Data
    galaxy_names, observed_radius, observed_speed, e_Vobs, observed_Vgas, e_Vgas = \
        read_rotation_data('OBSERVED_Rotation.txt')
        
    gs_galaxy_names, M2L, e_M2L, M2Lb, L36, e_L36, Lbulge, Rdisk, Reff, Distance_to_G, e_D, Mbar, e_Mbar = \
        read_mass_data('OBSERVED_Mass.txt')

    # 2. Initialize Storage 
    n_galaxies = len(galaxy_names)
    total_baryonic_mass = [None] * n_galaxies
    total_baryonic_mass_err = [None] * n_galaxies
    total_baryonic_mass_direct = [None] * n_galaxies
    total_baryonic_mass_direct_err = [None] * n_galaxies
    baryonic_mass_ratio = np.zeros(n_galaxies)
    total_gas_mass = [None] * n_galaxies
    total_gas_mass_err = [None] * n_galaxies
    enclosed_masses = [None] * n_galaxies
    enclosed_masses_err = [None] * n_galaxies
    enclosed_gas_masses = [None] * n_galaxies
    enclosed_gas_masses_err = [None] * n_galaxies
    radius_uncertainties = [None] * n_galaxies
    
    gs_Rdisk_ordered = np.zeros(n_galaxies)
    gs_Reff_ordered = np.zeros(n_galaxies)
    M_disk_ordered = np.zeros(n_galaxies)
    M_bulge_ordered = np.zeros(n_galaxies)
    
    All_Vobs = [None] * n_galaxies
    All_e_Vobs = [None] * n_galaxies
    All_Robs = [None] * n_galaxies
    All_e_Robs = [None] * n_galaxies
    All_Enclosed_Masses = [None] * n_galaxies
    All_e_Enclosed_Masses = [None] * n_galaxies

    # 3. Process Each Galaxy
    for g in range(n_galaxies):
        current_galaxy = galaxy_names[g]
        try:
            idx = gs_galaxy_names.index(current_galaxy)
        except ValueError:
            print(f"Warning: {current_galaxy} not found in mass data. Skipping.")
            continue

        # ===== DIRECT BARYONIC MASS =====
        total_baryonic_mass_direct[g] = Mbar[idx] * 1e9
        total_baryonic_mass_direct_err[g] = e_Mbar[idx] * 1e9

        D = Distance_to_G[idx]
        e_D_val = e_D[idx]

        # Radius uncertainties
        R_kpc = observed_radius[g]
        e_R_kpc = R_kpc * (e_D_val / D)
        radius_uncertainties[g] = e_R_kpc * 3.086e+19

        # Stellar component calculations
        L_disk = L36[idx] - Lbulge[idx]
        e_L_disk = np.sqrt(e_L36[idx]**2)

        M_disk = M2L[idx] * L_disk * 1e9
        e_M_disk = M_disk * (e_M2L[idx] / M2L[idx])

        M_bulge = M2Lb[idx] * Lbulge[idx] * 1e9
        e_M_bulge = 0.0

        total_stellar_mass = M_disk + M_bulge
        total_stellar_mass_err = np.sqrt(e_M_disk**2 + e_M_bulge**2)

        gs_Rdisk_ordered[g] = Rdisk[idx]
        gs_Reff_ordered[g] = Reff[idx]
        M_disk_ordered[g] = M_disk
        M_bulge_ordered[g] = M_bulge

        # Enclosed mass loop
        radii = R_kpc
        Vgas_vals = observed_Vgas[g]
        enc_mass = np.zeros_like(radii)
        enc_mass_err = np.zeros_like(radii)
        enc_gas_mass = np.zeros_like(radii)
        enc_gas_mass_err = np.zeros_like(radii)
        fractional_Vgas_error = 0.05

        for r_idx in range(len(radii)):
            r = radii[r_idx]
            Vgas_val = Vgas_vals[r_idx]

            enc_disk, e_enc_disk = compute_enclosed_disk_mass(r, M_disk, e_M_disk, Rdisk[idx])
            enc_bulge, e_enc_bulge = compute_enclosed_bulge_mass(r, M_bulge, e_M_bulge, Reff[idx])

            if Vgas_val > 0:
                M_enc_gas = (Vgas_val**2 * r) / G_kpc
                enc_gas_mass[r_idx] = M_enc_gas
                e_Vgas_val = fractional_Vgas_error * Vgas_val
                e_M_enc_gas = M_enc_gas * (2 * e_Vgas_val / Vgas_val)
                enc_gas_mass_err[r_idx] = e_M_enc_gas
            else:
                enc_gas_mass[r_idx] = 0.0
                enc_gas_mass_err[r_idx] = 0.0

            enc_mass[r_idx] = enc_disk + enc_bulge + enc_gas_mass[r_idx]
            enc_mass_err[r_idx] = np.sqrt(e_enc_disk**2 + e_enc_bulge**2 + enc_gas_mass_err[r_idx]**2)

        # Total gas mass (at largest radius)
        total_gas_mass[g] = enc_gas_mass[-1]
        total_gas_mass_err[g] = enc_gas_mass_err[-1]

        # Total baryonic mass (calculated)
        total_baryonic_mass[g] = total_stellar_mass + total_gas_mass[g]
        total_baryonic_mass_err[g] = np.sqrt(total_stellar_mass_err**2 + total_gas_mass_err[g]**2)

        # Comparison metric
        if total_baryonic_mass_direct[g] > 0:
            baryonic_mass_ratio[g] = total_baryonic_mass[g] / total_baryonic_mass_direct[g]
        else:
            baryonic_mass_ratio[g] = np.nan

        # Store results
        enclosed_masses[g] = enc_mass
        enclosed_masses_err[g] = enc_mass_err
        enclosed_gas_masses[g] = enc_gas_mass
        enclosed_gas_masses_err[g] = enc_gas_mass_err

        # Unit conversions & final storage
        Robs_m = observed_radius[g] * 3.086e+19
        Vobs_m_s = observed_speed[g] * 1e3
        e_Vobs_m_s = e_Vobs[g] * 1e3

        Enclosed_Mass_kg = enclosed_masses[g] * solar_mass_kg
        e_Enclosed_Mass_kg = enclosed_masses_err[g] * solar_mass_kg

        All_Vobs[g] = Vobs_m_s
        All_e_Vobs[g] = e_Vobs_m_s
        All_Robs[g] = Robs_m
        All_e_Robs[g] = radius_uncertainties[g]
        All_Enclosed_Masses[g] = Enclosed_Mass_kg
        All_e_Enclosed_Masses[g] = e_Enclosed_Mass_kg

    # 4. Save to .mat
    save_dict = {
        'galaxy_names': np.array(galaxy_names, dtype=object),
        'total_baryonic_mass': total_baryonic_mass,
        'total_baryonic_mass_err': total_baryonic_mass_err,
        'total_baryonic_mass_direct': total_baryonic_mass_direct,
        'total_baryonic_mass_direct_err': total_baryonic_mass_direct_err,
        'baryonic_mass_ratio': baryonic_mass_ratio,
        'total_gas_mass': total_gas_mass,
        'total_gas_mass_err': total_gas_mass_err,
        'enclosed_masses': enclosed_masses,
        'enclosed_masses_err': enclosed_masses_err,
        'enclosed_gas_masses': enclosed_gas_masses,
        'enclosed_gas_masses_err': enclosed_gas_masses_err,
        'Distance_to_G': Distance_to_G,
        'radius_uncertainties': radius_uncertainties,
        'gs_Rdisk_ordered': gs_Rdisk_ordered,
        'gs_Reff_ordered': gs_Reff_ordered,
        'M_disk_ordered': M_disk_ordered,
        'M_bulge_ordered': M_bulge_ordered,
        'All_Vobs': All_Vobs,
        'All_Robs': All_Robs,
        'All_Enclosed_Masses': All_Enclosed_Masses,
        'All_e_Vobs': All_e_Vobs,
        'All_e_Robs': All_e_Robs,
        'All_e_Enclosed_Masses': All_e_Enclosed_Masses
    }

    savemat('galaxy_mass_properties.mat', save_dict)
    print("✅ Successfully saved galaxy_mass_properties.mat")


    # ===== SECTION 2: MEZZI ANALYSIS =====
import numpy as np
from scipy.io import loadmat
from scipy.interpolate import PchipInterpolator
import matplotlib.pyplot as plt
import os

# --- 2.2 Constants ---
G = 6.67430e-11
c = 2.99792458e8
arcsec = 1 / 206265
Mpc = 3.085677581e22
pc = 3.085677581e16
Msun = 1.9885e30
 

# --- 2.3 Load Observational Data ---
print("[LOAD] Loading observational data...")
mat_data = loadmat('galaxy_mass_properties.mat', simplify_cells=True)

galaxy_names = [x.decode('utf-8') if isinstance(x, bytes) else str(x) for x in mat_data['galaxy_names'].flatten()]
All_Robs = [np.atleast_1d(x) for x in mat_data['All_Robs']]
All_Enclosed_Masses = [np.atleast_1d(x) for x in mat_data['All_Enclosed_Masses']]
All_Vobs = [np.atleast_1d(x) for x in mat_data['All_Vobs']]
Distance_to_G = np.atleast_1d(mat_data['Distance_to_G']).flatten()

numGalaxies = len(All_Robs)
print(f"[LOAD] Analyzing {numGalaxies} galaxies with Mezzi corrections...")

# --- HELPER: Mass Profile Scaling ---
def mass_profile_from_mass(r_obs, M_enc, scale_factor):
    """
    Reconstruct true mass profile via geometric projection.
    
    [M_TRUE_CUM, M_OBS_CUM, MASS_RATIO] = MASSPROFILEFROMMASS(R_OBS, M_ENC, SCALEFACTOR)
    reconstructs the true mass profile from observed coordinates by applying
    a geometric projection based on the Mezzi scale factor.

    PHYSICAL MODEL
    ==============
    For an axisymmetric system, the annular area element dA transforms 
    according to the Jacobian of the coordinate transformation. Observed 
    coordinates r_obs are related to true coordinates r_true by the Mezzi 
    scale factor ζ(r) (represented by SCALEFACTOR):

        r_true = r_obs / ζ(r_obs)

    Consider a discrete annular shell bounded by radii r_j and r_(j+1). 
    The observed and true areas of this shell are, respectively:

        Delta_A_obs  = pi * (r_(j+1)^2 - r_j^2)                           (15)
        Delta_A_true = pi * (r_true_(j+1)^2 - r_true_j^2)                 (16)

    The area ratio R_j = Delta_A_true / Delta_A_obs quantifies the 
    geometric distortion of the j-th shell. For a Mezzi scale factor 
    with ζ < 1, the shells expand in true coordinates (R_j > 1).

    DENSITY INVARIANCE APPROXIMATION
    =================================
    The reconstruction utilizes a local approximation of density invariance. 
    We posit that the observed surface density Sigma_obs provides the 
    correct mass density for the local shell, but the coordinate grid is 
    compressed, mapping the true extent onto smaller observed radii.

    By projecting this invariant density onto the true, expanded metric, 
    the true surface density is approximated by the observed value:

        Sigma_true approx Sigma_obs                                        (17)

    Consequently, the mass element for the expanded true shell is 
    approximated by applying the observed surface density to the larger 
    true geometry:

        Delta_M_true,j = Sigma_obs,j * Delta_A_true,j                     (19)

    Substituting the area ratio R_j, this can be expressed as:

        Delta_M_true,j = Delta_M_obs,j * R_j                              (19b)

    Since R_j > 1 for the Mezzi effect, the true shell mass exceeds the 
    observed shell mass.

    ALGORITHM
    =========
    Practical implementation employs a discrete shell algorithm:
      1. Discretize the observed profile into N-1 annular shells.
      2. Compute observed surface density: Sigma_obs,j = Delta_M_obs,j / Delta_A_obs,j
      3. Transform shell boundaries: r_true,j = r_obs,j / ζ(r_obs,j)
      4. Compute true shell areas and area ratios R_j.
      5. Redistribute mass: Delta_M_true,j = Sigma_obs,j * Delta_A_true,j
      6. Reconstruct cumulative profile via Eq. (20).

    INPUTS
    ======
    r_obs        : array-like, shape (N,)
        Observed radii in strictly increasing order.
    M_enc        : array-like, shape (N,)
        Enclosed mass M(<r_obs) at observed radii.
    scale_factor : array-like, shape (N,)
        Mezzi scale factor ζ(r_obs). Values < 1 indicate shells expand
        in true coordinates (R_j > 1).

    OUTPUTS
    =======
    M_true_cum   : ndarray, shape (N,)
        Cumulative mass M_true(<r_true) in true coordinates.
    M_obs_cum    : ndarray, shape (N,)
        Original observed cumulative mass.
    mass_ratio   : float
        Mass ratio η quantifying the net geometric bias:
        η = M_total,true / M_total,obs                          

    VALIDATION AND BOUNDARY CONDITIONS
    ===================================
    The innermost enclosed mass M_1 = M_obs(<r_obs,1) is treated as a central
    boundary condition. The assumption that the central region is unaffected 
    by geometric distortion ensures M_true(<r_true,1) = M_obs(<r_obs,1).

    For the Mezzi effect with ζ < 1, the expansion of outer shells yields 
    R_j > 1 throughout the profile, resulting in η > 1. The true total mass 
    exceeds the observed estimate because the observed aperture encompasses 
    a smaller physical region than the true galactic extent.

    NOTES AND CAVEATS
    =================
    * Surface density is assumed uniform within each shell (valid for fine binning).
    * Negative shell masses (numerical noise) are clipped to zero.
    * Non-monotonic M_ENC generates warning.
    """
    # Ensure 1D numpy arrays
    r_obs = np.asarray(r_obs).ravel()
    M_enc = np.asarray(M_enc).ravel()
    scale_factor = np.asarray(scale_factor).ravel()
    
    N = len(r_obs)
    if N < 2:
        raise ValueError("Input arrays must have at least 2 elements to form shells.")
        
    # Check for decreasing M_enc (Non-Monotonicity)
    dM_check = np.diff(M_enc)
    if np.any(dM_check < 0):
        warnings.warn(
            f"MassProfileFromMass:DecreasingMass - M_enc has {np.sum(dM_check < 0)} decreasing segment(s). "
            f"Max decrease: {np.max(np.abs(dM_check[dM_check < 0])):.4e}",
            RuntimeWarning
        )
        
    # 1. Convert enclosed mass to surface density in observed shells
    dM_obs_shell = np.diff(M_enc)
    dA_obs_shell = np.pi * (r_obs[1:]**2 - r_obs[:-1]**2)
    
    # Safe division to avoid NaNs from zero-area shells
    safe_dA_obs = np.where(dA_obs_shell > 0, dA_obs_shell, 1e-50)
    Sigma_shell = dM_obs_shell / safe_dA_obs
    r_obs_mid = (r_obs[:-1] + r_obs[1:]) / 2.0  
    
    # 2. Coordinate transformation
    r_true = r_obs / scale_factor
    
    # 3. Calculate true shell masses
    dA_true_shell = np.pi * (r_true[1:]**2 - r_true[:-1]**2)
    dM_true_shell = Sigma_shell * dA_true_shell
    dM_true_shell = np.maximum(dM_true_shell, 0.0)  # Clip negatives to zero
    r_true_mid = (r_true[:-1] + r_true[1:]) / 2.0   # Local variable
    
    # 4. EXACT CUMULATIVE RECONSTRUCTION
    M_true_cum = np.zeros(N)
    M_true_cum[0] = M_enc[0]
    M_true_cum[1:] = M_enc[0] + np.cumsum(dM_true_shell)
    
    # 5. Observed mass
    M_obs_cum = M_enc.copy()
    
    # 6. Calculate mass ratio
    M_total_true = M_true_cum[-1]
    M_total_obs = M_enc[-1]
    mass_ratio = M_total_true / M_total_obs if M_total_obs != 0 else np.nan
    
    return M_true_cum, M_obs_cum, mass_ratio


# --- 2.4 Main Galaxy Processing ---
Results = []
IndFigures = False  # Set to True to generate individual galaxy figures

for g in range(numGalaxies):
    print(f"[Geometric Scaling] Galaxy {g+1}/{numGalaxies}: {galaxy_names[g]}")
    res = {'GalaxyName': galaxy_names[g], 'galaxy_index': g+1}
    
    # Basic Parameters & QC
    Distance_to_Galaxy = Distance_to_G[g] * Mpc
    r_obs_I = All_Robs[g].copy()
    M_enc_obs = All_Enclosed_Masses[g].copy()
    
    valid_idx = (r_obs_I > 0) & np.isfinite(r_obs_I) & np.isfinite(M_enc_obs)
    r_obs_I = r_obs_I[valid_idx]
    M_enc_obs = M_enc_obs[valid_idx]
    
    if len(r_obs_I) == 0:
        continue
        
    r_max = np.max(r_obs_I)
    Mtot_obs = M_enc_obs[-1]
    
    # Observed Mass Interpolator
    M_interp_obs = PchipInterpolator(r_obs_I, M_enc_obs, extrapolate=True)
    def M_func_obs(x):
        return np.where(x <= r_max, M_interp_obs(x), Mtot_obs)
    
    # PG Flow & Mezzi Scale Factor
    v_flow = lambda x: -np.sqrt(2 * G * M_func_obs(x) / np.maximum(x, 1e-10 * pc))
    integrand_func = lambda x: v_flow(x) / (c * np.maximum(x, 1e-10 * pc))
    
    K_kinematic_compression = np.ones_like(r_obs_I)
    for i in range(len(r_obs_I)):
        r_i = r_obs_I[i]
        if abs(r_i - Distance_to_Galaxy) < 1e-8 * max(r_i, Distance_to_Galaxy):
            A_accumulated_potential = 0.0
        else:
            r_min_path = min(r_i, Distance_to_Galaxy)
            r_max_path = max(r_i, Distance_to_Galaxy)
            r_path = np.logspace(np.log10(max(r_min_path, 1e-3 * pc)), np.log10(r_max_path), 500)
            A_accumulated_potential = np.trapezoid(integrand_func(r_path), r_path)
        K_kinematic_compression[i] = np.exp(A_accumulated_potential)
        
    Mezzi_scale_factor_theory = np.exp(Compliance * (K_kinematic_compression - 1))
    
    # Mass Profile Scaling using geometric projection
    M_true_cum, M_obs_cum, mass_ratio = mass_profile_from_mass(r_obs_I, M_enc_obs, Mezzi_scale_factor_theory)
    M_enc_true = M_true_cum
    
    r_obs = r_obs_I
    r_true = r_obs / Mezzi_scale_factor_theory
    
    # True Mass Interpolator
    M_interp_true = PchipInterpolator(r_true, M_enc_true, extrapolate=True)
    def M_func_true(x):
        return np.where(x <= r_max, M_interp_true(x), M_enc_true[-1])
        
    # Velocity Recovery Test
    v_obs = All_Vobs[g][valid_idx]
    v_true = v_obs.copy()
    v_Newton_obs = np.sqrt(G * M_enc_obs / r_obs)
    v_Newton_true = np.sqrt(G * M_enc_true / r_true)
    v_Newton_true_To_observer_Frame = v_Newton_true.copy()
    
    residual_naive = v_true - v_Newton_obs
    residual_corrected = v_true - v_Newton_true_To_observer_Frame
    rms_naive = np.sqrt(np.mean(residual_naive**2))
    rms_corrected = np.sqrt(np.mean(residual_corrected**2))
    
    # Density Profiles
    r_fine_obs = np.logspace(np.log10(np.min(r_obs)), np.log10(r_max * 1.1), 2000)
    M_fine_obs = M_interp_obs(r_fine_obs)
    dMdr_obs = np.gradient(M_fine_obs, r_fine_obs)
    rho_fine_obs = dMdr_obs / (4 * np.pi * r_fine_obs**2 + 1e-50)
    
    r_fine_true = np.logspace(np.log10(np.min(r_true)), np.log10(np.max(r_true) * 1.1), 2000)
    M_fine_true = M_interp_true(r_fine_true)
    dMdr_true = np.gradient(M_fine_true, r_fine_true)
    rho_fine_true = dMdr_true / (4 * np.pi * r_fine_true**2 + 1e-50)
    
    # Surface Density (Abel-type integral)
    R_grid_obs = r_obs
    Sigma_obs = np.zeros_like(R_grid_obs)
    r_upper_obs = np.max(r_fine_obs)
    for j in range(len(R_grid_obs)):
        R = R_grid_obs[j]
        if R <= 0 or R > r_max or np.isnan(R): continue
        r_int = np.linspace(R * 1.0001, r_upper_obs, 1000)
        rho_int = np.interp(r_int, r_fine_obs, rho_fine_obs, left=0, right=0)
        integrand_vals = rho_int * r_int / np.sqrt(np.maximum(r_int**2 - R**2, 1e-20))
        Sigma_obs[j] = 2 * np.trapezoid(integrand_vals, r_int)
        
    R_grid_true = r_true
    Sigma_true = np.zeros_like(R_grid_true)
    r_upper_true = np.max(r_fine_true)
    for j in range(len(R_grid_true)):
        R = R_grid_true[j]
        if R <= 0 or R > np.max(r_true) or np.isnan(R): continue
        r_int = np.linspace(R * 1.0001, r_upper_true, 1000)
        rho_int = np.interp(r_int, r_fine_true, rho_fine_true, left=0, right=0)
        integrand_vals = rho_int * r_int / np.sqrt(np.maximum(r_int**2 - R**2, 1e-20))
        Sigma_true[j] = 2 * np.trapezoid(integrand_vals, r_int)
        
    Sigma_true_on_obs_grid = np.zeros_like(R_grid_obs)
    for j in range(len(R_grid_obs)):
        R_obs_val = R_grid_obs[j]
        idx = np.argmin(np.abs(r_obs - R_obs_val))
        if 0 <= idx < len(Mezzi_scale_factor_theory):
            Mezzi_at_R = Mezzi_scale_factor_theory[idx]
            R_true_val = R_obs_val / Mezzi_at_R
            Sigma_val = np.interp(R_true_val, R_grid_true, Sigma_true, left=0, right=0)
            Sigma_true_on_obs_grid[j] = Sigma_val / (Mezzi_at_R**2 + 1e-20)
            
    # Deflection Angles
    alpha_rad_obs = np.zeros_like(R_grid_obs)
    for j in range(len(R_grid_obs)):
        R = R_grid_obs[j]
        if R < 1e-8 * pc or np.isnan(R): continue
        alpha_rad_obs[j] = (4 * G / c**2) * (M_func_obs(R) / R)
        
    alpha_rad_true = np.zeros_like(R_grid_true)
    for j in range(len(R_grid_true)):
        R = R_grid_true[j]
        if R < 1e-8 * pc or np.isnan(R) or R > np.max(r_true): continue
        alpha_rad_true[j] = (4 * G / c**2) * (M_func_true(R) / R)
        
    # Store Results
    res.update({
        'Distance_to_Galaxy': Distance_to_Galaxy, 'r_obs': r_obs, 'M_enc_obs': M_enc_obs,
        'v_obs': v_obs, 'r_max': r_max, 'Mtot_obs': Mtot_obs, 'PG': K_kinematic_compression,
        'Mezzi_th': Mezzi_scale_factor_theory, 'C': Compliance, 'r_true': r_true,
        'M_enc_true': M_enc_true, 'Mtot_true': M_enc_true[-1], 'v_true': v_true,
        'v_Newton_obs': v_Newton_obs, 'v_Newton_true': v_Newton_true,
        'v_Newton_true_To_obs': v_Newton_true_To_observer_Frame,
        'residual_naive': residual_naive, 'residual_corrected': residual_corrected,
        'rms_naive': rms_naive, 'rms_corrected': rms_corrected,
        'r_fine_obs': r_fine_obs, 'M_fine_obs': M_fine_obs, 'dMdr_obs': dMdr_obs,
        'rho_fine_obs': rho_fine_obs, 'r_fine_true': r_fine_true,
        'M_fine_true': M_fine_true, 'dMdr_true': dMdr_true, 'rho_fine_true': rho_fine_true,
        'R_grid_obs': R_grid_obs, 'Sigma_obs': Sigma_obs, 'r_upper_obs': r_upper_obs,
        'R_grid_true': R_grid_true, 'Sigma_true': Sigma_true, 'r_upper_true': r_upper_true,
        'Sigma_true_on_obs_grid': Sigma_true_on_obs_grid,
        'Sigma_obs_Msun_pc2': Sigma_obs * pc**2 / Msun,
        'Sigma_true_Msun_pc2': Sigma_true * pc**2 / Msun,
        'alpha_rad_obs': alpha_rad_obs, 'alpha_rad_true': alpha_rad_true,
        'alpha_obs_arcsec': alpha_rad_obs / arcsec,
        'alpha_true_arcsec': alpha_rad_true / arcsec,
        'max_alpha_obs_arcsec': np.max(alpha_rad_obs / arcsec),
        'max_alpha_true_arcsec': np.max(alpha_rad_true / arcsec),
        'Sigma_ratio': Sigma_true / (Sigma_obs + 1e-20),
        'alpha_ratio_true': alpha_rad_true / (alpha_rad_obs + 1e-20),
        'v_flow_obs': v_flow(r_obs),
        'integrand_obs': integrand_func(r_obs)
    })
    Results.append(res)

# ===== SECTION 2.5: VISUALIZATION =====
print("[VIZ] Setting up plot directories...")
baseDir = 'Results'
folders = ['RotationCurve', 'MassProfile', 'SurfaceDensity', 'Deflection']
for folder in folders:
    os.makedirs(os.path.join(baseDir, folder), exist_ok=True)

if IndFigures:
    print("[VIZ] Generating figures...")
    for g, res in enumerate(Results):
        g_idx = g + 1
        gal_name = res['GalaxyName']
        
        # --- 2.5.1 Rotation Curve ---
        r_pc_obs = res['r_obs'] / pc
        v_obs_kms = res['v_obs'] / 1e3
        v_newt_obs_kms = res['v_Newton_obs'] / 1e3
        v_true_kms = res['v_Newton_true_To_obs'] / 1e3
        
        fig, ax = plt.subplots(figsize=(8, 6))
        ax.scatter(r_pc_obs, v_obs_kms, c=[0.2, 0.4, 0.8], alpha=0.7, label='Observed', s=60, zorder=3)
        ax.plot(r_pc_obs, v_newt_obs_kms, 'b--', linewidth=2, label='Newtonian (Obs Frame)')
        ax.plot(r_pc_obs, v_true_kms, 'r-', linewidth=2.5, label='Newton (True Frame)')
        ax.set_xscale('log')
        ax.grid(True, which='both', ls='--', alpha=0.5)
        ax.set_xlabel('Radius (pc)', fontsize=16, fontweight='bold')
        ax.set_ylabel('Circular Velocity (km/s)', fontsize=16, fontweight='bold')
        ax.set_title(f'Rotation Curve — {gal_name}', fontsize=14, fontweight='bold')
        ax.legend(loc='best', fontsize=10, frameon=True)
        
        improvement = 100 * (1 - res['rms_corrected'] / (res['rms_naive'] + 1e-20))
        annot = (f"RMS (naive) = {res['rms_naive']/1e3:.1f} km/s\n"
                 f"RMS (corr) = {res['rms_corrected']/1e3:.1f} km/s\n"
                 f"Improvement = {improvement:.0f}%")
        ax.text(0.02, 0.70, annot, transform=ax.transAxes, fontsize=10,
                va='top', bbox=dict(boxstyle='square', facecolor='white', alpha=0.9, ec='gray'))
        plt.tight_layout()
        fig.savefig(os.path.join(baseDir, 'RotationCurve', f'Fig_2_5_1_RotationCurve_{g_idx}.png'), dpi=300)
        plt.close(fig)
        
        # --- 2.5.2 Mass Profile ---
        r_pc_true = res['r_true'] / pc
        M_obs_Msun = res['M_enc_obs'] / Msun
        M_true_Msun = res['M_enc_true'] / Msun
        
        fig, ax = plt.subplots(figsize=(8, 6))
        ax.loglog(r_pc_obs, M_obs_Msun, 'bo-', linewidth=2, markersize=6, label='Observed')
        ax.loglog(r_pc_true, M_true_Msun, 'rs-', linewidth=2, markersize=6, label='Mezzi-Corrected')
        ax.grid(True, which='both', ls='--', alpha=0.5)
        ax.set_xlabel('Radius (pc)', fontsize=16, fontweight='bold')
        ax.set_ylabel(r'Enclosed Mass ($M_\odot$)', fontsize=16, fontweight='bold')
        ax.set_title(f'Mass Profile — {gal_name}', fontsize=14, fontweight='bold')
        ax.legend(loc='best', fontsize=10, frameon=True)
        
        ratio_mass = res['Mtot_true'] / res['Mtot_obs']
        annot = (f"M_tot(obs) = {res['Mtot_obs']/Msun:.2e} Msun\n"
                 f"M_tot(true) = {res['Mtot_true']/Msun:.2e} Msun\n"
                 f"Ratio = {ratio_mass:.2f}")
        ax.text(0.02, 0.70, annot, transform=ax.transAxes, fontsize=10,
                va='top', bbox=dict(boxstyle='square', facecolor='white', alpha=0.9, ec='gray'))
        plt.tight_layout()
        fig.savefig(os.path.join(baseDir, 'MassProfile', f'Fig_2_5_2_MassProfile_{g_idx}.png'), dpi=300)
        plt.close(fig)
        
        # --- 2.5.3 Surface Density ---
        R_pc_obs = res['R_grid_obs'] / pc
        R_pc_true = res['R_grid_true'] / pc
        Sigma_obs_Msun = res['Sigma_obs_Msun_pc2']
        Sigma_true_Msun = res['Sigma_true_Msun_pc2']
        
        fig, ax = plt.subplots(figsize=(8, 6))
        ax.loglog(R_pc_obs, Sigma_obs_Msun, 'bo-', linewidth=2, markersize=6, label='Observed')
        ax.loglog(R_pc_true, Sigma_true_Msun, 'rs-', linewidth=2, markersize=6, label='Mezzi-Corrected')
        ax.grid(True, which='both', ls='--', alpha=0.5)
        ax.set_xlabel('Projected Radius (pc)', fontsize=16, fontweight='bold')
        ax.set_ylabel(r'Surface Density ($M_\odot/pc^2$)', fontsize=16, fontweight='bold')
        ax.set_title(f'Surface Density — {gal_name}', fontsize=14, fontweight='bold')
        ax.legend(loc='best', fontsize=10, frameon=True)
        
        core_ratio = Sigma_true_Msun[0] / (Sigma_obs_Msun[0] + 1e-20)
        max_dev = np.max(np.abs((Sigma_true_Msun - Sigma_obs_Msun) / (Sigma_obs_Msun + 1e-20))) * 100
        annot = f"Σ ratio (core) = {core_ratio:.2f}\nMax deviation = {max_dev:.0f}%"
        ax.text(0.02, 0.70, annot, transform=ax.transAxes, fontsize=10,
                va='top', bbox=dict(boxstyle='square', facecolor='white', alpha=0.9, ec='gray'))
        plt.tight_layout()
        fig.savefig(os.path.join(baseDir, 'SurfaceDensity', f'Fig_2_5_3_SurfaceDensity_{g_idx}.png'), dpi=300)
        plt.close(fig)
        
        # --- 2.5.4 Deflection Angles ---
        alpha_obs_arc = res['alpha_obs_arcsec']
        alpha_true_arc = res['alpha_true_arcsec']
        
        fig, ax = plt.subplots(figsize=(8, 6))
        ax.semilogx(r_pc_obs, alpha_obs_arc, 'bo-', linewidth=2, markersize=6, label='Observed')
        ax.semilogx(r_pc_true, alpha_true_arc, 'rs-', linewidth=2, markersize=6, label='Mezzi-Corrected')
        ax.set_yscale('log')
        ax.grid(True, which='both', ls='--', alpha=0.5)
        ax.set_xlabel('Impact Parameter (pc)', fontsize=16, fontweight='bold')
        ax.set_ylabel('Deflection Angle (arcsec)', fontsize=16, fontweight='bold')
        ax.set_title(f'Gravitational Deflection — {gal_name}', fontsize=14, fontweight='bold')
        ax.legend(loc='best', fontsize=10, frameon=True)
        plt.tight_layout()
        fig.savefig(os.path.join(baseDir, 'Deflection', f'Fig_2_5_4_Deflection_{g_idx}.png'), dpi=300)
        plt.close(fig)

print("[DONE] Mezzi Analysis & Visualization Complete.")

print("[2.6] Mezzi Scale Factor Analysis...")
analysisDir = os.path.join(baseDir, 'MezziAnalysis')
os.makedirs(analysisDir, exist_ok=True)

galaxySummary = []
for g in range(numGalaxies):
    res = Results[g]
    r_obs = res['r_obs']
    M_enc_obs = res['M_enc_obs']
    M_enc_true = res['M_enc_true']
    Mezzi_th = res['Mezzi_th']
    PG = res['PG']
    v_obs = res['v_obs']
    v_Newton_obs = res['v_Newton_obs']
    v_Newton_true = res['v_Newton_true_To_obs']
    
    summary = {
        'name': galaxy_names[g],
        'PG_mean': np.mean(PG),
        'PG_range': np.max(PG) - np.min(PG),
        'Mezzi_mean': np.mean(Mezzi_th),
        'Mezzi_range': np.max(Mezzi_th) - np.min(Mezzi_th),
        'M_total_obs_Msun': M_enc_obs[-1] / Msun,
        'M_total_true_Msun': M_enc_true[-1] / Msun,
        'mass_ratio': M_enc_true[-1] / (M_enc_obs[-1] + 1e-20),
        'v_central_obs_kms': v_obs[0] / 1e3,
        'v_flat_obs_kms': v_obs[np.argmin(np.abs(r_obs - r_obs[-1]/2))] / 1e3,
        'v_flat_newt_kms': v_Newton_obs[np.argmin(np.abs(r_obs - r_obs[-1]/2))] / 1e3,
        'v_flat_true_kms': v_Newton_true[np.argmin(np.abs(r_obs - r_obs[-1]/2))] / 1e3,
        'rms_naive_kms': res['rms_naive'] / 1e3,
        'rms_corrected_kms': res['rms_corrected'] / 1e3,
        'improvement_percent': 100 * (1 - res['rms_corrected'] / (res['rms_naive'] + 1e-20))
    }
    galaxySummary.append(summary)

# Extract arrays
mass_ratios = np.array([s['mass_ratio'] for s in galaxySummary])
PG_means = np.array([s['PG_mean'] for s in galaxySummary])
Mezzi_means = np.array([s['Mezzi_mean'] for s in galaxySummary])
galaxy_masses = np.array([s['M_total_obs_Msun'] for s in galaxySummary])
log_masses = np.log10(galaxy_masses)

# ===== FIG 2.6.2: Mezzi Scale Factor vs Normalized Radius (Inset as Main) =====
print("[2.6.2] Plotting Mezzi Scale Factor vs Normalized Radius...")
cmap = plt.colormaps['viridis']   

# Generate mass-based colors (N galaxies × 4 RGBA values)
mass_colors = cmap(np.linspace(0, 1, numGalaxies))

# Sort galaxies by mass for color assignment  
sort_idx = np.argsort(log_masses)  # Indices that would sort log_masses ascending
color_idx = np.argsort(sort_idx)   # color_idx[g] = color index for original galaxy g

# Create figure with same dimensions  
fig, ax = plt.subplots(figsize=(8, 6))

# Plot Mezzi scale factor curves vs NORMALIZED radius 
for g in range(numGalaxies):
    # Normalized radius: r / r_max for each galaxy
    r_norm = Results[g]['r_obs'] / np.max(Results[g]['r_obs'])
    Mezzi = Results[g]['Mezzi_th']
    
    # Ensure 1D arrays for plotting
    r_norm = np.asarray(r_norm).ravel()
    Mezzi = np.asarray(Mezzi).ravel()
    
    this_color = mass_colors[color_idx[g]]  # RGBA array [R, G, B, A]
    ax.plot(r_norm, Mezzi, linewidth=1.2, color=this_color)

# Configure axes  
ax.set_xlim(0, 1)  # Normalized radius always [0, 1]
ax.set_xlabel('Normalized Radius ($r/r_{max}$)', fontsize=22, fontweight='bold')
ax.set_ylabel('Mezzi Scale Factor', fontsize=22, fontweight='bold')
ax.grid(True, which='both', ls='--', alpha=0.5)
ax.tick_params(labelsize=11)

for spine in ax.spines.values():
    spine.set_linewidth(1.5)

# Colorbar with proper normalization - ticks in DATA space
norm = plt.Normalize(vmin=np.min(log_masses), vmax=np.max(log_masses))
sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])  # Required dummy array for colorbar

cbar = plt.colorbar(sm, ax=ax, location='right', pad=0.02)
cbar.set_label('$\\log_{10}(M_{obs}/M_\\odot)$', fontsize=14, fontweight='bold')

# Set ticks at actual data values (ScalarMappable auto-normalizes these)
tick_values = np.linspace(np.min(log_masses), np.max(log_masses), 5)
cbar.set_ticks(tick_values)  # Pass data values, not normalized positions
cbar.set_ticklabels([f"{x:.1f}" for x in tick_values])

# Tight layout to prevent label cutoff
plt.tight_layout()

# Save with updated filename to reflect content change
fig.savefig(os.path.join(analysisDir, 'Fig_2_6_2_Mezzi_vs_NormalizedRadius.png'), 
            dpi=300, bbox_inches='tight')
plt.close(fig)

print(f"[2.6.2] Saved: Fig_2_6_2_Mezzi_vs_NormalizedRadius.png")


# ===== FIG 2.6.3: Total Mass Comparison =====
print("[2.6.3] Plotting Mass Comparison...")
cmap_turbo = get_cmap('turbo')
M_obs_all = np.array([s['M_total_obs_Msun'] for s in galaxySummary])
M_true_all = np.array([s['M_total_true_Msun'] for s in galaxySummary])
correction_factor = M_true_all / (M_obs_all + 1e-20)
sort_idx = np.argsort(correction_factor)
log_M_obs_sorted = np.log10(M_obs_all[sort_idx])
log_M_true_sorted = np.log10(M_true_all[sort_idx])
sorted_corr = correction_factor[sort_idx]

corr_min, corr_max = np.min(sorted_corr), np.max(sorted_corr)
corr_norm = (sorted_corr - corr_min) / (corr_max - corr_min)
color_indices = np.clip(np.round(corr_norm * 255).astype(int), 0, 255)
line_colors = cmap_turbo(color_indices)

fig = plt.figure(figsize=(7.2, 5.5))
ax_main = fig.add_axes([0.10, 0.12, 0.75, 0.82])
x_pos = np.arange(1, numGalaxies+1)

# Connecting lines
for g in range(numGalaxies):
    alpha = 0.15 if sorted_corr[g] > np.percentile(sorted_corr, 75) else 1.0
    ax_main.plot([x_pos[g], x_pos[g]], [log_M_obs_sorted[g], log_M_true_sorted[g]], 
                 color=[*line_colors[g][:3], alpha], linewidth=2.5)

# Data points
ax_main.scatter(x_pos, log_M_obs_sorted, s=70, c=[0.95]*3, marker='o', 
                edgecolors=[0.15]*3, linewidths=1.2, label='Observed Mass')
ax_main.scatter(x_pos, log_M_true_sorted, s=70, c=line_colors, marker='^', 
                edgecolors=[0.15]*3, linewidths=1.2, label='True (Mezzi-Corrected)')

ax_main.set_xlabel('Galaxy (sorted by $M_{true}/M_{obs}$)', fontsize=22, fontweight='bold')
ax_main.set_ylabel('$\\log_{10}(M_{obs}/M_\\odot)$', fontsize=22, fontweight='bold')
ax_main.grid(axis='y', ls='--', alpha=0.5)
ax_main.legend(loc='upper left', fontsize=10, frameon=True, edgecolor=[0.3]*3)

x_ticks = np.arange(1, numGalaxies+1, 5) if numGalaxies > 20 else x_pos
ax_main.set_xticks(x_ticks)
ax_main.set_xticklabels(x_ticks, rotation=45)
ax_main.tick_params(labelsize=10)

cbar = plt.colorbar(plt.cm.ScalarMappable(cmap=cmap_turbo, norm=Normalize(corr_min, corr_max)), 
                    ax=ax_main, fraction=0.03, pad=0.04)
cbar.set_label('$M_{true}/M_{obs}$', fontsize=20, fontweight='bold')
if corr_max / corr_min < 10:
    cbar_ticks = np.linspace(corr_min, corr_max, 5)
else:
    cbar_ticks = np.array([1, 2, 5, 10, 20, 50])
    cbar_ticks = cbar_ticks[(cbar_ticks >= corr_min) & (cbar_ticks <= corr_max)]
cbar.set_ticks(cbar_ticks)
cbar.set_ticklabels([f"{t:.1f}×" for t in cbar_ticks])

# Inset Histogram
ax_inset = fig.add_axes([0.52, 0.22, 0.28, 0.20])
counts, edges = np.histogram(sorted_corr, bins=20)
ax_inset.hist(edges[:-1], bins=edges, weights=counts, facecolor=[0.5, 0.55, 0.6], 
              edgecolor=[0.2]*3, linewidth=0.6, alpha=0.85)
ax_inset.set_xlabel('$M_{true}/M_{obs}$', fontsize=11, fontweight='bold')
ax_inset.set_ylabel('Count', fontsize=11, fontweight='bold')
ax_inset.grid(ls='-', alpha=0.8)
ax_inset.tick_params(labelsize=9)
if corr_max / corr_min > 5:
    ax_inset.set_xscale('log')
    ax_inset.xaxis.set_major_formatter(plt.FuncFormatter(lambda x, _: f"{x:.0f}"))

fig.savefig(os.path.join(analysisDir, 'Fig_2_6_3_Mass_Comparison.png'), dpi=600, bbox_inches='tight')
plt.close(fig)

# ===== FIG 2.6.4: Velocity Residuals =====
print("[2.6.4] Plotting Velocity Residuals...")
v_flat_obs = np.array([s['v_flat_obs_kms'] for s in galaxySummary])
v_flat_mezzi = np.array([s['v_flat_true_kms'] for s in galaxySummary])
residual_mezzi = v_flat_mezzi - v_flat_obs
log_v_obs = np.log10(v_flat_obs)
log_residual_mezzi = np.log10(v_flat_mezzi / (v_flat_obs + 1e-20))

# Create subplots once (no nexttile needed)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6), gridspec_kw={'wspace': 0.25})
sort_idx = np.argsort(log_masses)
color_idx = np.argsort(sort_idx)

# === PANEL (a): Linear Residuals ===
for g in range(numGalaxies):
    ax1.scatter(v_flat_obs[g], residual_mezzi[g], s=80, c=[mass_colors[color_idx[g]]],
                marker='o', edgecolors='k', linewidths=1, zorder=3)
ax1.axhline(0, color='k', ls='--', lw=2, alpha=0.6)
ax1.axhline(np.mean(residual_mezzi), color='r', lw=2, alpha=0.9, label='Mean')
ax1.axhline(np.mean(residual_mezzi) + np.std(residual_mezzi), color='r', ls=':', lw=1.5, alpha=0.5, label='±1σ')
ax1.axhline(np.mean(residual_mezzi) - np.std(residual_mezzi), color='r', ls=':', lw=1.5, alpha=0.5)
ax1.set_xlabel('$v_{obs}$ (km/s)', fontsize=20, fontweight='bold')
ax1.set_ylabel('$v_{Newt,True} - v_{obs}$ (km/s)', fontsize=20, fontweight='bold')
ax1.grid(True, ls='--', alpha=0.5)
ax1.tick_params(labelsize=12)

stats_a = (f"Linear Statistics:\nN = {numGalaxies} galaxies\n"
           f"RMS residual = {np.sqrt(np.mean(residual_mezzi**2)):.1f} km/s\n"
           f"Std dev = {np.std(residual_mezzi):.1f} km/s\n"
           f"Mean bias = {np.mean(residual_mezzi):+.1f} km/s\n"
           f"Median bias = {np.median(residual_mezzi):+.1f} km/s")
ax1.text(0.02, 0.98, stats_a, transform=ax1.transAxes, fontsize=12, va='top', ha='left',
         bbox=dict(facecolor='white', alpha=0.9, edgecolor='gray'))

# === PANEL (b): Logarithmic Residuals (Dex) ===
for g in range(numGalaxies):
    ax2.scatter(log_v_obs[g], log_residual_mezzi[g], s=80, c=[mass_colors[color_idx[g]]],
                marker='o', edgecolors='k', linewidths=1, zorder=3)
ax2.axhline(0, color='k', ls='--', lw=2, alpha=0.6)
ax2.axhline(np.mean(log_residual_mezzi), color='b', lw=2, alpha=0.9, label='Mean')
ax2.axhline(np.std(log_residual_mezzi), color='b', ls=':', lw=1.5, alpha=0.5, label='±1σ')
ax2.axhline(-np.std(log_residual_mezzi), color='b', ls=':', lw=1.5, alpha=0.5)
ax2.set_xlabel('$\\log_{10}(v_{obs})$', fontsize=20, fontweight='bold')
ax2.set_ylabel('$\\log(v_{Newt,True}/v_{obs})$ (dex)', fontsize=20, fontweight='bold')
ax2.grid(True, ls='--', alpha=0.5)
ax2.tick_params(labelsize=12)

# Colorbar for Panel B
norm = plt.Normalize(vmin=np.min(log_masses), vmax=np.max(log_masses))
sm = plt.cm.ScalarMappable(cmap=plt.colormaps['viridis'], norm=norm)
sm.set_array([])
cbar = plt.colorbar(sm, ax=ax2, fraction=0.03, pad=0.04)
cbar.set_label('$\\log_{10}(M_{obs}/M_\\odot)$', fontsize=12, fontweight='bold')

within_10 = np.sum(np.abs(log_residual_mezzi) < np.log10(1.10)) / numGalaxies * 100
within_20 = np.sum(np.abs(log_residual_mezzi) < np.log10(1.20)) / numGalaxies * 100
stats_b = (f"Logarithmic Statistics:\nN = {numGalaxies} galaxies\n"
           f"RMS = {np.sqrt(np.mean(log_residual_mezzi**2)):.3f} dex\n"
           f"Scatter = {np.std(log_residual_mezzi):.3f} dex\n"
           f"Mean bias = {np.mean(log_residual_mezzi):+.3f} dex\n"
           f"Median bias = {np.median(log_residual_mezzi):+.3f} dex\n"
           f"Within 10%: {within_10:.1f}% of galaxies\n"
           f"Within 20%: {within_20:.1f}% of galaxies")
ax2.text(0.02, 0.98, stats_b, transform=ax2.transAxes, fontsize=12, va='top', ha='left',
         bbox=dict(facecolor='white', alpha=0.9, edgecolor='gray'))

plt.tight_layout()
fig.savefig(os.path.join(analysisDir, 'Fig_2_6_4_Mezzi_Residuals.png'), dpi=300, bbox_inches='tight')
plt.close(fig)

# ===== FIG 2.6.5: Dex Error Distribution =====
print("[2.6.5] Plotting Dex Error Distribution...")
from scipy.stats import norm as scipy_norm  # Explicit import to avoid shadowing

fig, (ax3, ax4) = plt.subplots(1, 2, figsize=(16, 7), gridspec_kw={'wspace': 0.15})

# === PANEL (a): Histogram ===
bin_edges = np.linspace(np.min(log_residual_mezzi)-0.03, np.max(log_residual_mezzi)+0.03, 12)
counts, _, patches = ax3.hist(log_residual_mezzi, bins=bin_edges, facecolor=[0.25, 0.55, 0.35], 
                              edgecolor='k', linewidth=2, alpha=0.75)
x_norm = np.linspace(bin_edges[0], bin_edges[-1], 200)

mean_val = np.mean(log_residual_mezzi)
std_val = np.std(log_residual_mezzi)
y_norm = scipy_norm.pdf(x_norm, loc=mean_val, scale=std_val)
y_norm_scaled = y_norm * numGalaxies * (bin_edges[1]-bin_edges[0])

ax3.plot(x_norm, y_norm_scaled, 'r-', linewidth=3, label='Gaussian fit')
ax3.axvline(0, color='k', ls='--', lw=2.5, alpha=0.6)
ax3.axvline(mean_val, color='b', lw=2.5, alpha=0.9, label=f"Mean: {mean_val:.3f}")
ax3.axvline(mean_val + std_val, color='b', ls=':', lw=2, alpha=0.5)
ax3.axvline(mean_val - std_val, color='b', ls=':', lw=2, alpha=0.5)
ax3.set_xlabel('$\\log(v_{Newton,true}) - \\log(v_{obs})$', fontsize=24, fontweight='bold')
ax3.set_ylabel('Number of Galaxies', fontsize=24, fontweight='bold')
ax3.grid(True, ls='--', alpha=0.5)
ax3.tick_params(labelsize=14)

norm_text = (f"Distribution:\nN = {numGalaxies}\n"
             f"Mean = {mean_val:.3f} dex\n"
             f"Median = {np.median(log_residual_mezzi):.3f} dex\n"
             f"Skewness = {skew(log_residual_mezzi):.2f}")
ax3.text(0.02, 0.98, norm_text, transform=ax3.transAxes, fontsize=14, va='top', ha='left',
         bbox=dict(facecolor='white', alpha=0.9, edgecolor='gray'))

# === PANEL (b): CDF ===
abs_log_res = np.abs(log_residual_mezzi)
sorted_abs = np.sort(abs_log_res)
cdf_y = np.arange(1, numGalaxies+1) / numGalaxies
ax4.plot(sorted_abs, cdf_y, 'b-', linewidth=4, label='Cumulative')
p50, p68, p95 = np.percentile(abs_log_res, [50, 68, 95])
ax4.axvline(p50, color='g', ls='--', lw=2.5, label=f"50%: {p50:.3f} dex")
ax4.axvline(p68, color='r', ls='--', lw=2.5, label=f"68%: {p68:.3f} dex")
ax4.axvline(p95, color='m', ls='--', lw=2.5, label=f"95%: {p95:.3f} dex")
ax4.set_xlabel('$|\\log(v_{Newton,true}) - \\log(v_{obs})|$', fontsize=24, fontweight='bold')
ax4.set_ylabel('Cumulative Fraction', fontsize=24, fontweight='bold')
ax4.set_xlim(0, np.max(sorted_abs)*1.05)
ax4.set_ylim(0, 1.05)
ax4.grid(True, ls='--', alpha=0.5)
ax4.legend(loc='lower right', fontsize=12)
ax4.tick_params(labelsize=14)

perc_text = (f"Percentiles:\nN = {numGalaxies}\n"
             f"50% (median) = {p50:.3f} dex\n"
             f"68% (1σ) = {p68:.3f} dex\n"
             f"95% (2σ) = {p95:.3f} dex\n"
             f"Mean |error| = {np.mean(abs_log_res):.3f} dex")
ax4.text(0.98, 0.02, perc_text, transform=ax4.transAxes, fontsize=14, va='bottom', ha='right',
         bbox=dict(facecolor='white', alpha=0.9, edgecolor='gray'))

plt.tight_layout()
fig.savefig(os.path.join(analysisDir, 'Fig_2_6_5_Mezzi_DexError.png'), dpi=300, bbox_inches='tight')
plt.close(fig)

# ===== SECTION 2.7: INNER REGION ANALYSIS =====
print("[2.7] Inner Region Analysis...")
InnerAnalysis = []
core_fraction = 0.3

for g in range(numGalaxies):
    res = Results[g]
    r_obs = res['r_obs']
    v_obs = res['v_obs']
    v_newt_obs = res['v_Newton_obs']
    v_newt_mezzi = res['v_Newton_true_To_obs']
    
    N = min(len(r_obs), len(v_obs), len(v_newt_obs), len(v_newt_mezzi))
    r_obs, v_obs = r_obs[:N], v_obs[:N]
    v_newt_obs, v_newt_mezzi = v_newt_obs[:N], v_newt_mezzi[:N]
    
    valid = (r_obs > 0) & (v_obs > 0) & np.isfinite(v_obs) & np.isfinite(v_newt_obs) & np.isfinite(v_newt_mezzi)
    if np.sum(valid) < 10:
        InnerAnalysis.append({'valid': False, 'galaxy': galaxy_names[g]})
        continue
        
    r, v, vn, vm = r_obs[valid], v_obs[valid], v_newt_obs[valid], v_newt_mezzi[valid]
    r_core = core_fraction * np.max(r)
    core_mask = r < r_core
    outer_mask = ~core_mask
    
    res_naive = v - vn
    res_mezzi = v - vm
    
    stats = {}
    if np.sum(core_mask) > 3:
        stats['N_core'] = np.sum(core_mask)
        stats['rms_naive_core'] = np.sqrt(np.mean(res_naive[core_mask]**2))
        stats['rms_mezzi_core'] = np.sqrt(np.mean(res_mezzi[core_mask]**2))
        stats['bias_naive_core'] = np.mean(res_naive[core_mask])
        stats['bias_mezzi_core'] = np.mean(res_mezzi[core_mask])
        stats['std_naive_core'] = np.std(res_naive[core_mask])
        stats['std_mezzi_core'] = np.std(res_mezzi[core_mask])
    if np.sum(outer_mask) > 3:
        stats['N_outer'] = np.sum(outer_mask)
        stats['rms_naive_outer'] = np.sqrt(np.mean(res_naive[outer_mask]**2))
        stats['rms_mezzi_outer'] = np.sqrt(np.mean(res_mezzi[outer_mask]**2))
        stats['bias_naive_outer'] = np.mean(res_naive[outer_mask])
        stats['bias_mezzi_outer'] = np.mean(res_mezzi[outer_mask])
        stats['std_naive_outer'] = np.std(res_naive[outer_mask])
        stats['std_mezzi_outer'] = np.std(res_mezzi[outer_mask])
        
    InnerAnalysis.append({'valid': True, 'galaxy': galaxy_names[g], 'stats': stats, 'r': r, 'v': v, 
                          'v_newt': vn, 'v_mezzi': vm, 'core_mask': core_mask, 'outer_mask': outer_mask})

# ===== FIG 2.7.1: Bias Distribution =====
print("[2.7.1] Plotting Bias Comparison...")
valid_galaxies = [i for i, ia in enumerate(InnerAnalysis) if ia.get('valid')]

bias_n_core, bias_m_core, bias_n_outer, bias_m_outer = [], [], [], []
rms_n_core, rms_m_core, rms_n_outer, rms_m_outer = [], [], [], []

for g in valid_galaxies:
    s = InnerAnalysis[g]['stats']
    if s.get('N_core', 0) > 3:
        bias_n_core.append(s['bias_naive_core']/1e3)
        bias_m_core.append(s['bias_mezzi_core']/1e3)
        rms_n_core.append(s['rms_naive_core']/1e3)
        rms_m_core.append(s['rms_mezzi_core']/1e3)
    if s.get('N_outer', 0) > 3:
        bias_n_outer.append(s['bias_naive_outer']/1e3)
        bias_m_outer.append(s['bias_mezzi_outer']/1e3)
        rms_n_outer.append(s['rms_naive_outer']/1e3)
        rms_m_outer.append(s['rms_mezzi_outer']/1e3)

if len(bias_n_core) > 0:
    fig, ax = plt.subplots(figsize=(7, 6))
    naive_color = [0.85, 0.325, 0.098]
    mezzi_color = [0.466, 0.674, 0.188]

    bp = ax.boxplot([bias_n_core, bias_m_core, bias_n_outer, bias_m_outer],
                    patch_artist=True, widths=0.22, positions=[0.85, 1.15, 1.85, 2.15])

    # Apply colors to boxplot elements
    box_colors = [naive_color, mezzi_color, naive_color, mezzi_color]
    for patch, color in zip(bp['boxes'], box_colors):
        patch.set_facecolor([*color[:3], 0.35])
        patch.set_edgecolor(color)
    for i, whisker in enumerate(bp['whiskers']):
        whisker.set_color(box_colors[i//2])
    for i, cap in enumerate(bp['caps']):
        cap.set_color(box_colors[i//2])
    for median in bp['medians']:
        median.set_color('k')
        median.set_linewidth(2)

    ax.axhline(0, color='k', ls='--', lw=2, alpha=0.6)
    ax.set_xticks([1, 2])
    ax.set_xticklabels([f'Inner ($r < {core_fraction:.2f}$)', f'Outer ($r > {core_fraction:.2f}$)'])
    ax.set_ylabel('Bias $\\langle v_{obs} - v_{Newton,True} \\rangle$ [km/s]', fontsize=14)
    ax.grid(True, ls='--', alpha=0.5)
    ax.tick_params(labelsize=12)

    # Custom legend
    from matplotlib.patches import Patch
    legend_elements = [
        Patch(facecolor=[*naive_color, 0.35], edgecolor=naive_color, label='Naive Newtonian'),
        Patch(facecolor=[*mezzi_color, 0.35], edgecolor=mezzi_color, label='True Frame Newtonian')
    ]
    ax.legend(handles=legend_elements, loc='best', fontsize=11)

    # Compute IQRs and Mean RMS for annotation (safely)
    iqr_n_c = np.percentile(bias_n_core, 75) - np.percentile(bias_n_core, 25)
    iqr_m_c = np.percentile(bias_m_core, 75) - np.percentile(bias_m_core, 25)
    iqr_n_o = np.percentile(bias_n_outer, 75) - np.percentile(bias_n_outer, 25)
    iqr_m_o = np.percentile(bias_m_outer, 75) - np.percentile(bias_m_outer, 25)

    stats_text = (
        "Median Bias [IQR] (Mean RMS):\n"
        f"Inner Naive: {np.median(bias_n_core):+.2f} [{iqr_n_c:.2f}] ({np.mean(rms_n_core):.2f}) km/s\n"
        f"Inner Mezzi: {np.median(bias_m_core):+.2f} [{iqr_m_c:.2f}] ({np.mean(rms_m_core):.2f}) km/s\n"
        f"Outer Naive: {np.median(bias_n_outer):+.2f} [{iqr_n_o:.2f}] ({np.mean(rms_n_outer):.2f}) km/s\n"
        f"Outer Mezzi: {np.median(bias_m_outer):+.2f} [{iqr_m_o:.2f}] ({np.mean(rms_m_outer):.2f}) km/s"
    )
    ax.text(0.02, 0.98, stats_text, transform=ax.transAxes, fontsize=10, va='top', ha='left',
            bbox=dict(facecolor='white', alpha=0.9, edgecolor='gray'))

    fig.savefig(os.path.join(analysisDir, 'Fig_2_7_1_Bias_Comparison.png'), dpi=300, bbox_inches='tight')
    plt.close(fig)

# ===== SECTION 2.8: SAVE SUMMARY =====
print("[2.8] Saving Summary Data...")
# Convert to savemat-friendly format
def prep_for_savemat(data):
    if not isinstance(data, list) or len(data) == 0:
        return {}
    if not isinstance(data[0], dict):
        return np.array(data)

    # 1. Collect all unique keys across every dictionary
    all_keys = sorted(set().union(*(d.keys() for d in data if isinstance(d, dict))))
    out = {}
    
    for k in all_keys:
        vals = [d.get(k) for d in data if isinstance(d, dict)]
        non_none = [v for v in vals if v is not None]
        
        if not non_none:
            # All None -> default to numeric NaN
            out[k] = np.full(len(vals), np.nan, dtype=float)
            continue
            
        # 2. Handle numpy arrays (common in your Results/InnerAnalysis)
        if any(isinstance(v, np.ndarray) for v in non_none):
            cleaned = [v if v is not None else np.array([]) for v in vals]
            out[k] = np.array(cleaned, dtype=object)
            
        # 3. Try numeric conversion
        else:
            try:
                cleaned = [np.nan if v is None else float(v) for v in vals]
                out[k] = np.array(cleaned, dtype=float)
                
            # 4. Fallback to strings/objects
            except (ValueError, TypeError):
                cleaned = ['' if v is None else str(v) for v in vals]
                out[k] = np.array(cleaned, dtype=object)
                
    return out


# ===== SECTION 3: COMPREHENSIVE SCALING ANALYSIS =====
print("[3.0] Starting Comprehensive Scaling Analysis...")
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from scipy.stats import ttest_1samp, pearsonr, linregress

ScalingDir = os.path.join(baseDir, 'ScalingAnalysis')
os.makedirs(ScalingDir, exist_ok=True)

# Predefine colormaps & mass colors to avoid scope errors
cmap = plt.colormaps['viridis']
mass_colors_global = cmap(np.linspace(0, 1, numGalaxies))
sort_idx_global = np.argsort(log_masses)
_, color_idx_global = np.sort(sort_idx_global, axis=0), np.argsort(sort_idx_global)

# --- 3.1 Rotation Curve Analysis ---
RC_stats = []
all_rms_naive = np.zeros(numGalaxies)
all_rms_true = np.zeros(numGalaxies)
all_R2_naive = np.zeros(numGalaxies)
all_R2_true = np.zeros(numGalaxies)

for g in range(numGalaxies):
    v_obs = Results[g]['v_obs']
    v_newt = Results[g]['v_Newton_obs']
    v_true = Results[g]['v_Newton_true_To_obs']
    
    SS_tot = np.sum((v_obs - np.mean(v_obs))**2)
    R2_naive = 1 - np.sum((v_obs - v_newt)**2) / (SS_tot + 1e-20)
    R2_true = 1 - np.sum((v_obs - v_true)**2) / (SS_tot + 1e-20)
    
    sigma_v = 0.05 * v_obs
    chi2_naive = np.sum(((v_obs - v_newt) / sigma_v)**2) / (len(v_obs) - 1)
    chi2_true = np.sum(((v_obs - v_true) / sigma_v)**2) / (len(v_obs) - 1)
    
    RC_stats.append({
        'R2_naive': R2_naive, 'R2_true': R2_true,
        'chi2_naive': chi2_naive, 'chi2_true': chi2_true,
        'v_max_obs': np.max(v_obs)/1e3, 'v_max_newt': np.max(v_newt)/1e3,
        'v_max_true': np.max(v_true)/1e3
    })
    all_rms_naive[g] = Results[g]['rms_naive'] / 1e3
    all_rms_true[g] = Results[g]['rms_corrected'] / 1e3
    all_R2_naive[g] = R2_naive
    all_R2_true[g] = R2_true

# --- 3.2 Enclosed Mass Correction Analysis ---
mass_scaling = []
total_mass_obs = np.zeros(numGalaxies)
total_mass_true = np.zeros(numGalaxies)
mass_concentration_obs = np.zeros(numGalaxies)
mass_concentration_true = np.zeros(numGalaxies)

for g in range(numGalaxies):
    r_obs = Results[g]['r_obs']
    r_true = Results[g]['r_true']
    M_obs = Results[g]['M_enc_obs']
    M_true = Results[g]['M_enc_true']
    
    r_half = r_obs[-1] / 2
    idx_half_obs = np.argmin(np.abs(r_obs - r_half))
    idx_half_true = np.argmin(np.abs(r_true - r_half))
    
    M_half_obs = M_obs[idx_half_obs]
    M_half_true = M_true[idx_half_true]
    
    C_obs = M_half_obs / (M_obs[-1] + 1e-20)
    C_true = M_half_true / (M_true[-1] + 1e-20)
    
    mass_scaling.append({
        'M_total_obs': M_obs[-1]/Msun, 'M_total_true': M_true[-1]/Msun,
        'M_half_obs': M_half_obs/Msun, 'M_half_true': M_half_true/Msun,
        'concentration_obs': C_obs, 'concentration_true': C_true,
        'mass_ratio': M_true[-1]/(M_obs[-1]+1e-20)
    })
    total_mass_obs[g] = M_obs[-1]/Msun
    total_mass_true[g] = M_true[-1]/Msun
    mass_concentration_obs[g] = C_obs
    mass_concentration_true[g] = C_true

# --- 3.3 Deflection Angle Analysis ---
deflection_stats = []
max_alpha_obs = np.zeros(numGalaxies)
max_alpha_true = np.zeros(numGalaxies)

for g in range(numGalaxies):
    alpha_obs = Results[g]['alpha_rad_obs'] / arcsec
    alpha_true = Results[g]['alpha_rad_true'] / arcsec
    
    max_alpha_obs[g] = np.max(alpha_obs)
    max_alpha_true[g] = np.max(alpha_true)
    
    deflection_stats.append({
        'max_alpha_obs_arcsec': max_alpha_obs[g],
        'max_alpha_true_arcsec': max_alpha_true[g],
        'alpha_ratio': max_alpha_true[g]/(max_alpha_obs[g]+1e-20)
    })

# --- 3.4 Surface Density Analysis ---
surface_stats = []
Sigma_central_obs = np.zeros(numGalaxies)
Sigma_central_true = np.zeros(numGalaxies)
Sigma_scale_obs = np.zeros(numGalaxies)
Sigma_scale_true = np.zeros(numGalaxies)

for g in range(numGalaxies):
    R_obs = Results[g]['R_grid_obs']
    R_true = Results[g]['R_grid_true']
    Sigma_obs = Results[g]['Sigma_obs'] * pc**2 / Msun
    Sigma_true = Results[g]['Sigma_true'] * pc**2 / Msun
    
    Sigma_central_obs[g] = Sigma_obs[0]
    Sigma_central_true[g] = Sigma_true[0]
    
    Sigma_half_obs = Sigma_obs[0] / 2
    Sigma_half_true = Sigma_true[0] / 2
    
    idx_half_obs = np.argmin(np.abs(Sigma_obs - Sigma_half_obs))
    idx_half_true = np.argmin(np.abs(Sigma_true - Sigma_half_true))
    
    R_scale_obs = R_obs[idx_half_obs] / pc
    R_scale_true = R_true[idx_half_true] / pc
    
    R_eff = R_obs[-1] / 2
    idx_eff = np.argmin(np.abs(R_obs - R_eff))
    Sigma_int_obs = np.trapezoid(Sigma_obs[:idx_eff], R_obs[:idx_eff]) / 1e6
    Sigma_int_true = np.trapezoid(Sigma_true[:idx_eff], R_true[:idx_eff]) / 1e6
    
    surface_stats.append({
        'Sigma0_obs': Sigma_central_obs[g], 'Sigma0_true': Sigma_central_true[g],
        'R_scale_obs_pc': R_scale_obs, 'R_scale_true_pc': R_scale_true,
        'Sigma_ratio': Sigma_central_true[g]/(Sigma_central_obs[g]+1e-20)
    })
    Sigma_scale_obs[g] = R_scale_obs
    Sigma_scale_true[g] = R_scale_true

# --- 3.5 Outer Region Statistics & Fig 3.5.1 ---
improvement_percents = np.array([g['improvement_percent'] for g in galaxySummary])
significant_galaxies = np.argsort(improvement_percents)[::-1][:min(3, numGalaxies)]

all_rms_true_outer = np.full(numGalaxies, np.nan)
all_R2_true_outer = np.full(numGalaxies, np.nan)
chi2_true_outer = np.full(numGalaxies, np.nan)
outer_fraction = np.zeros(numGalaxies)
num_outer_points = np.zeros(numGalaxies)

for g in range(numGalaxies):
    r_norm = Results[g]['r_obs'] / np.max(Results[g]['r_obs'])
    v_obs = Results[g]['v_obs']
    v_true = Results[g]['v_Newton_true_To_obs']
    
    outer_idx = np.where(r_norm > 0.3)[0]
    num_outer_points[g] = len(outer_idx)
    outer_fraction[g] = len(outer_idx) / len(r_norm)
    
    if len(outer_idx) >= 3:
        v_obs_o, v_true_o = v_obs[outer_idx], v_true[outer_idx]
        res_o = v_obs_o - v_true_o
        all_rms_true_outer[g] = np.sqrt(np.mean(res_o**2)) / 1e3
        
        SS_tot_o = np.sum((v_obs_o - np.mean(v_obs_o))**2)
        all_R2_true_outer[g] = 1 - np.sum(res_o**2)/(SS_tot_o + 1e-20)
        
        sigma_v_o = 0.05 * v_obs_o
        chi2_true_outer[g] = np.sum(((v_obs_o - v_true_o)/sigma_v_o)**2) / (len(outer_idx)-1)

# Fig 3.5.1: Deflection Correlation
fig, ax = plt.subplots(figsize=(7.5, 6.5))
min_val = min(np.min(max_alpha_obs), np.min(max_alpha_true)) * 0.95
max_val = max(np.max(max_alpha_obs), np.max(max_alpha_true)) * 1.05
ax.plot([min_val, max_val], [min_val, max_val], 'k--', linewidth=2, label='1:1 line')

sc = ax.scatter(max_alpha_obs, max_alpha_true, s=100, c=mass_colors_global[color_idx_global],
                marker='o', edgecolors='k', linewidths=0.8, zorder=3)

idx_obs_max = np.argmax(max_alpha_obs)
idx_mezzi_max = np.argmax(max_alpha_true)
x_range, y_range = max_val-min_val, max_val-min_val
ax.text(max_alpha_obs[idx_obs_max]+0.6*x_range, max_alpha_true[idx_obs_max], 
        galaxy_names[idx_obs_max], fontsize=10, fontweight='bold', ha='right',
        bbox=dict(facecolor='white', alpha=0.8, edgecolor='none'))
ax.text(max_alpha_obs[idx_mezzi_max]+0.6*x_range, max_alpha_true[idx_mezzi_max]+0.03*y_range, 
        galaxy_names[idx_mezzi_max], fontsize=10, fontweight='bold', ha='right',
        bbox=dict(facecolor='white', alpha=0.8, edgecolor='none'))

norm = plt.Normalize(vmin=np.min(log_masses), vmax=np.max(log_masses))
sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])  
cbar = plt.colorbar(sm, ax=ax, location='right', pad=0.02)
cbar.set_label('$\\log_{10}(M_{obs}/M_\\odot)$', fontsize=12, fontweight='bold')
tick_values = np.linspace(np.min(log_masses), np.max(log_masses), 5)
cbar.set_ticks(tick_values)   
cbar.set_ticklabels([f"{x:.1f}" for x in tick_values])  

ax.set_xlim([min_val, max_val])
ax.set_ylim([min_val, max_val])
ax.set_aspect('equal')
ax.set_xlabel('Observed $\\alpha_{\\max}$ (arcsec)', fontsize=20, fontweight='bold')
ax.set_ylabel('True $\\alpha_{\\max}$ (arcsec)', fontsize=20, fontweight='bold')
ax.grid(True, ls='--', alpha=0.5)
ax.tick_params(labelsize=12)
fig.tight_layout()
fig.savefig(os.path.join(ScalingDir, 'Fig_3_5_1_Deflection_Correlation.png'), dpi=400, bbox_inches='tight')
plt.close(fig)

# --- 3.6 Scaling Factor Impact ---
Sigma0_obs = np.array([s['Sigma0_obs'] for s in surface_stats])
Sigma0_true = np.array([s['Sigma0_true'] for s in surface_stats])
R_s_obs = np.array([s['R_scale_obs_pc'] for s in surface_stats])
R_s_true = np.array([s['R_scale_true_pc'] for s in surface_stats])

Ratio_Sigma = Sigma0_true / (Sigma0_obs + 1e-20)
Ratio_Radius = R_s_true / (R_s_obs + 1e-20)
Ratio_Sigma_predicted = (R_s_obs / (R_s_true + 1e-20))**2
Consistency_residual = Ratio_Sigma / (Ratio_Sigma_predicted + 1e-20)

mass_edges = [7.0, 9.0, 10.5, 12.0]
mass_labels = ['Low-Mass', 'Intermediate', 'High-Mass']
regime_colors = np.array([[0.8, 0.4, 0.2], [0.2, 0.6, 0.8], [0.4, 0.2, 0.6]])

global_stats = {
    'mean_Sigma': np.mean(Ratio_Sigma), 'std_Sigma': np.std(Ratio_Sigma),
    'mean_Radius': np.mean(Ratio_Radius), 'std_Radius': np.std(Ratio_Radius),
    'mean_Consistency': np.mean(Consistency_residual), 'std_Consistency': np.std(Consistency_residual)
}

regime_stats_36 = []
for r in range(3):
    idx = (log_masses >= mass_edges[r]) & (log_masses < mass_edges[r+1])
    n = np.sum(idx)
    reg = {'count': n, 'mass_range': [np.min(log_masses[idx]) if n>0 else np.nan, np.max(log_masses[idx]) if n>0 else np.nan]}
    if n >= 3:
        reg['mean_Radius_ratio'] = np.mean(Ratio_Radius[idx])
        reg['std_Radius_ratio'] = np.std(Ratio_Radius[idx])
        reg['median_Radius_ratio'] = np.median(Ratio_Radius[idx])
        reg['mean_Sigma_ratio'] = np.mean(Ratio_Sigma[idx])
        reg['std_Sigma_ratio'] = np.std(Ratio_Sigma[idx])
        reg['median_Sigma_ratio'] = np.median(Ratio_Sigma[idx])
        r_c, p_c = pearsonr(log_masses[idx], Ratio_Radius[idx])
        reg['corr_Radius_Mass'] = r_c
        reg['pval_Radius_Mass'] = p_c
        reg['mean_Consistency'] = np.mean(Consistency_residual[idx])
        reg['std_Consistency'] = np.std(Consistency_residual[idx])
        reg['extreme_radius'] = np.where(idx & ((Ratio_Radius > 2) | (Ratio_Radius < 0.5)))[0]
        reg['extreme_density'] = np.where(idx & ((Ratio_Sigma > 2) | (Ratio_Sigma < 0.5)))[0]
    else:
        for k in ['mean_Radius_ratio','std_Radius_ratio','median_Radius_ratio','mean_Sigma_ratio',
                  'std_Sigma_ratio','median_Sigma_ratio','corr_Radius_Mass','pval_Radius_Mass',
                  'mean_Consistency','std_Consistency','extreme_radius','extreme_density']:
            reg[k] = np.nan
    regime_stats_36.append(reg)

# Global Regression for 3.6.2
x = log_masses
y = Ratio_Radius
p_global = np.polyfit(x, y, 1)
slope_global, intercept_global = p_global[0], p_global[1]
R_global, P_global = np.corrcoef(x, y)[0,1], None # placeholder
# Bootstrap
rng = np.random.default_rng(42)
boot_slopes = []
for _ in range(500):
    idx_b = rng.choice(numGalaxies, size=numGalaxies, replace=True)
    boot_slopes.append(np.polyfit(x[idx_b], y[idx_b], 1)[0])
ci_slope = np.percentile(boot_slopes, [2.5, 97.5])

# Fig 3.6.2: Mass-Regime Scatter
fig, ax = plt.subplots(figsize=(8.5, 6.5))
for r in range(3):
    idx = (x >= mass_edges[r]) & (x < mass_edges[r+1])
    if np.any(idx):
        ax.scatter(x[idx], y[idx], s=60, c=[regime_colors[r]], marker='o',
                   alpha=0.7, edgecolors='k', linewidths=0.8,
                   label=f"{mass_labels[r]} (N={np.sum(idx)}, mean={np.mean(y[idx]):.2f})")

ax.axhline(1, color='k', ls='--', linewidth=1.5)
for r in range(3):
    idx = (x >= mass_edges[r]) & (x < mass_edges[r+1])
    if np.any(idx):
        x_anno, y_anno = np.mean(x[idx]), np.mean(y[idx])
        offset_x = -0.1*(np.max(x[idx])-np.min(x[idx])) - 0.15
        offset_y = 0.15*(np.max(y[idx])-np.min(y[idx])) + 0.3
        ax.text(x_anno+offset_x, y_anno+offset_y, f"mean = {np.mean(y[idx]):.2f}",
                color=regime_colors[r], fontweight='bold', fontsize=9,
                bbox=dict(facecolor='white', edgecolor=regime_colors[r], pad=3),
                ha='center', va='bottom')

ax.legend(loc='upper left', fontsize=9, frameon=True, edgecolor='k')
ax.set_xlabel('$\\log_{10}(M_{obs}/M_\\odot)$', fontsize=22, fontweight='bold')
ax.set_ylabel('$R_s^{True} / R_s^{Obs}$', fontsize=22, fontweight='bold')
ax.set_yscale('log')
ax.set_xlim([7.2, 12.1])
ax.set_ylim([0.6, 9.3])
ax.set_yticks([0.6, 0.7, 0.8, 0.9, 1, 1.2, 1.5, 2, 2.5, 3, 4, 5, 6, 7, 8, 9])
ax.grid(True, ls='--', alpha=0.5)
ax.tick_params(labelsize=11)
fig.tight_layout()
fig.savefig(os.path.join(ScalingDir, 'Fig_3_6_2_RadiusRatio_vs_Mass.png'), dpi=400, bbox_inches='tight')
plt.close(fig)

# Fig 3.6.3: Mass Regime Consistency
log_R_ratio = np.log10(Ratio_Radius)
log_Sigma_ratio = np.log10(Ratio_Sigma)
p_obs, intercept_obs, r_val, p_val, se_obs = linregress(log_R_ratio, log_Sigma_ratio)
observed_slope = p_obs

fig, ax = plt.subplots(figsize=(8.5, 6.5))
for r in range(3):
    idx = (log_masses >= mass_edges[r]) & (log_masses < mass_edges[r+1])
    if np.any(idx):
        ax.scatter(log_R_ratio[idx], log_Sigma_ratio[idx], s=60, c=[regime_colors[r]],
                   marker='o', alpha=0.7, edgecolors='k', linewidths=0.8,
                   label=f"{mass_labels[r]} (N={np.sum(idx)})")
    # Regime slope lines
    if np.sum(idx) >= 3:
        slope_r, intercept_r, _, _, _ = linregress(log_R_ratio[idx], log_Sigma_ratio[idx])
        x_reg = np.linspace(np.min(log_R_ratio[idx]), np.max(log_R_ratio[idx]), 50)
        ax.plot(x_reg, slope_r*x_reg + intercept_r, '--', color=regime_colors[r], linewidth=1.5)

x_line = np.linspace(np.min(log_R_ratio)-0.3, np.max(log_R_ratio)+0.3, 100)
ax.plot(x_line, intercept_obs + observed_slope*x_line, 'k-', linewidth=2)
ax.text(np.max(x_line)-0.4, intercept_obs+observed_slope*np.max(x_line),
        f"global slope = {observed_slope:.2f}", color='k', fontweight='bold', fontsize=9,
        bbox=dict(facecolor='white', edgecolor='k', pad=3), ha='left', va='top')

ax.legend(loc='upper left', fontsize=12, frameon=True, edgecolor='k')
ax.set_xlabel('$\\log_{10}(R_s^{True}/R_s^{obs})$', fontsize=22, fontweight='bold')
ax.set_ylabel('$\\log_{10}(\\Sigma_0^{True}/\\Sigma_0^{obs})$', fontsize=22, fontweight='bold')
ax.set_xlim(np.min(log_R_ratio)-0.2, np.max(log_R_ratio)+0.2)
ax.set_ylim(np.min(log_Sigma_ratio)-0.2, np.max(log_Sigma_ratio)+0.4)
ax.grid(True, ls='--', alpha=0.5)
ax.tick_params(labelsize=11)
fig.tight_layout()
fig.savefig(os.path.join(ScalingDir, 'Fig_3_6_3_MassRegime_Consistency.png'), dpi=400, bbox_inches='tight')
plt.close(fig)

# --- 3.6.4 Figure: Mass-Stratified Boxplots ---
x_positions = [1, 2.5, 4]
valid_positions = []
valid_colors = []
valid_labels = []
radius_data = []
sigma_data = []

# 1. Manually group data by mass regime (Matplotlib expects a list of arrays)
for b in range(3):
    idx = (log_masses >= mass_edges[b]) & (log_masses < mass_edges[b+1])
    if np.any(idx):
        radius_data.append(Ratio_Radius[idx])
        sigma_data.append(Ratio_Sigma[idx])
        valid_positions.append(x_positions[b])
        valid_colors.append(regime_colors[b])
        valid_labels.append(mass_labels[b])

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11.5, 5.8))

# 2. Robust styling function for multi-colored boxplots
def style_boxplot(ax, bp, positions, colors):
    for patch, pos, col in zip(bp['boxes'], positions, colors):
        patch.set_facecolor([*col, 0.35])
        patch.set_edgecolor(col)
        patch.set_linewidth(1.5)
    for line in bp['whiskers'] + bp['caps']:
        x_mid = np.mean(line.get_xdata())
        idx = np.argmin(np.abs(np.array(positions) - x_mid))
        line.set_color(colors[idx])
        line.set_linewidth(1.5)
    for median in bp['medians']:
        median.set_color('k')
        median.set_linewidth(2)
    for flier in bp['fliers']:
        x_mid = np.mean(flier.get_xdata())
        idx = np.argmin(np.abs(np.array(positions) - x_mid))
        flier.set_markerfacecolor(colors[idx])
        flier.set_markeredgecolor(colors[idx])
        flier.set_markersize(4)

# Panel 1: Radius Ratio
bp1 = ax1.boxplot(radius_data, positions=valid_positions, widths=0.5, 
                  patch_artist=True, showfliers=True)
style_boxplot(ax1, bp1, valid_positions, valid_colors)
ax1.axhline(1, color='k', ls='--', linewidth=2, alpha=0.6)
for pos, vals, col in zip(valid_positions, radius_data, valid_colors):
    if len(vals) > 0:
        ax1.text(pos, np.max(vals)*1.1, f"N={len(vals)}\n{np.mean(vals):.2f}±{np.std(vals):.2f}",
                 ha='center', fontsize=9, bbox=dict(facecolor='white', edgecolor=col, pad=3))
ax1.set_yscale('log')
ax1.set_xticks(valid_positions)
ax1.set_xticklabels(valid_labels)
ax1.set_ylabel('$R_s^{True} / R_s^{obs}$', fontsize=14, fontweight='bold')
ax1.grid(True, ls='--', alpha=0.5)
ax1.tick_params(labelsize=12)

# Panel 2: Density Ratio
bp2 = ax2.boxplot(sigma_data, positions=valid_positions, widths=0.5, 
                  patch_artist=True, showfliers=True)
style_boxplot(ax2, bp2, valid_positions, valid_colors)
ax2.axhline(1, color='k', ls='--', linewidth=2, alpha=0.6)
for pos, vals, col in zip(valid_positions, sigma_data, valid_colors):
    if len(vals) > 0:
        ax2.text(pos, np.max(vals)*1.13, f"N={len(vals)}\n{np.mean(vals):.2f}±{np.std(vals):.2f}",
                 ha='center', fontsize=9, bbox=dict(facecolor='white', edgecolor=col, pad=3))
ax2.set_yscale('log')
ax2.set_xticks(valid_positions)
ax2.set_xticklabels(valid_labels)
ax2.set_ylabel('$\\Sigma_0^{True} / \\Sigma_0^{obs}$', fontsize=14, fontweight='bold')
ax2.grid(True, ls='--', alpha=0.5)
ax2.tick_params(labelsize=12)

fig.tight_layout()
fig.savefig(os.path.join(ScalingDir, 'Fig_3_6_4_MassStratified_Summary.png'), dpi=400, bbox_inches='tight')
plt.close(fig)

# --- 3.7 Lensing-Dynamics Connection & Fig 3.7.2 ---
v_flat_obs = np.array([g['v_flat_obs_kms'] for g in galaxySummary])
v_flat_true = np.array([g['v_flat_true_kms'] for g in galaxySummary])
M_baryonic = np.array([g['M_total_obs_Msun'] for g in galaxySummary])
M_true_all = np.array([g['M_total_true_Msun'] for g in galaxySummary])
alpha_max_obs = np.array([d['max_alpha_obs_arcsec'] for d in deflection_stats])
alpha_max_true = np.array([d['max_alpha_true_arcsec'] for d in deflection_stats])

g_bar_obs, g_bar_true = np.zeros(numGalaxies), np.zeros(numGalaxies)
g_dyn_obs, g_dyn_true = np.zeros(numGalaxies), np.zeros(numGalaxies)

for g in range(numGalaxies):
    r_char_o = np.max(Results[g]['r_obs'])/2 if np.max(Results[g]['r_obs'])>0 else Results[g]['r_obs'][len(Results[g]['r_obs'])//2]
    r_char_t = np.max(Results[g]['r_true'])/2 if np.max(Results[g]['r_true'])>0 else Results[g]['r_obs'][len(Results[g]['r_obs'])//2]
    g_bar_obs[g] = G * M_baryonic[g] * Msun / r_char_o**2
    g_bar_true[g] = G * M_true_all[g] * Msun / r_char_t**2
    g_dyn_obs[g] = (v_flat_obs[g]*1e3)**2 / r_char_o
    g_dyn_true[g] = (v_flat_true[g]*1e3)**2 / r_char_t

valid_cross = np.isfinite(v_flat_obs) & (v_flat_obs>0) & np.isfinite(alpha_max_true) & (alpha_max_true>0)
log_v_obs = np.log10(v_flat_obs[valid_cross])
log_alpha_true = np.log10(alpha_max_true[valid_cross])
p_cross = np.polyfit(log_v_obs, log_alpha_true, 1)
power_law_index_cross, alpha_v_intercept_cross = p_cross[0], p_cross[1]

fig, ax = plt.subplots(figsize=(8, 6.5))
for r in range(3):
    idx = (log_masses >= mass_edges[r]) & (log_masses < mass_edges[r+1]) & valid_cross
    if np.any(idx):
        ax.scatter(v_flat_obs[idx], alpha_max_true[idx], s=90, c=[regime_colors[r]],
                   marker='o', alpha=0.75, edgecolors='k', linewidths=1.2,
                   label=f"{mass_labels[r]} (N={np.sum(idx)})")

v_theory = np.linspace(np.min(v_flat_obs[valid_cross]), np.max(v_flat_obs[valid_cross]), 100)
alpha_fit = 10**alpha_v_intercept_cross * v_theory**power_law_index_cross
ax.plot(v_theory, alpha_fit, 'k-', linewidth=2.5, label=f"Fit: $\\alpha \\propto v_{{obs}}^{{{power_law_index_cross:.2f}}}$")

ax.set_xscale('log')
ax.set_yscale('log')
ax.set_xlabel('$v_{obs}$ (km/s)', fontsize=30, fontweight='bold')
ax.set_ylabel('True $\\alpha_{max}$ (arcsec)', fontsize=30, fontweight='bold')
ax.legend(loc='upper left', fontsize=10, frameon=True, edgecolor='k')
ax.grid(True, ls='--', alpha=0.5)
ax.tick_params(labelsize=11)
fig.tight_layout()
fig.savefig(os.path.join(ScalingDir, 'Fig_3_7_2_TrueDeflection_vs_ObservedVelocity.png'), dpi=400, bbox_inches='tight')
plt.close(fig)

# --- 3.8 Rank Order & Scaling Verification ---
log_Sigma0_obs = np.log10(np.array([s['Sigma0_obs'] for s in surface_stats]))
log_Sigma0_true = np.log10(np.array([s['Sigma0_true'] for s in surface_stats]))
log_M_obs = np.log10(np.array([m['M_total_obs'] for m in mass_scaling]))

rank_stats = []
for r in range(3):
    idx = (log_M_obs >= mass_edges[r]) & (log_M_obs < mass_edges[r+1])
    n = np.sum(idx)
    if n < 3:
        rank_stats.append({'n': n, 'slope': np.nan, 'intercept': np.nan, 'R2': np.nan, 'sigma_1to1': np.nan, 'pearson_r': np.nan})
        continue
    p_r = np.polyfit(log_Sigma0_obs[idx], log_Sigma0_true[idx], 1)
    y_fit = np.polyval(p_r, log_Sigma0_obs[idx])
    R2_r = 1 - np.sum((log_Sigma0_true[idx]-y_fit)**2)/np.sum((log_Sigma0_true[idx]-np.mean(log_Sigma0_true[idx]))**2)
    sigma_1to1 = np.std(log_Sigma0_true[idx] - log_Sigma0_obs[idx])
    r_c, _ = pearsonr(log_Sigma0_obs[idx], log_Sigma0_true[idx])
    rank_stats.append({'n': n, 'slope': p_r[0], 'intercept': p_r[1], 'R2': R2_r, 'sigma_1to1': sigma_1to1, 'pearson_r': r_c})

# Fig 3.8.1
fig, ax = plt.subplots(figsize=(9, 7))
for r in range(3):
    idx = (log_M_obs >= mass_edges[r]) & (log_M_obs < mass_edges[r+1])
    if np.any(idx):
        lab = f"{mass_labels[r]}: slope={rank_stats[r]['slope']:.3f}, r={rank_stats[r]['pearson_r']:.3f} (N={rank_stats[r]['n']})"
        ax.scatter(log_Sigma0_obs[idx], log_Sigma0_true[idx], s=65, c=[regime_colors[r]],
                   marker='o', alpha=0.75, edgecolors='k', linewidths=0.8, label=lab)
        if rank_stats[r]['n'] >= 3:
            x_fit = np.linspace(np.min(log_Sigma0_obs[idx]), np.max(log_Sigma0_obs[idx]), 50)
            ax.plot(x_fit, rank_stats[r]['intercept'] + rank_stats[r]['slope']*x_fit, '--', color=regime_colors[r], linewidth=1.5)

x_lim = [np.min(log_Sigma0_obs)-0.15, np.max(log_Sigma0_obs)+0.15]
ax.plot(x_lim, x_lim, 'k--', linewidth=2, label='1:1 line')
ax.legend(loc='lower right', fontsize=9, frameon=True, edgecolor='k')
ax.set_xlabel('$\\log_{10}(\\Sigma_0^{obs} / M_\\odot \\, pc^{-2})$', fontsize=30, fontweight='bold')
ax.set_ylabel('$\\log_{10}(\\Sigma_0^{true} / M_\\odot \\, pc^{-2})$', fontsize=30, fontweight='bold')
ax.grid(True, ls='--', alpha=0.5)
ax.set_aspect('equal', 'box')
ax.tick_params(labelsize=11)
fig.tight_layout()
fig.savefig(os.path.join(ScalingDir, 'Fig_3_8_1_RankOrderPreservation.png'), dpi=400, bbox_inches='tight')
plt.close(fig)

# Fig 3.8.2
R_ratio = np.array([s['R_scale_true_pc'] for s in surface_stats]) / (np.array([s['R_scale_obs_pc'] for s in surface_stats])+1e-20)
M_ratio = np.array([m['mass_ratio'] for m in mass_scaling])
log_dev = np.log10(M_ratio) - np.log10(R_ratio)
p_fit = np.polyfit(np.log10(R_ratio), np.log10(M_ratio), 1)
measured_exp = p_fit[0]

fig, ax = plt.subplots(figsize=(8, 6))
sc = ax.scatter(R_ratio, M_ratio, s=80, c=mass_colors_global[color_idx_global], marker='o', edgecolors='k')
R_range = np.linspace(np.min(R_ratio)*0.9, np.max(R_ratio)*1.1, 200)
ax.plot(R_range, R_range, 'k--', linewidth=2, label='1:1 Line')
ax.plot(R_range, 10**(p_fit[1]) * R_range**measured_exp, 'r-', linewidth=2, label=f"Fit (slope={measured_exp:.2f})")
ax.set_xscale('log')
ax.set_yscale('log')
ax.set_xlim(1, 10)
ax.set_ylim(1, 10)
ax.set_xlabel('$R_s^{True}/R_s^{obs}$', fontsize=16, fontweight='bold')
ax.set_ylabel('$M_{true}/M_{obs}$', fontsize=16, fontweight='bold')
ax.grid(True, ls='--', alpha=0.5)
ax.legend()
norm = plt.Normalize(vmin=np.min(log_masses), vmax=np.max(log_masses))
sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])  
cbar = plt.colorbar(sm, ax=ax, location='right', pad=0.02)
cbar.set_label('$\\log_{10}(M_{obs}/M_\\odot)$', fontsize=12, fontweight='bold')
tick_values = np.linspace(np.min(log_masses), np.max(log_masses), 5)
cbar.set_ticks(tick_values)   
cbar.set_ticklabels([f"{x:.1f}" for x in tick_values])  
ax.text(0.13, 0.69, f"N = {numGalaxies} galaxies\nExponent = {measured_exp:.3f}\nMean log-dev = {np.mean(log_dev):.3f} dex\nStd dev = {np.std(log_dev):.3f} dex",
        transform=ax.transAxes, fontsize=12, va='top', bbox=dict(facecolor='white', alpha=0.9, edgecolor='gray'))
fig.tight_layout()
fig.savefig(os.path.join(ScalingDir, 'Fig_3_8_2_Scaling_Verification.png'), dpi=300, bbox_inches='tight')
plt.close(fig)

# Fig 3.8.3 & 3.8.4 (Helper function to avoid repetition)
def plot_ratio_vs_zeta(z_data, y_data, y_label, filename):
    mask = np.isfinite(z_data) & np.isfinite(y_data) & (z_data>0) & (y_data>0)
    z, y = z_data[mask], y_data[mask]
    lg = log_masses[mask]
    sort_idx = np.argsort(lg)
    color_idx = np.argsort(sort_idx)
    colors = cmap(np.linspace(0,1,len(z)))[color_idx]
    
    fig, ax = plt.subplots(figsize=(8, 6))
    ax.scatter(z, y, s=80, c=colors, marker='o', edgecolors='k')
    p = np.polyfit(np.log10(z), np.log10(y), 1)
    z_fit = np.linspace(np.min(z)*0.9, np.max(z)*1.1, 200)
    ax.plot(z_fit, 10**p[1]*z_fit**p[0], 'r-', linewidth=2, label=f"Fit: slope={p[0]:.2f}")
    r_c = np.corrcoef(np.log10(z), np.log10(y))[0,1]
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.set_xlabel('Mezzi Scale Factor $\\zeta$', fontsize=22, fontweight='bold')
    ax.set_ylabel(y_label, fontsize=22, fontweight='bold')
    ax.grid(True, ls='--', alpha=0.5)
    ax.legend(loc='lower left', fontsize=10, frameon=True)
    norm = plt.Normalize(vmin=np.min(log_masses), vmax=np.max(log_masses))
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])  
    cbar = plt.colorbar(sm, ax=ax, location='right', pad=0.02)
    cbar.set_label('$\\log_{10}(M_{obs}/M_\\odot)$', fontsize=12, fontweight='bold')
    tick_values = np.linspace(np.min(log_masses), np.max(log_masses), 5)
    cbar.set_ticks(tick_values)   
    cbar.set_ticklabels([f"{x:.1f}" for x in tick_values])  
    cbar.set_ticks(tick_values)   
    cbar.set_ticklabels([f"{x:.1f}" for x in tick_values])  
    ax.text(0.05, 0.85, f"N = {len(z)} galaxies\nSlope = {p[0]:.3f} (r = {r_c:.3f})",
            transform=ax.transAxes, fontsize=12, va='top', bbox=dict(facecolor='white', alpha=0.9, edgecolor='gray'))
    fig.tight_layout()
    fig.savefig(os.path.join(ScalingDir, filename), dpi=400, bbox_inches='tight')
    plt.close(fig)

plot_ratio_vs_zeta(Mezzi_means, M_ratio, '$M_{true}/M_{obs}$', 'Fig_3_8_3_Mass_vs_Zeta.png')
plot_ratio_vs_zeta(Mezzi_means, alpha_max_true/(alpha_max_obs+1e-20), '$\\alpha_{true}/\\alpha_{obs}$', 'Fig_3_8_4_Deflection_vs_Zeta.png')

print("[DONE] Section 3: Comprehensive Scaling Analysis Complete.")


# ===== SECTION 4: RAR VALIDATION ANALYSIS =====
print("[4.0] Starting RAR Validation Analysis...")
import numpy as np
import matplotlib.pyplot as plt
from scipy.io import savemat

a0_RAR = 1.20e-10
g_dagger = a0_RAR

def RAR_empirical(g):
    return g / (1 - np.exp(-np.sqrt(g / a0_RAR)))

def RAR_newtonian(g_bar):
    return g_bar

def RAR_mond(g_bar):
    return np.sqrt(g_bar * g_dagger)

RAR_Analysis = []
valid_galaxies = []

# Precompute galaxy colors by mass to avoid scope issues
cmap = plt.colormaps['viridis']
mass_rank = np.argsort(np.argsort(log_masses))
galaxy_colors = cmap(np.linspace(0, 1, numGalaxies))[mass_rank]

for g in range(numGalaxies):
    r_obs = Results[g]['r_obs'].ravel()
    r_true = Results[g]['r_true'].ravel()
    M_obs = Results[g]['M_enc_obs'].ravel()
    M_true = Results[g]['M_enc_true'].ravel()
    v_obs = Results[g]['v_obs'].ravel()
    v_Newton_true = Results[g]['v_Newton_true_To_obs'].ravel()
    Mezzi_th = Results[g]['Mezzi_th'].ravel()

    N = min(len(r_obs), len(r_true), len(M_obs), len(M_true), len(v_obs), len(v_Newton_true))
    r_obs, r_true, M_obs, M_true, v_obs, v_Newton_true, Mezzi_th = [a[:N] for a in [r_obs, r_true, M_obs, M_true, v_obs, v_Newton_true, Mezzi_th]]

    valid = (r_obs > 1e3) & (M_obs > 0) & (r_true > 1e3) & (M_true > 0) & \
            (v_obs > 0) & (v_Newton_true > 0) & \
            np.isfinite(r_obs) & np.isfinite(M_obs) & np.isfinite(r_true) & np.isfinite(M_true) & \
            np.isfinite(v_obs) & np.isfinite(v_Newton_true)
            
    if np.sum(valid) < 5:
        RAR_Analysis.append({'valid': False, 'galaxy': galaxy_names[g]})
        continue

    r_o, r_t = r_obs[valid], r_true[valid]
    M_o, M_t = M_obs[valid], M_true[valid]
    v_o, v_t = v_obs[valid], v_Newton_true[valid]
    Mezzi_local = Mezzi_th[valid]

    # Case 1: Newtonian Consistency (True Frame)
    g_bar_C1 = G * M_t / r_t**2
    g_dyn_C1 = v_t**2 / r_t
    
    # Case 2: Empirical Mapping (Mixed Frame)
    g_bar_C2 = G * M_t / r_t**2
    g_dyn_C2 = v_o**2 / r_o
    
    # Case 3: Baryonic Shift (Geometry only)
    g_bar_C3 = G * M_t / r_o**2
    g_dyn_C3 = v_o**2 / r_o
    
    # Observed
    g_bar_obs = G * M_o / r_o**2
    g_dyn_obs = v_o**2 / r_o

    def calc_scatter(gd, gp): return np.std(np.log10(gd) - np.log10(gp))

    RAR_Analysis.append({
        'valid': True, 'galaxy': galaxy_names[g], 'N': np.sum(valid),
        'g_bar_C1': g_bar_C1, 'g_dyn_C1': g_dyn_C1,
        'g_bar_C2': g_bar_C2, 'g_dyn_C2': g_dyn_C2,
        'g_bar_C3': g_bar_C3, 'g_dyn_C3': g_dyn_C3,
        'g_bar_obs': g_bar_obs, 'g_dyn_obs': g_dyn_obs,
        'Mezzi': Mezzi_local, 'r_norm': r_o / np.max(r_o),
        'scatter_C1': calc_scatter(g_dyn_C1, g_bar_C1),
        'scatter_C2': calc_scatter(g_dyn_C2, RAR_empirical(g_bar_C2)),
        'scatter_C3': calc_scatter(g_dyn_C3, RAR_mond(g_bar_C3)),
        'scatter_obs': calc_scatter(g_dyn_obs, RAR_empirical(g_bar_obs))
    })
    valid_galaxies.append(g)

# --- 4.2 Global Analysis ---
all_data = {'C1': ([], []), 'C2': ([], []), 'C3': ([], []), 'Obs': ([], [])}
for g in valid_galaxies:
    res = RAR_Analysis[g]
    mask = (res['g_bar_obs'] > 1e-15) & (res['g_bar_obs'] < 1e-5) & (res['g_dyn_obs'] > 1e-15) & (res['g_dyn_obs'] < 1e-5)
    all_data['C1'][0].append(res['g_bar_C1'][mask]); all_data['C1'][1].append(res['g_dyn_C1'][mask])
    all_data['C2'][0].append(res['g_bar_C2'][mask]); all_data['C2'][1].append(res['g_dyn_C2'][mask])
    all_data['C3'][0].append(res['g_bar_C3'][mask]); all_data['C3'][1].append(res['g_dyn_C3'][mask])
    all_data['Obs'][0].append(res['g_bar_obs'][mask]); all_data['Obs'][1].append(res['g_dyn_obs'][mask])

# Flatten arrays
for key in all_data:
    all_data[key] = (np.concatenate(all_data[key][0]), np.concatenate(all_data[key][1]))

def global_scatter(g_bar, g_dyn, ref_func):
    valid = (g_bar > 1e-15) & (g_bar < 1e-5) & (g_dyn > 1e-15) & (g_dyn < 1e-5)
    return np.std(np.log10(g_dyn[valid]) - np.log10(ref_func(g_bar[valid])))

scatter_global = {
    'C1': global_scatter(*all_data['C1'], RAR_newtonian),
    'C2': global_scatter(*all_data['C2'], RAR_empirical),
    'C3': global_scatter(*all_data['C3'], RAR_mond),
    'Obs': global_scatter(*all_data['Obs'], RAR_empirical)
}

print(f"[RESULT] Global scatter metrics:\n  C1: {scatter_global['C1']:.3f}\n  C2: {scatter_global['C2']:.3f}\n  C3: {scatter_global['C3']:.3f}\n  Obs: {scatter_global['Obs']:.3f}")

# --- 4.3 Generate Figures ---
rarDir = os.path.join(baseDir, 'RAR_Figures')
os.makedirs(rarDir, exist_ok=True)
g_ref = np.logspace(-2, 2, 100) * 1e-10

cases = [
    ('C1', 'Case 1 - Newtonian Consistency', [0.2, 0.4, 0.8], 
     'DEF:\n$g_{bar}=GM_{True}/r_{True}^2$\n$g_{dyn}=v_{Newton,True}^2/r_{True}$', 
     'Scatter vs Newtonian', 'C1'),
    ('C2', 'Case 2 - Empirical Mapping', [0.9, 0.5, 0.1],
     'DEF:\n$g_{bar}=GM_{True}/r_{True}^2$\n$g_{dyn}=v_{obs}^2/r_{obs}$',
     'Scatter vs RAR', 'C2'),
    ('C3', 'Case 3 - Baryonic Scale Shift', [0.6, 0.2, 0.6],
     'DEF:\n$g_{bar}=GM_{True}/r_{obs}^2$\n$g_{dyn}=v_{obs}^2/r_{obs}$',
     'Scatter vs deep-MOND', 'C3')
]

for key, title, color, defn, scat_label, scat_key in cases:
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.scatter(all_data['Obs'][0]*1e10, all_data['Obs'][1]*1e10, s=15, c=[0.6]*3, alpha=0.25, label='Observed frame')
    ax.scatter(all_data[key][0]*1e10, all_data[key][1]*1e10, s=20, c=[color], alpha=0.7, label=title.split(' - ')[1])
    
    ax.plot(g_ref*1e10, g_ref*1e10, 'k--', lw=3, label='Newtonian')
    ax.plot(g_ref*1e10, RAR_empirical(g_ref)*1e10, 'k-', lw=3, label='Empirical RAR')
    ax.plot(g_ref*1e10, np.sqrt(g_ref*1e10)*np.sqrt(1.2), 'k-.', lw=2.5, label='deep-MOND')
    
    ax.set_xscale('log'); ax.set_yscale('log')
    ax.set_xlim(0.01, 100); ax.set_ylim(0.01, 100)
    ax.set_xlabel(r'$g_{bar} [10^{-10} m/s^2]$', fontsize=16, fontweight='bold')
    ax.set_ylabel(r'$g_{dyn} [10^{-10} m/s^2]$', fontsize=16, fontweight='bold')
    ax.legend(loc='lower right', fontsize=10)
    ax.grid(True, ls='--', alpha=0.5)
    
    ax.text(0.02, 0.98, f"{title}\n\n{defn}\n\n{scat_label}: {scatter_global[scat_key]:.3f} dex",
            transform=ax.transAxes, fontsize=11, fontfamily='monospace', va='top',
            bbox=dict(facecolor='white', edgecolor=color, linewidth=1.5))
            
    fig.savefig(os.path.join(rarDir, f'Fig_4_3_{key}.png'), dpi=300, bbox_inches='tight')
    plt.close(fig)

# ===== 4.4 C2 Residual Analysis =====
all_g_bar_C2, all_g_dyn_C2 = all_data['C2']
residuals = np.log10(all_g_dyn_C2) - np.log10(RAR_empirical(all_g_bar_C2))
valid_res = np.isfinite(residuals) & np.isfinite(all_g_bar_C2)
all_g_bar_C2, all_g_dyn_C2, residuals = all_g_bar_C2[valid_res], all_g_dyn_C2[valid_res], residuals[valid_res]

# Radius-matched data reconstruction for Fig 4.4
all_r_res, all_res_r, all_mass_res, all_gal_idx_res = [], [], [], []
for g in valid_galaxies:
    r_orig = Results[g]['r_obs']
    M_orig = Results[g]['M_enc_obs']
    v_orig = Results[g]['v_obs']
    
    valid_mask = (r_orig > 1e3) & (M_orig > 0) & np.isfinite(r_orig) & np.isfinite(M_orig) & np.isfinite(v_orig)
    if not np.any(valid_mask): continue
        
    r_v, M_v, v_v = r_orig[valid_mask], M_orig[valid_mask], v_orig[valid_mask]
    g_bar_v = G * M_v / r_v**2
    g_dyn_v = v_v**2 / r_v
    acc_valid = (g_bar_v > 1e-15) & (g_bar_v < 1e-5) & (g_dyn_v > 1e-15) & (g_dyn_v < 1e-5)
    
    if np.sum(acc_valid) < 5: continue
    r_m, g_bar_m, g_dyn_m = r_v[acc_valid], g_bar_v[acc_valid], g_dyn_v[acc_valid]
    res_m = np.log10(g_dyn_m) - np.log10(RAR_empirical(g_bar_m))
    
    valid_final = np.isfinite(res_m) & np.isfinite(r_m) & (r_m > 1e3)
    n_pts = np.sum(valid_final)
    if n_pts == 0: continue
        
    all_r_res.extend(r_m[valid_final]/1e3)
    all_res_r.extend(res_m[valid_final])
    all_mass_res.extend([galaxySummary[g]['M_total_obs_Msun']] * n_pts)
    all_gal_idx_res.extend([g] * n_pts) 

all_r_res = np.array(all_r_res)
all_res_r = np.array(all_res_r)
all_mass_res = np.array(all_mass_res)
all_gal_idx_res = np.array(all_gal_idx_res)

if len(all_res_r) > 50:
    fig, ax = plt.subplots(figsize=(9, 6))
    norm = plt.Normalize(vmin=np.min(log_masses), vmax=np.max(log_masses))
    point_colors = galaxy_colors[all_gal_idx_res]
    
    ax.scatter(all_res_r, all_r_res, s=20, c=point_colors, alpha=0.6)
    ax.axvline(0, color='r', ls='--', lw=1.5)
    ax.axvline(np.median(all_res_r), color='b', ls=':', lw=1.5, label=f"Median: {np.median(all_res_r):.2f}")
    ax.set_yscale('log')
    ax.set_xlabel(r'Residual: $\log_{10}(g_{obs}/g_{RAR})$ [dex]', fontsize=12, fontweight='bold')
    ax.set_ylabel(r'$r_{obs}$ [km]', fontsize=14, fontweight='bold')
    ax.grid(True, ls='--', alpha=0.5)
    ax.legend(fontsize=10)
    
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])
    cbar = plt.colorbar(sm, ax=ax)
    cbar.set_label(r'$\log_{10}(M_{obs}/M_\odot)$', fontsize=14, fontweight='bold')
    cbar.set_ticks(np.linspace(0, 1, 5))
    cbar.set_ticklabels([f"{x:.1f}" for x in np.linspace(np.min(log_masses), np.max(log_masses), 5)])
    
    ax.text(0.65, 0.5, f"N = {len(all_res_r)} points\nσ = {np.std(all_res_r):.3f} dex\nMedian = {np.median(all_res_r):+.2f} dex",
            transform=ax.transAxes, fontsize=11, fontfamily='monospace', va='center',
            bbox=dict(facecolor='white', edgecolor=[0.9, 0.5, 0.1], linewidth=1.5))
            
    fig.savefig(os.path.join(rarDir, 'Fig_4_4_Res_vs_Radius.png'), dpi=300, bbox_inches='tight')
    plt.close(fig)

# ===== 4.5 Acceleration Ratio Analysis =====
ratioDir = os.path.join(baseDir, 'AccelerationRatioAnalysis')
os.makedirs(ratioDir, exist_ok=True)
RatioAnalysis = []

for g in valid_galaxies:
    res = RAR_Analysis[g]
    g_bar_C = res['g_bar_C2']
    g_dyn_C = res['g_dyn_C2']
    r_norm = res['r_norm']
    
    R_ratio = g_dyn_C**2 / (g_bar_C + 1e-20)
    valid_R = (g_bar_C > 1e-15) & (R_ratio > 1e-15) & (R_ratio < 1e-5) & np.isfinite(R_ratio)
    if np.sum(valid_R) < 5: continue
    
    g_bar_v, g_dyn_v, R_v, r_v = g_bar_C[valid_R], g_dyn_C[valid_R], R_ratio[valid_R], r_norm[valid_R]
    slope, intercept = np.polyfit(np.log10(g_bar_v), np.log10(R_v), 1)[:2]
    
    classif = 'Transition'
    if slope > 0.7: classif = 'Newtonian'
    elif slope < 0.3: classif = 'MOND'
        
    RatioAnalysis.append({
        'valid': True, 'galaxy': galaxy_names[g], 'g_bar': g_bar_v, 'g_dyn': g_dyn_v,
        'R_ratio': R_v, 'R_over_a0': R_v/a0_RAR, 'r_norm': r_v, 'slope': slope, 'class': classif
    })

# Ensemble Data for Fig 4.5
all_r_ens, all_R_ens, all_mass_ens = [], [], []
for g_idx, r in enumerate(RatioAnalysis):
    n_pts = len(r['r_norm'])
    all_r_ens.extend(r['r_norm'])
    all_R_ens.extend(r['R_over_a0'])
    all_mass_ens.extend([galaxySummary[valid_galaxies[g_idx]]['M_total_obs_Msun']] * n_pts)

all_r_ens = np.array(all_r_ens)
all_R_ens = np.array(all_R_ens)
all_mass_ens = np.array(all_mass_ens)

if len(all_r_ens) > 0:
    fig, ax = plt.subplots(figsize=(12, 6))
    
    # Normalize mass values for coloring
    norm = plt.Normalize(vmin=np.min(log_masses), vmax=np.max(log_masses))
    point_colors = cmap(norm(np.log10(all_mass_ens)))
    
    ax.scatter(all_r_ens, all_R_ens, s=15, c=point_colors, alpha=0.6)
    ax.axhline(1, color='r', ls='--', lw=2, label='MOND: R/a₀ = 1')
    ax.axhline(10, color='k', ls=':', lw=2, label='Newtonian: R/a₀ = 10')
    
    ax.set_yscale('log')
    ax.set_xlabel(r'Normalized Radius $r/r_{max}$', fontsize=16, fontweight='bold')
    ax.set_ylabel(r'$R/a_0 = (g_{dyn}^2/g_{bar})/a_0$', fontsize=16, fontweight='bold')
    ax.legend(loc='best', fontsize=11)
    ax.grid(True, ls='--', alpha=0.5)
    
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])  
    
    cbar = plt.colorbar(sm, ax=ax, location='right', pad=0.02)
    cbar.set_label(r'$\log_{10}(M_{obs}/M_\odot)$', fontsize=14, fontweight='bold')
    
    tick_values = np.linspace(np.min(log_masses), np.max(log_masses), 5)
    cbar.set_ticks(tick_values)   
    cbar.set_ticklabels([f"{x:.1f}" for x in tick_values])   
    
    fig.savefig(os.path.join(ratioDir, 'Fig_4_5_Radial_Profiles.png'), dpi=300, bbox_inches='tight')
    plt.close(fig)

print("[DONE] Section 4: RAR Validation & Ratio Analysis Complete.")


# ===== SECTION 5: BTFR ANALYSIS =====
print("[5.0] Starting BTFR Analysis...")
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import pearsonr
from scipy.io import savemat

# Reference slopes
slope_newtonian = 2.0
slope_btfr = 4.0

BTFR = []
valid_btfr_galaxies = []

for g in range(numGalaxies):
    r_obs = Results[g]['r_obs'].ravel()
    r_true = Results[g]['r_true'].ravel()
    M_obs = Results[g]['M_enc_obs'].ravel()
    M_true = Results[g]['M_enc_true'].ravel()
    v_obs = Results[g]['v_obs'].ravel()
    v_Newton_true = Results[g]['v_Newton_true_To_obs'].ravel()
    Mezzi_th = Results[g]['Mezzi_th'].ravel()
    
    N = min(len(r_obs), len(r_true), len(M_obs), len(M_true), len(v_obs), len(v_Newton_true), len(Mezzi_th))
    r_obs, r_true, M_obs, M_true, v_obs, v_Newton_true, Mezzi_th = [a[:N] for a in [r_obs, r_true, M_obs, M_true, v_obs, v_Newton_true, Mezzi_th]]
    
    valid = (r_obs > 1e3*pc) & (M_obs > 0) & (r_true > 0) & (M_true > 0) & \
            (v_obs > 0) & (v_Newton_true > 0) & np.isfinite(Mezzi_th) & \
            (Mezzi_th > 0) & (Mezzi_th <= 1)
    
    if np.sum(valid) < 10:
        BTFR.append({'valid': False, 'galaxy': galaxy_names[g]})
        continue
    
    r_o, r_t = r_obs[valid], r_true[valid]
    M_o, M_t = M_obs[valid], M_true[valid]
    v_o, v_t = v_obs[valid], v_Newton_true[valid]
    mezzi = Mezzi_th[valid]
    
    outer_fraction = 0.3
    r_thresh_o = outer_fraction * np.max(r_o)
    r_thresh_t = outer_fraction * np.max(r_t)
    outer_mask_o = r_o >= r_thresh_o
    outer_mask_t = r_t >= r_thresh_t
    
    if np.sum(outer_mask_o) < 5 or np.sum(outer_mask_t) < 5:
        BTFR.append({'valid': False, 'galaxy': galaxy_names[g]})
        continue
    
    M_tot_obs = M_o[-1] / Msun
    M_tot_true = M_t[-1] / Msun
    v_flat_obs = np.mean(v_o[outer_mask_o]) / 1e3
    v_flat_true_T = np.mean(v_t[outer_mask_t]) / 1e3
    v_flat_true_E = v_flat_obs  # Empirical mapping
    
    BTFR.append({
        'valid': True, 'galaxy': galaxy_names[g], 'N_outer': np.sum(valid),
        'tot': {
            'logM_obs': np.log10(M_tot_obs), 'logv_obs': np.log10(v_flat_obs),
            'logM_T': np.log10(M_tot_true), 'logv_T': np.log10(v_flat_true_T),
            'logM_E': np.log10(M_tot_true), 'logv_E': np.log10(v_flat_true_E)
        },
        'Mezzi': np.mean(np.log10(mezzi[outer_mask_t])),
        'M_ratio': M_tot_true / M_tot_obs
    })
    valid_btfr_galaxies.append(g)
    print(f"[BTFR] Galaxy {g+1}/{numGalaxies}: {galaxy_names[g]}")

# --- 5.2 Global Analysis ---
tot_logM_T = np.array([BTFR[g]['tot']['logM_T'] for g in valid_btfr_galaxies])
tot_logv_T = np.array([BTFR[g]['tot']['logv_T'] for g in valid_btfr_galaxies])
tot_logM_E = np.array([BTFR[g]['tot']['logM_E'] for g in valid_btfr_galaxies])
tot_logv_E = np.array([BTFR[g]['tot']['logv_E'] for g in valid_btfr_galaxies])

N_gal = len(tot_logM_T)

# Fits
p_tot_T = np.polyfit(tot_logv_T, tot_logM_T, 1)
p_tot_E = np.polyfit(tot_logv_E, tot_logM_E, 1)

# Bootstrap CIs
rng = np.random.default_rng(42)
n_boot = 1000
boot_tot_T, boot_tot_E = [], []
for _ in range(n_boot):
    idx = rng.choice(N_gal, size=N_gal, replace=True)
    if len(idx) > 3:
        boot_tot_T.append(np.polyfit(tot_logv_T[idx], tot_logM_T[idx], 1)[0])
        boot_tot_E.append(np.polyfit(tot_logv_E[idx], tot_logM_E[idx], 1)[0])
ci_tot_T = np.percentile(boot_tot_T, [2.5, 97.5])
ci_tot_E = np.percentile(boot_tot_E, [2.5, 97.5])

# Scatter metrics
def calc_scatter_ref(logM, logv, slope):
    return np.std(logM - (slope * logv + np.mean(logM - slope*logv)))

s_tot_T_s2 = calc_scatter_ref(tot_logM_T, tot_logv_T, 2.0)
s_tot_T_s4 = calc_scatter_ref(tot_logM_T, tot_logv_T, 4.0)
s_tot_E_s2 = calc_scatter_ref(tot_logM_E, tot_logv_E, 2.0)
s_tot_E_s4 = calc_scatter_ref(tot_logM_E, tot_logv_E, 4.0)

Global_BTFR_Results = {
    'tot': {
        'slope_T': p_tot_T[0], 'slope_E': p_tot_E[0],
        'ci_T': ci_tot_T, 'ci_E': ci_tot_E,
        'scatter_s2': [s_tot_T_s2, s_tot_E_s2],
        'scatter_s4': [s_tot_T_s4, s_tot_E_s4],
        'N_gal': N_gal
    },
    'N_galaxies': N_gal
}

# --- 5.3 Generate Figures ---
btfrDir = os.path.join(baseDir, 'BTFR_Figures')
os.makedirs(btfrDir, exist_ok=True)

case_names = ['True Frame', 'Empirical Mapping']
case_colors = np.array([[0.2, 0.4, 0.8], [0.9, 0.5, 0.1]])
scenario_desc = ['M_{True} & v_{Newton} at r_{True}', 'M_{True}(r_{True}) vs v_{obs}(r_{obs})']
s_tot_values = [s_tot_T_s4, s_tot_E_s4]
slopes_vals = [p_tot_T[0], p_tot_E[0]]
ci_vals = np.array([ci_tot_T, ci_tot_E])

# Fig 5.3.8: BTFR Slopes
fig, ax = plt.subplots(figsize=(10, 7))
x_pos = [1, 2]
for i in range(2):
    ax.errorbar(x_pos[i], slopes_vals[i], 
                yerr=[[slopes_vals[i]-ci_vals[i,0]], [ci_vals[i,1]-slopes_vals[i]]],
                fmt='o', linewidth=2.5, markersize=14, 
                markerfacecolor=case_colors[i], color=case_colors[i]*0.7, capsize=10,
                label=f"{case_names[i]} (scatter={s_tot_values[i]:.3f})")
    ax.text(x_pos[i]+0.1, slopes_vals[i], f"{slopes_vals[i]:.2f}", 
            ha='left', fontsize=10, fontweight='bold', color=case_colors[i], va='center')

ax.axhline(2.0, ls='--', color=[0.2, 0.7, 0.3], lw=2.5, alpha=0.6, label='Newtonian (M ∝ v²)')
ax.axhline(4.0, ls='--', color=[0.8, 0.2, 0.2], lw=2.5, alpha=0.6, label='BTFR (M ∝ v⁴)')
ax.text(0.5, 2.15, 'Newtonian (M ∝ v²)', color=[0.2, 0.7, 0.3], fontsize=10, fontweight='bold')
ax.text(0.5, 4.15, 'BTFR (M ∝ v⁴)', color=[0.8, 0.2, 0.2], fontsize=10, fontweight='bold')

ax.set_xticks(x_pos)
ax.set_xticklabels(case_names, fontsize=11)
ax.set_ylabel('Fitted Slope (log M vs log v)', fontsize=13, fontweight='bold')
for i in range(2):
    ax.text(x_pos[i], 0.2, scenario_desc[i], ha='center', fontsize=8, color=[0.4]*3)
ax.set_ylim([0, 6.5])
ax.set_xlim([0.5, 2.5])
ax.grid(True, ls='--', alpha=0.5)
ax.legend(loc='upper center', fontsize=10, edgecolor=[0.5]*3)
fig.tight_layout()
fig.savefig(os.path.join(btfrDir, 'Fig_5_3_8_BTFR_Slopes_Total_Mass.png'), dpi=300, bbox_inches='tight')
plt.close(fig)

# Fig 5.3.9: Log V vs Log M Scatter
fig, (ax_T, ax_E) = plt.subplots(1, 2, figsize=(12, 5.5))
v_range = np.linspace(np.min([tot_logv_T, tot_logv_E]), np.max([tot_logv_T, tot_logv_E]), 100)

# True Frame
ax_T.scatter(tot_logv_T, tot_logM_T, s=50, c=[case_colors[0]], alpha=0.7, edgecolors='none', label=f'Scatter vs v⁴ = {s_tot_T_s4:.2f} dex')
ax_T.plot(v_range, np.polyval(p_tot_T, v_range), 'k-', lw=2.5, label=f"Fit (slope={p_tot_T[0]:.2f})")
v_ref = np.linspace(np.min(tot_logv_T), np.max(tot_logv_T), 100)
ax_T.plot(v_ref, 2.0*v_ref + np.mean(tot_logM_T - 2.0*tot_logv_T), '--', color=[0.2,0.7,0.3], lw=2, label='Newtonian')
ax_T.plot(v_ref, 4.0*v_ref + np.mean(tot_logM_T - 4.0*tot_logv_T), '--', color=[0.8,0.2,0.2], lw=2, label='BTFR')
ax_T.text(0.05, 0.85, scenario_desc[0], transform=ax_T.transAxes, fontsize=11, fontweight='bold', 
          color=case_colors[0], bbox=dict(facecolor='white', alpha=0.8, edgecolor='none'))
ax_T.set_xlabel('log₁₀(v_flat) [km/s]', fontsize=16)
ax_T.set_ylabel('log₁₀(M_tot) [M_☉]', fontsize=16)
ax_T.grid(True, ls='--', alpha=0.5)
ax_T.legend(loc='lower right', fontsize=11)

# Empirical Mapping
ax_E.scatter(tot_logv_E, tot_logM_E, s=50, c=[case_colors[1]], alpha=0.7, edgecolors='none', label=f'Scatter vs v⁴ = {s_tot_E_s4:.2f} dex')
ax_E.plot(v_range, np.polyval(p_tot_E, v_range), 'k-', lw=2.5, label=f"Fit (slope={p_tot_E[0]:.2f})")
v_ref = np.linspace(np.min(tot_logv_E), np.max(tot_logv_E), 100)
ax_E.plot(v_ref, 2.0*v_ref + np.mean(tot_logM_E - 2.0*tot_logv_E), '--', color=[0.2,0.7,0.3], lw=2, label='Newtonian')
ax_E.plot(v_ref, 4.0*v_ref + np.mean(tot_logM_E - 4.0*tot_logv_E), '--', color=[0.8,0.2,0.2], lw=2, label='BTFR')
ax_E.text(0.05, 0.85, scenario_desc[1], transform=ax_E.transAxes, fontsize=11, fontweight='bold', 
          color=case_colors[1], bbox=dict(facecolor='white', alpha=0.8, edgecolor='none'))
ax_E.set_xlabel('log₁₀(v_flat) [km/s]', fontsize=16)
ax_E.set_ylabel('log₁₀(M_tot) [M_☉]', fontsize=16)
ax_E.grid(True, ls='--', alpha=0.5)
ax_E.legend(loc='lower right', fontsize=11)

fig.tight_layout()
fig.savefig(os.path.join(btfrDir, 'Fig_5_3_9_LogV_LogM_Scatter.png'), dpi=300, bbox_inches='tight')
plt.close(fig)

# --- 5.3.10-17: β-Offset vs Mezzi Scale Factor ---
residuals_obs = np.array([BTFR[g]['tot']['logM_obs'] - 4*BTFR[g]['tot']['logv_obs'] for g in valid_btfr_galaxies])
residuals_E = np.array([BTFR[g]['tot']['logM_E'] - 4*BTFR[g]['tot']['logv_E'] for g in valid_btfr_galaxies])
Mezzi_values = np.array([10**BTFR[g]['Mezzi'] for g in valid_btfr_galaxies])
galaxy_masses_btfr = np.array([galaxySummary[g]['M_total_obs_Msun'] for g in valid_btfr_galaxies])

log_masses_btfr = np.log10(galaxy_masses_btfr)
mass_bin_edges = [8.5, 9.5, 10.5, 12.0]
mass_bin_labels = ['Low Mass', 'Intermediate', 'High Mass']
mass_bin_colors = np.array([[0.8, 0.3, 0.2], [0.6, 0.4, 0.8], [0.2, 0.5, 0.8]])

mass_category = np.digitize(log_masses_btfr, mass_bin_edges[1:-1]) + 1
mass_category[log_masses_btfr < mass_bin_edges[0]] = 1
mass_category[log_masses_btfr >= mass_bin_edges[-1]] = 3

def compute_fit(alpha, res):
    valid = np.isfinite(alpha) & np.isfinite(res)
    if np.sum(valid) < 3:
        return {'slope': [np.nan, np.nan], 'r': np.nan, 'mean_res': np.nan, 'std_res': np.nan, 'N': 0}
    slope, intercept = np.polyfit(np.log10(alpha[valid]), res[valid], 1)
    r_val, _ = pearsonr(np.log10(alpha[valid]), res[valid])
    return {'slope': [slope, intercept], 'r': r_val, 'mean_res': np.mean(res[valid]), 
            'std_res': np.std(res[valid]), 'N': np.sum(valid)}

valid_obs = np.isfinite(Mezzi_values) & np.isfinite(residuals_obs)
valid_E = np.isfinite(Mezzi_values) & np.isfinite(residuals_E)
stats_obs = compute_fit(Mezzi_values[valid_obs], residuals_obs[valid_obs])
stats_E = compute_fit(Mezzi_values[valid_E], residuals_E[valid_E])

# Bin stats
bin_stats_obs = {'sigma': [], 'scatter': [], 'N': [], 'centroid_zeta': [], 'centroid_res': []}
bin_stats_E = {'sigma': [], 'scatter': [], 'N': [], 'centroid_zeta': [], 'centroid_res': []}
for b in range(3):
    cat_mask = (mass_category == b+1)
    # Observed
    valid_bin = valid_obs & cat_mask
    if np.sum(valid_bin) >= 3:
        bin_stats_obs['sigma'].append(np.std(residuals_obs[valid_bin]))
        bin_stats_obs['N'].append(np.sum(valid_bin))
        bin_stats_obs['centroid_zeta'].append(np.median(Mezzi_values[valid_bin]))
        bin_stats_obs['centroid_res'].append(np.mean(residuals_obs[valid_bin]))
        fit_vals = np.polyval(stats_obs['slope'], np.log10(Mezzi_values[valid_bin]))
        bin_stats_obs['scatter'].append(np.std(residuals_obs[valid_bin] - fit_vals))
    else:
        for k in bin_stats_obs: bin_stats_obs[k].append(np.nan)
    # Empirical
    valid_bin = valid_E & cat_mask
    if np.sum(valid_bin) >= 3:
        bin_stats_E['sigma'].append(np.std(residuals_E[valid_bin]))
        bin_stats_E['N'].append(np.sum(valid_bin))
        bin_stats_E['centroid_zeta'].append(np.median(Mezzi_values[valid_bin]))
        bin_stats_E['centroid_res'].append(np.mean(residuals_E[valid_bin]))
        fit_vals = np.polyval(stats_E['slope'], np.log10(Mezzi_values[valid_bin]))
        bin_stats_E['scatter'].append(np.std(residuals_E[valid_bin] - fit_vals))
    else:
        for k in bin_stats_E: bin_stats_E[k].append(np.nan)

# Fig 5.3.16A: Observed Frame β-Offset vs ζ
def plot_beta_offset(ax, Mezzi_vals, residuals, stats, bin_stats, valid_mask, title, ylabel, colors):
    for b in range(3):
        cat_mask = (mass_category == b+1) & valid_mask
        if np.sum(cat_mask) > 0:
            ax.scatter(Mezzi_vals[cat_mask], residuals[cat_mask], s=80, c=[colors[b]], 
                      alpha=0.8, edgecolors='white', linewidths=0.5)
    
    if stats['N'] > 0:
        xfit = np.linspace(np.min(np.log10(Mezzi_vals[valid_mask])), 
                          np.max(np.log10(Mezzi_vals[valid_mask])), 100)
        yfit = np.polyval(stats['slope'], xfit)
        ax.plot(10**xfit, yfit, '--k', lw=2)
    
    ax.set_xscale('log')
    ax.set_xlabel('Mezzi Scale Factor ζ', fontsize=16)
    ax.set_ylabel(ylabel, fontsize=16)
    ax.grid(True, ls='--', alpha=0.5)
    
    # Tags with scatter values
    y_all = np.concatenate([residuals[valid_obs], residuals_E[valid_E]])
    y_range = np.max(y_all) - np.min(y_all) if len(y_all) > 1 else 1
    base_offset = 0.3 * y_range
    tag_positions = []
    for b in range(3):
        if not np.isnan(bin_stats['scatter'][b]):
            tag_x = bin_stats['centroid_zeta'][b]
            sign = 1 if b % 2 == 0 else -1
            tag_y = bin_stats['centroid_res'][b] + sign * base_offset
            if tag_positions:
                if np.min(np.abs(tag_y - np.array(tag_positions)[:,1])) < 0.5*base_offset:
                    tag_y += sign * base_offset
            tag_positions.append([tag_x, tag_y])
            ax.text(tag_x, tag_y, f"scatter={bin_stats['scatter'][b]:.2f} dex", 
                   ha='center', fontsize=9, fontweight='bold', color=colors[b],
                   bbox=dict(facecolor='white', alpha=0.9, edgecolor=colors[b], pad=3))
    
    # Legend
    legend_text = [
        f"Fit: slope={stats['slope'][0]:.2f}, int={stats['slope'][1]:.2f}, r={stats['r']:.2f}",
        f"Mean β = {stats['mean_res']:.2f} dex"
    ] + [f"{mass_bin_labels[b]} (N={bin_stats['N'][b]}): scatter={bin_stats['scatter'][b]:.2f} dex" 
         for b in range(3) if not np.isnan(bin_stats['scatter'][b])]
    ax.legend(legend_text, loc='best', fontsize=10, frameon=True, edgecolor=[0.5]*3)

fig, ax_obs = plt.subplots(figsize=(9.5, 5.5))
plot_beta_offset(ax_obs, Mezzi_values, residuals_obs, stats_obs, bin_stats_obs, valid_obs,
                'Observed Frame', 'Offset β = log₁₀M_Obs - 4log₁₀v_Obs [dex]', mass_bin_colors)
fig.tight_layout()
fig.savefig(os.path.join(btfrDir, 'Fig_5_3_16A_BetaOffset_Observed.png'), dpi=300, bbox_inches='tight')
plt.close(fig)

# Fig 5.3.17E: Empirical Mapping β-Offset vs ζ
fig, ax_E = plt.subplots(figsize=(9.5, 5.5))
plot_beta_offset(ax_E, Mezzi_values, residuals_E, stats_E, bin_stats_E, valid_E,
                'Empirical Mapping', 'Offset β = log₁₀M_True - 4log₁₀v_Obs [dex]', mass_bin_colors)
fig.tight_layout()
fig.savefig(os.path.join(btfrDir, 'Fig_5_3_17E_BetaOffset_True.png'), dpi=300, bbox_inches='tight')
plt.close(fig)

# Save results
save_dict = {}
if 'prep_for_savemat' in globals():
    # prep_for_savemat returns a dict, so we pass it directly without indexing
    save_dict['Global_BTFR_Results'] = prep_for_savemat([Global_BTFR_Results])
    save_dict['BTFR'] = prep_for_savemat(BTFR)
else:
    save_dict['Global_BTFR_Results'] = Global_BTFR_Results
    save_dict['BTFR'] = BTFR

