#!/usr/bin/env python3
"""
3D+3D ROTATION CURVE SOLVER
===========================

Computes galaxy rotation curves using the 3D+3D discrete spacetime framework.
Zero free parameters per galaxy - only universal constants λ₂, β.

Features:
- SPARC database integration
- Baryonic mass modeling (disk + bulge + gas)
- Q-field Yukawa profiles
- Breathing scale effects
- Statistical analysis tools

Author: Simone Calzighetti & Lucy (Claude AI)
Date: December 2025
Version: 1.0

Theory Reference: Paper I, Section 4; Paper II, Section 12
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import quad, simpson
from scipy.interpolate import interp1d
from scipy.special import iv, kv  # Modified Bessel functions
from scipy.optimize import minimize_scalar
import warnings
warnings.filterwarnings('ignore')

# ============================================================================
# PHYSICAL CONSTANTS
# ============================================================================

class Constants:
    """Physical constants in CGS and convenient units."""
    
    # Fundamental
    G = 6.674e-8           # cm³/(g·s²)
    c = 2.998e10           # cm/s
    M_sun = 1.989e33       # g
    kpc = 3.086e21         # cm
    pc = 3.086e18          # cm
    year = 3.156e7         # s
    
    # 3D+3D parameters (FIXED - not fitted!)
    lambda_2 = 4.30        # kpc - fundamental breathing scale
    beta = 0.49            # dimensionless coupling (universal)
    M_crit = 2.43e10       # M_sun - critical mass
    
    # Compactification radii
    L4 = 15.1              # light-years
    L5 = 9.6               # light-years
    
    # Derived Q-field masses
    m2_eV = 1.47e-24       # eV
    m3_eV = 2.32e-24       # eV
    
    # Golden ratio
    phi = (1 + np.sqrt(5)) / 2  # 1.618...

# ============================================================================
# BARYONIC MASS MODELS
# ============================================================================

class ExponentialDisk:
    """Freeman exponential disk model."""
    
    def __init__(self, M_disk, R_d, h_z=0.3):
        """
        Parameters:
        -----------
        M_disk : float
            Total disk mass [M_sun]
        R_d : float
            Scale radius [kpc]
        h_z : float
            Vertical scale height [kpc], default 0.3
        """
        self.M_disk = M_disk * Constants.M_sun
        self.R_d = R_d * Constants.kpc
        self.h_z = h_z * Constants.kpc
        self.Sigma_0 = self.M_disk / (2 * np.pi * self.R_d**2)
    
    def surface_density(self, R):
        """Surface density Σ(R) [g/cm²]."""
        x = R / self.R_d
        return self.Sigma_0 * np.exp(-x)
    
    def V_circular(self, R):
        """
        Circular velocity from disk [cm/s].
        
        Uses Freeman (1970) formula with Bessel functions.
        """
        y = R / (2 * self.R_d)
        
        # Avoid numerical issues at small R
        if np.isscalar(R):
            if R < 0.01 * self.R_d:
                return 0.0
        else:
            result = np.zeros_like(R)
            mask = R >= 0.01 * self.R_d
            y = y[mask] if hasattr(y, '__len__') else y
        
        # Bessel function combination
        I0 = iv(0, y)
        I1 = iv(1, y)
        K0 = kv(0, y)
        K1 = kv(1, y)
        
        V_sq = 4 * np.pi * Constants.G * self.Sigma_0 * self.R_d * y**2 * (
            I0 * K0 - I1 * K1
        )
        
        if np.isscalar(R):
            return np.sqrt(max(V_sq, 0))
        else:
            result[mask] = np.sqrt(np.maximum(V_sq, 0))
            return result


class HernquistBulge:
    """Hernquist spherical bulge model."""
    
    def __init__(self, M_bulge, a):
        """
        Parameters:
        -----------
        M_bulge : float
            Total bulge mass [M_sun]
        a : float
            Scale radius [kpc]
        """
        self.M_bulge = M_bulge * Constants.M_sun
        self.a = a * Constants.kpc
    
    def M_enclosed(self, r):
        """Enclosed mass [g]."""
        return self.M_bulge * (r / (r + self.a))**2
    
    def V_circular(self, R):
        """Circular velocity [cm/s]."""
        M_enc = self.M_enclosed(R)
        return np.sqrt(Constants.G * M_enc / R)


class GasDisk:
    """HI gas disk with exponential profile."""
    
    def __init__(self, M_gas, R_gas):
        """
        Parameters:
        -----------
        M_gas : float
            Total gas mass [M_sun]
        R_gas : float
            Gas scale radius [kpc]
        """
        self.M_gas = M_gas * Constants.M_sun
        self.R_gas = R_gas * Constants.kpc
        self.Sigma_0 = self.M_gas / (2 * np.pi * self.R_gas**2)
    
    def V_circular(self, R):
        """Circular velocity from gas [cm/s]."""
        y = R / (2 * self.R_gas)
        
        if np.isscalar(R):
            if R < 0.01 * self.R_gas:
                return 0.0
        
        I0 = iv(0, y)
        I1 = iv(1, y)
        K0 = kv(0, y)
        K1 = kv(1, y)
        
        V_sq = 4 * np.pi * Constants.G * self.Sigma_0 * self.R_gas * y**2 * (
            I0 * K0 - I1 * K1
        )
        
        return np.sqrt(np.maximum(V_sq, 0))


# ============================================================================
# Q-FIELD CONTRIBUTION
# ============================================================================

class QFieldContribution:
    """
    Q-field contribution to rotation curve.
    
    The Q₂ and Q₃ fields create additional gravitational effects
    through their coupling to baryonic matter.
    """
    
    def __init__(self, M_baryonic, lambda_b=Constants.lambda_2, beta=Constants.beta):
        """
        Parameters:
        -----------
        M_baryonic : float
            Total baryonic mass [M_sun]
        lambda_b : float
            Breathing scale [kpc]
        beta : float
            Coupling constant
        """
        self.M_b = M_baryonic * Constants.M_sun
        self.lambda_b = lambda_b * Constants.kpc
        self.beta = beta
    
    def form_factor(self, x):
        """
        Form factor F_Q(x) where x = R/λ_b.
        
        F_Q(x) = 1 - (1 + x)exp(-x)
        
        Encodes Yukawa suppression of Q-field.
        """
        return 1.0 - (1.0 + x) * np.exp(-x)
    
    def V_Q_squared(self, R, M_enclosed):
        """
        Q-field velocity squared contribution [cm²/s²].
        
        V²_Q = β² × GM_enc/R × F_Q(R/λ_b)
        """
        x = R / self.lambda_b
        F = self.form_factor(x)
        
        return self.beta**2 * Constants.G * M_enclosed / R * F
    
    def enhancement_factor(self, R, M_enclosed):
        """
        Total gravitational enhancement factor.
        
        g_total/g_Newton = 1 + β² × F_Q(R/λ_b)
        """
        x = R / self.lambda_b
        F = self.form_factor(x)
        
        return 1.0 + self.beta**2 * F


# ============================================================================
# COMPLETE ROTATION CURVE MODEL
# ============================================================================

class RotationCurve3D3D:
    """
    Complete 3D+3D rotation curve model.
    
    V²_rot = V²_bar + V²_Q
    
    where V²_bar = V²_disk + V²_bulge + V²_gas (baryonic)
    and V²_Q = Q-field contribution
    """
    
    def __init__(self, 
                 M_disk=1e10, R_disk=3.0,
                 M_bulge=0, R_bulge=0.5,
                 M_gas=1e9, R_gas=6.0,
                 lambda_b=Constants.lambda_2,
                 beta=Constants.beta):
        """
        Parameters (all masses in M_sun, radii in kpc):
        """
        # Initialize components
        self.disk = ExponentialDisk(M_disk, R_disk) if M_disk > 0 else None
        self.bulge = HernquistBulge(M_bulge, R_bulge) if M_bulge > 0 else None
        self.gas = GasDisk(M_gas, R_gas) if M_gas > 0 else None
        
        # Total baryonic mass
        self.M_total = (M_disk + M_bulge + M_gas) * Constants.M_sun
        
        # Q-field
        M_bar_total = M_disk + M_bulge + M_gas
        self.Q_field = QFieldContribution(M_bar_total, lambda_b, beta)
        
        # Store parameters
        self.params = {
            'M_disk': M_disk,
            'R_disk': R_disk,
            'M_bulge': M_bulge,
            'R_bulge': R_bulge,
            'M_gas': M_gas,
            'R_gas': R_gas,
            'lambda_b': lambda_b,
            'beta': beta
        }
    
    def V_baryonic(self, R):
        """
        Total baryonic rotation velocity [km/s].
        """
        R_cm = R * Constants.kpc
        
        V_sq = 0.0
        
        if self.disk is not None:
            V_d = self.disk.V_circular(R_cm)
            V_sq += V_d**2
        
        if self.bulge is not None:
            V_b = self.bulge.V_circular(R_cm)
            V_sq += V_b**2
        
        if self.gas is not None:
            V_g = self.gas.V_circular(R_cm)
            V_sq += V_g**2
        
        # Convert to km/s
        return np.sqrt(V_sq) / 1e5
    
    def M_enclosed(self, R):
        """
        Total enclosed baryonic mass at radius R [g].
        """
        R_cm = R * Constants.kpc
        
        M_enc = 0.0
        
        if self.disk is not None:
            # Approximate disk enclosed mass
            x = R / self.params['R_disk']
            M_enc += self.disk.M_disk * (1 - (1 + x) * np.exp(-x))
        
        if self.bulge is not None:
            M_enc += self.bulge.M_enclosed(R_cm)
        
        if self.gas is not None:
            x = R / self.params['R_gas']
            M_enc += self.gas.M_gas * (1 - (1 + x) * np.exp(-x))
        
        return M_enc
    
    def V_rotation(self, R):
        """
        Total 3D+3D rotation velocity [km/s].
        
        This is the main prediction: V²_rot = V²_bar + V²_Q
        """
        R_cm = R * Constants.kpc
        M_enc = self.M_enclosed(R)
        
        # Baryonic contribution
        V_bar = self.V_baryonic(R)
        V_bar_cgs = V_bar * 1e5
        
        # Q-field contribution
        V_Q_sq = self.Q_field.V_Q_squared(R_cm, M_enc)
        
        # Total
        V_total_sq = V_bar_cgs**2 + V_Q_sq
        
        return np.sqrt(V_total_sq) / 1e5  # km/s
    
    def compute_curve(self, R_array):
        """
        Compute rotation curve over array of radii.
        
        Returns:
        --------
        dict with V_bar, V_total, V_Q for each R
        """
        V_bar = np.array([self.V_baryonic(r) for r in R_array])
        V_total = np.array([self.V_rotation(r) for r in R_array])
        
        # Q-field contribution
        V_Q = np.sqrt(V_total**2 - V_bar**2)
        
        return {
            'R': R_array,
            'V_bar': V_bar,
            'V_total': V_total,
            'V_Q': V_Q
        }
    
    def plot(self, R_max=30, n_points=100, observed=None, figsize=(10, 6)):
        """
        Plot rotation curve.
        
        Parameters:
        -----------
        R_max : float
            Maximum radius [kpc]
        observed : dict, optional
            Dictionary with 'R', 'V', 'V_err' for observed data
        """
        R = np.linspace(0.1, R_max, n_points)
        result = self.compute_curve(R)
        
        fig, ax = plt.subplots(figsize=figsize)
        
        # Plot components
        ax.plot(R, result['V_bar'], 'b--', label='Baryonic', linewidth=2)
        ax.plot(R, result['V_Q'], 'g:', label='Q-field', linewidth=2)
        ax.plot(R, result['V_total'], 'r-', label='3D+3D Total', linewidth=2.5)
        
        # Plot observed if provided
        if observed is not None:
            ax.errorbar(observed['R'], observed['V'], yerr=observed['V_err'],
                       fmt='ko', markersize=5, capsize=3, label='Observed')
        
        # Breathing scale marker
        ax.axvline(self.params['lambda_b'], color='gray', linestyle='--', 
                   alpha=0.5, label=f'λ₂ = {self.params["lambda_b"]:.1f} kpc')
        
        ax.set_xlabel('Radius [kpc]', fontsize=12)
        ax.set_ylabel('Rotation Velocity [km/s]', fontsize=12)
        ax.set_title('3D+3D Rotation Curve (Zero Free Parameters)', fontsize=14)
        ax.legend(loc='lower right')
        ax.set_xlim(0, R_max)
        ax.set_ylim(0, None)
        ax.grid(True, alpha=0.3)
        
        plt.tight_layout()
        return fig, ax


# ============================================================================
# SPARC DATABASE INTERFACE
# ============================================================================

class SPARCGalaxy:
    """
    Interface for SPARC galaxy data.
    
    SPARC format: http://astroweb.cwru.edu/SPARC/
    """
    
    def __init__(self, name, distance, inclination,
                 R_obs, V_obs, V_err,
                 V_gas, V_disk, V_bulge=None,
                 M_L_disk=0.5, M_L_bulge=0.7):
        """
        Parameters from SPARC catalog.
        
        Parameters:
        -----------
        name : str
            Galaxy name
        distance : float
            Distance [Mpc]
        inclination : float
            Inclination [degrees]
        R_obs : array
            Observed radii [kpc]
        V_obs : array
            Observed rotation velocity [km/s]
        V_err : array
            Velocity errors [km/s]
        V_gas, V_disk, V_bulge : arrays
            Component velocities [km/s]
        M_L_disk, M_L_bulge : float
            Mass-to-light ratios
        """
        self.name = name
        self.distance = distance
        self.inclination = inclination
        
        self.R = np.array(R_obs)
        self.V_obs = np.array(V_obs)
        self.V_err = np.array(V_err)
        
        self.V_gas = np.array(V_gas) if V_gas is not None else np.zeros_like(R_obs)
        self.V_disk = np.array(V_disk) * np.sqrt(M_L_disk)
        self.V_bulge = np.array(V_bulge) * np.sqrt(M_L_bulge) if V_bulge is not None else np.zeros_like(R_obs)
        
        # Compute baryonic velocity
        self.V_bar = np.sqrt(self.V_gas**2 + self.V_disk**2 + self.V_bulge**2)
    
    def fit_3d3d(self, lambda_b=Constants.lambda_2, beta=Constants.beta):
        """
        Fit 3D+3D model to observed rotation curve.
        
        Note: This uses the UNIVERSAL parameters - no fitting per galaxy!
        
        Returns:
        --------
        dict with V_pred, residuals, RMS, chi2
        """
        # Q-field enhancement
        x = self.R / lambda_b
        F_Q = 1.0 - (1.0 + x) * np.exp(-x)
        
        # Predicted velocity
        V_pred_sq = self.V_bar**2 * (1 + beta**2 * F_Q)
        V_pred = np.sqrt(V_pred_sq)
        
        # Statistics
        residuals = self.V_obs - V_pred
        RMS = np.sqrt(np.mean(residuals**2))
        
        chi2 = np.sum((residuals / self.V_err)**2)
        chi2_red = chi2 / (len(self.R) - 2)  # Reduced chi2
        
        return {
            'V_pred': V_pred,
            'residuals': residuals,
            'RMS': RMS,
            'chi2': chi2,
            'chi2_red': chi2_red
        }
    
    def plot(self, lambda_b=Constants.lambda_2, beta=Constants.beta, figsize=(10, 6)):
        """Plot observed vs 3D+3D prediction."""
        
        result = self.fit_3d3d(lambda_b, beta)
        
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=figsize, 
                                        gridspec_kw={'height_ratios': [3, 1]})
        
        # Main plot
        ax1.errorbar(self.R, self.V_obs, yerr=self.V_err, 
                    fmt='ko', markersize=4, capsize=2, label='Observed')
        ax1.plot(self.R, self.V_bar, 'b--', label='Baryonic', linewidth=1.5)
        ax1.plot(self.R, result['V_pred'], 'r-', label='3D+3D', linewidth=2)
        
        ax1.axvline(lambda_b, color='gray', linestyle=':', alpha=0.5)
        ax1.set_ylabel('V [km/s]')
        ax1.set_title(f'{self.name} - RMS = {result["RMS"]:.1f} km/s')
        ax1.legend()
        ax1.set_xlim(0, self.R.max() * 1.1)
        
        # Residuals
        ax2.errorbar(self.R, result['residuals'], yerr=self.V_err,
                    fmt='ko', markersize=4, capsize=2)
        ax2.axhline(0, color='r', linestyle='-')
        ax2.axhline(result['RMS'], color='gray', linestyle='--', alpha=0.5)
        ax2.axhline(-result['RMS'], color='gray', linestyle='--', alpha=0.5)
        ax2.set_xlabel('R [kpc]')
        ax2.set_ylabel('Residual [km/s]')
        ax2.set_xlim(0, self.R.max() * 1.1)
        
        plt.tight_layout()
        return fig


# ============================================================================
# BATCH ANALYSIS
# ============================================================================

def analyze_sparc_batch(galaxies, lambda_b=Constants.lambda_2, beta=Constants.beta):
    """
    Analyze batch of SPARC galaxies.
    
    Parameters:
    -----------
    galaxies : list of SPARCGalaxy
        List of galaxy objects
    
    Returns:
    --------
    DataFrame-like dict with results
    """
    results = {
        'name': [],
        'RMS': [],
        'chi2_red': [],
        'N_points': [],
        'R_max': [],
        'V_flat': []
    }
    
    for gal in galaxies:
        fit = gal.fit_3d3d(lambda_b, beta)
        
        results['name'].append(gal.name)
        results['RMS'].append(fit['RMS'])
        results['chi2_red'].append(fit['chi2_red'])
        results['N_points'].append(len(gal.R))
        results['R_max'].append(gal.R.max())
        results['V_flat'].append(gal.V_obs[-1])
    
    # Summary statistics
    results['mean_RMS'] = np.mean(results['RMS'])
    results['median_RMS'] = np.median(results['RMS'])
    results['std_RMS'] = np.std(results['RMS'])
    
    return results


# ============================================================================
# DEMO / TESTING
# ============================================================================

def demo():
    """Demonstrate rotation curve solver."""
    
    print("="*70)
    print("3D+3D ROTATION CURVE SOLVER - Demo")
    print("="*70)
    print(f"\nUniversal Parameters (FIXED):")
    print(f"  λ₂ = {Constants.lambda_2} kpc (breathing scale)")
    print(f"  β  = {Constants.beta} (coupling constant)")
    print(f"  M_crit = {Constants.M_crit:.2e} M☉")
    
    # Create example galaxy (NGC 2403-like)
    print("\n" + "-"*70)
    print("Example: NGC 2403-like galaxy")
    print("-"*70)
    
    model = RotationCurve3D3D(
        M_disk=5e9,    # M_sun
        R_disk=2.5,    # kpc
        M_bulge=0,     # No bulge
        M_gas=3e9,     # M_sun
        R_gas=7.5      # kpc
    )
    
    # Compute curve
    R = np.linspace(0.5, 25, 50)
    result = model.compute_curve(R)
    
    print(f"\nResults at R = 10 kpc:")
    idx = np.argmin(np.abs(R - 10))
    print(f"  V_bar   = {result['V_bar'][idx]:.1f} km/s")
    print(f"  V_Q     = {result['V_Q'][idx]:.1f} km/s")
    print(f"  V_total = {result['V_total'][idx]:.1f} km/s")
    
    # Plot
    fig, ax = model.plot(R_max=25)
    plt.savefig('rotation_curve_demo.png', dpi=150, bbox_inches='tight')
    print("\n✓ Saved: rotation_curve_demo.png")
    
    print("\n" + "="*70)
    print("Demo complete!")
    print("="*70)
    
    return model, result


if __name__ == '__main__':
    model, result = demo()
    plt.show()
