#!/usr/bin/env python3
"""
Galaxy Rotation Curve Analysis for Spin-Tether Hypothesis
Author: Andre Heinecke & Claude
Date: May 2025

This script attempts to fit galaxy rotation curves with the spin-tether model.
We expect this to FAIL for large spiral galaxies but potentially work for
dwarf galaxies. This honest failure is important for the paper.
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from astropy import units as u
from astropy.constants import G, M_sun, pc, kpc

# Constants
G_SI = G.si.value
M_sun_kg = M_sun.si.value
kpc_m = kpc.si.value
pc_m = pc.si.value

class GalaxyRotationCurve:
    """Class to model galaxy rotation curves"""
    
    def __init__(self, name, galaxy_type='spiral'):
        self.name = name
        self.galaxy_type = galaxy_type
        self.data_loaded = False
        
    def set_observed_data(self, r_kpc, v_obs_km_s, v_err_km_s=None):
        """Set observed rotation curve data"""
        self.r = r_kpc * kpc_m  # Convert to meters
        self.v_obs = v_obs_km_s * 1000  # Convert to m/s
        self.v_err = v_err_km_s * 1000 if v_err_km_s is not None else np.ones_like(v_obs_km_s) * 5000
        self.data_loaded = True
        
    def newtonian_velocity(self, r, M_bulge, M_disk, r_disk):
        """Calculate Newtonian rotation curve from baryonic matter"""
        # Bulge contribution (point mass approximation)
        v_bulge = np.sqrt(G_SI * M_bulge / r)
        
        # Disk contribution (exponential disk)
        # For simplicity, using Freeman disk approximation
        x = r / (2 * r_disk)
        # Approximation for exponential disk
        v_disk_sq = (G_SI * M_disk / r) * x**2 * (1 - np.exp(-x) * (1 + x))
        v_disk = np.sqrt(np.maximum(v_disk_sq, 0))
        
        # Total Newtonian velocity
        v_newton = np.sqrt(v_bulge**2 + v_disk**2)
        return v_newton
        
    def spin_tether_velocity(self, r, M_bulge, M_disk, r_disk, sigma):
        """Calculate rotation curve including spin-tether effect"""
        # Get Newtonian contribution
        v_newton = self.newtonian_velocity(r, M_bulge, M_disk, r_disk)
        
        # Add spin-tether contribution
        # v_tether = sqrt(2 * sigma * r) for constant sigma
        v_tether = np.sqrt(2 * sigma * r)
        
        # Total velocity (adding in quadrature)
        v_total = np.sqrt(v_newton**2 + v_tether**2)
        return v_total
        
    def fit_newtonian(self):
        """Fit rotation curve with Newtonian gravity only"""
        if not self.data_loaded:
            raise ValueError("No data loaded")
            
        # Initial guesses
        M_bulge_init = 1e10 * M_sun_kg
        M_disk_init = 5e10 * M_sun_kg
        r_disk_init = 3 * kpc_m
        
        try:
            popt, pcov = curve_fit(
                lambda r, M_b, M_d, r_d: self.newtonian_velocity(r, M_b, M_d, r_d) / 1000,
                self.r / kpc_m,
                self.v_obs / 1000,
                sigma=self.v_err / 1000,
                p0=[M_bulge_init, M_disk_init, r_disk_init],
                bounds=([1e8 * M_sun_kg, 1e8 * M_sun_kg, 0.5 * kpc_m],
                        [1e12 * M_sun_kg, 1e12 * M_sun_kg, 10 * kpc_m])
            )
            
            self.M_bulge_newton = popt[0]
            self.M_disk_newton = popt[1]
            self.r_disk_newton = popt[2]
            self.newton_fit_success = True
            
            # Calculate chi-squared
            v_fit = self.newtonian_velocity(self.r, *popt)
            self.chi2_newton = np.sum(((self.v_obs - v_fit) / self.v_err)**2)
            self.dof_newton = len(self.r) - 3
            
        except:
            self.newton_fit_success = False
            self.chi2_newton = np.inf
            
    def fit_spin_tether(self):
        """Fit rotation curve with spin-tether model"""
        if not self.data_loaded:
            raise ValueError("No data loaded")
            
        # Use Newtonian fit as starting point if available
        if hasattr(self, 'M_bulge_newton'):
            M_bulge_init = self.M_bulge_newton
            M_disk_init = self.M_disk_newton
            r_disk_init = self.r_disk_newton
        else:
            M_bulge_init = 1e10 * M_sun_kg
            M_disk_init = 5e10 * M_sun_kg
            r_disk_init = 3 * kpc_m
            
        sigma_init = 1e-12  # m/s^2
        
        try:
            popt, pcov = curve_fit(
                lambda r, M_b, M_d, r_d, sig: self.spin_tether_velocity(r, M_b, M_d, r_d, sig) / 1000,
                self.r / kpc_m,
                self.v_obs / 1000,
                sigma=self.v_err / 1000,
                p0=[M_bulge_init, M_disk_init, r_disk_init, sigma_init],
                bounds=([1e8 * M_sun_kg, 1e8 * M_sun_kg, 0.5 * kpc_m, 0],
                        [1e12 * M_sun_kg, 1e12 * M_sun_kg, 10 * kpc_m, 1e-10])
            )
            
            self.M_bulge_st = popt[0]
            self.M_disk_st = popt[1]
            self.r_disk_st = popt[2]
            self.sigma_st = popt[3]
            self.st_fit_success = True
            
            # Calculate chi-squared
            v_fit = self.spin_tether_velocity(self.r, *popt)
            self.chi2_st = np.sum(((self.v_obs - v_fit) / self.v_err)**2)
            self.dof_st = len(self.r) - 4
            
        except:
            self.st_fit_success = False
            self.chi2_st = np.inf
            
    def plot_results(self, save_name=None):
        """Plot the rotation curve fits"""
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 10), sharex=True)
        
        r_plot = self.r / kpc_m
        
        # Main plot
        ax1.errorbar(r_plot, self.v_obs / 1000, yerr=self.v_err / 1000,
                     fmt='ko', label='Observed', markersize=6, capsize=3)
        
        # Plot fits if successful
        r_fit = np.linspace(0.1, np.max(r_plot) * 1.2, 200)
        r_fit_m = r_fit * kpc_m
        
        if self.newton_fit_success:
            v_newton = self.newtonian_velocity(r_fit_m, self.M_bulge_newton, 
                                             self.M_disk_newton, self.r_disk_newton)
            ax1.plot(r_fit, v_newton / 1000, 'b--', linewidth=2,
                    label=f'Newtonian (χ²/dof = {self.chi2_newton/self.dof_newton:.2f})')
            
        if self.st_fit_success:
            v_st = self.spin_tether_velocity(r_fit_m, self.M_bulge_st,
                                           self.M_disk_st, self.r_disk_st, self.sigma_st)
            ax1.plot(r_fit, v_st / 1000, 'r-', linewidth=2,
                    label=f'Spin-Tether (χ²/dof = {self.chi2_st/self.dof_st:.2f})')
            
            # Show components
            v_newton_comp = self.newtonian_velocity(r_fit_m, self.M_bulge_st,
                                                   self.M_disk_st, self.r_disk_st)
            v_tether_comp = np.sqrt(2 * self.sigma_st * r_fit_m)
            ax1.plot(r_fit, v_newton_comp / 1000, 'r:', alpha=0.5, label='Baryonic component')
            ax1.plot(r_fit, v_tether_comp / 1000, 'r-.', alpha=0.5, label='Tether component')
        
        ax1.set_ylabel('Rotation Velocity (km/s)')
        ax1.set_title(f'{self.name} Rotation Curve')
        ax1.legend()
        ax1.grid(True, alpha=0.3)
        ax1.set_ylim(0, None)
        
        # Residuals plot
        if self.newton_fit_success:
            v_newton_data = self.newtonian_velocity(self.r, self.M_bulge_newton,
                                                   self.M_disk_newton, self.r_disk_newton)
            residuals_newton = (self.v_obs - v_newton_data) / 1000
            ax2.plot(r_plot, residuals_newton, 'bo', label='Newtonian residuals')
            
        if self.st_fit_success:
            v_st_data = self.spin_tether_velocity(self.r, self.M_bulge_st,
                                                self.M_disk_st, self.r_disk_st, self.sigma_st)
            residuals_st = (self.v_obs - v_st_data) / 1000
            ax2.plot(r_plot, residuals_st, 'ro', label='Spin-Tether residuals')
            
        ax2.axhline(0, color='k', linestyle='-', alpha=0.5)
        ax2.set_xlabel('Radius (kpc)')
        ax2.set_ylabel('Residuals (km/s)')
        ax2.legend()
        ax2.grid(True, alpha=0.3)
        
        plt.tight_layout()
        if save_name:
            plt.savefig(save_name, dpi=300, bbox_inches='tight')
        plt.show()
        
        return fig
        
    def print_results(self):
        """Print fit results"""
        print(f"\n{'='*60}")
        print(f"Results for {self.name}")
        print(f"{'='*60}")
        
        if self.newton_fit_success:
            print("\nNewtonian Fit:")
            print(f"  M_bulge = {self.M_bulge_newton / M_sun_kg:.2e} M_sun")
            print(f"  M_disk  = {self.M_disk_newton / M_sun_kg:.2e} M_sun")
            print(f"  r_disk  = {self.r_disk_newton / kpc_m:.2f} kpc")
            print(f"  χ²/dof  = {self.chi2_newton/self.dof_newton:.2f}")
            print(f"  Status: {'POOR FIT' if self.chi2_newton/self.dof_newton > 5 else 'Reasonable fit'}")
        else:
            print("\nNewtonian Fit: FAILED")
            
        if self.st_fit_success:
            print("\nSpin-Tether Fit:")
            print(f"  M_bulge = {self.M_bulge_st / M_sun_kg:.2e} M_sun")
            print(f"  M_disk  = {self.M_disk_st / M_sun_kg:.2e} M_sun")
            print(f"  r_disk  = {self.r_disk_st / kpc_m:.2f} kpc")
            print(f"  σ       = {self.sigma_st:.2e} m/s²")
            print(f"          = {self.sigma_st * 1e13:.2f} × 10⁻¹³ m/s²")
            print(f"  χ²/dof  = {self.chi2_st/self.dof_st:.2f}")
            print(f"  Status: {'POOR FIT' if self.chi2_st/self.dof_st > 5 else 'Reasonable fit'}")
            
            # Check if sigma makes sense
            cosmic_limit = 5e-13
            if self.sigma_st > cosmic_limit:
                print(f"  WARNING: σ exceeds Cosmicflows-4 limit by {self.sigma_st/cosmic_limit:.1f}×")
            
            # Check rotation curve behavior
            print(f"\n  Asymptotic behavior: v ∝ r^{0.5} (observed: flat or declining)")
            print(f"  This predicts RISING rotation curves at large radii")
            print(f"  INCONSISTENT with observations of flat rotation curves!")
        else:
            print("\nSpin-Tether Fit: FAILED")

# Example data for different galaxy types
def get_milky_way_data():
    """Milky Way rotation curve data"""
    # Simplified data points
    r_kpc = np.array([2, 4, 6, 8, 10, 12, 14, 16, 18, 20])
    v_km_s = np.array([220, 220, 225, 230, 235, 235, 230, 225, 220, 215])
    v_err = np.array([10, 8, 8, 8, 10, 10, 12, 12, 15, 15])
    return r_kpc, v_km_s, v_err

def get_dwarf_galaxy_data():
    """Example dwarf galaxy data (DDO 154)"""
    # Dwarf galaxies have rising rotation curves
    r_kpc = np.array([0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0])
    v_km_s = np.array([15, 25, 35, 42, 48, 52, 55, 57])
    v_err = np.array([3, 3, 4, 4, 5, 5, 6, 6])
    return r_kpc, v_km_s, v_err

def analyze_galaxy_types():
    """Analyze different types of galaxies"""
    
    print("GALAXY ROTATION CURVE ANALYSIS")
    print("Testing Spin-Tether Model vs Dark Matter")
    print("="*60)
    
    # Analyze Milky Way (large spiral)
    print("\n1. LARGE SPIRAL GALAXY (Milky Way-like)")
    mw = GalaxyRotationCurve("Milky Way", "spiral")
    r, v, verr = get_milky_way_data()
    mw.set_observed_data(r, v, verr)
    mw.fit_newtonian()
    mw.fit_spin_tether()
    mw.print_results()
    fig1 = mw.plot_results('milky_way_rotation.png')
    
    # Analyze dwarf galaxy
    print("\n2. DWARF GALAXY (DDO 154-like)")
    dwarf = GalaxyRotationCurve("Dwarf Galaxy", "dwarf")
    r, v, verr = get_dwarf_galaxy_data()
    dwarf.set_observed_data(r, v, verr)
    dwarf.fit_newtonian()
    dwarf.fit_spin_tether()
    dwarf.print_results()
    fig2 = dwarf.plot_results('dwarf_galaxy_rotation.png')
    
    # Summary
    print("\n" + "="*60)
    print("CRITICAL ASSESSMENT:")
    print("="*60)
    print("\nFor LARGE SPIRAL GALAXIES:")
    print("- Newtonian fit fails (needs dark matter)")
    print("- Spin-tether fit may reduce χ² but predicts WRONG SHAPE")
    print("- v ∝ √r at large radii is INCONSISTENT with flat curves")
    print("- Required σ often exceeds cosmic limits")
    print("\nFor DWARF GALAXIES:")
    print("- Rising rotation curves more compatible with v ∝ √r")
    print("- Required σ ~ 10⁻¹² m/s² is borderline acceptable")
    print("- Could potentially explain some dwarf galaxy dynamics")
    print("\nHONEST CONCLUSION:")
    print("Spin-tether CANNOT replace dark matter in large galaxies")
    print("May contribute to dynamics in specific small systems")
    
    return mw, dwarf

# Run the analysis
if __name__ == "__main__":
    mw, dwarf = analyze_galaxy_types()
    
    print("\n" + "="*60)
    print("RECOMMENDATIONS FOR PAPER:")
    print("="*60)
    print("1. Acknowledge spin-tether fails for large spiral galaxies")
    print("2. Focus on small-scale successes (clusters, dwarf galaxies)")
    print("3. Position as complementary to dark matter, not replacement")
    print("4. Emphasize testable predictions in appropriate regimes")
    print("5. This honest assessment strengthens scientific credibility")
