import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# Typical stellar radii in solar radii for main sequence stars
stellar_data = {
    'O': {'radius': 15, 'temp': 35000, 'color': '#9bb0ff', 'mass': '16-50'},
    'B': {'radius': 7, 'temp': 20000, 'color': '#aabfff', 'mass': '2.1-16'},
    'A': {'radius': 2.5, 'temp': 9000, 'color': '#cad7ff', 'mass': '1.4-2.1'},
    'F': {'radius': 1.5, 'temp': 6750, 'color': '#f8f7ff', 'mass': '1.04-1.4'},
    'G': {'radius': 1.0, 'temp': 5600, 'color': '#fff4ea', 'mass': '0.8-1.04'},
    'K': {'radius': 0.8, 'temp': 4500, 'color': '#ffd2a1', 'mass': '0.45-0.8'},
    'M': {'radius': 0.3, 'temp': 3000, 'color': '#ff4422', 'mass': '0.08-0.45'}
}

# Create figure with black background
fig, ax = plt.subplots(figsize=(16, 10), facecolor='black')
ax.set_facecolor('black')

# Position stars horizontally with more spacing
classes = list(stellar_data.keys())
# Manually set positions to give O-class more room on left
x_positions = [.05, 7, 10.5, 11.85, 12.9, 14, 15]
# Manual label positions for each class
label_x_positions = [1, 7, 10.5, 11.85, 12.9, 14, 15]
label_y_offsets = [1, 3.5, 1.5, 1.5, 1.5, 1.5, 1.5]  # Adjust these values
y_position = 5

# Plot each star as a bubble
for i, (cls, x) in enumerate(zip(classes, x_positions)):
    data = stellar_data[cls]
    radius = data['radius']
    
    # Scale bubble size (area proportional to radius²)
    size = (radius ** 2) * 1200
    
    # Draw star with glow effect
    ax.scatter(x, y_position, s=size, c=data['color'], 
               alpha=0.9, edgecolors='white', linewidth=1.5)
    
    # Add outer glow
    ax.scatter(x, y_position, s=size * 1.3, c=data['color'], 
               alpha=0.3, edgecolors='none')
    
    # Class label above star
    label_x = label_x_positions[i]
    label_y = y_position + label_y_offsets[i]
    label_text = f'[Sun-like]\n{cls}-class' if cls == 'G' else f'{cls}-class'
    ax.text(label_x, label_y, label_text, 
            ha='center', fontsize=14, fontweight='bold', color='white')
    
    # Radius below star
    ax.text(x, y_position - 3.5, f'{radius} R☉', 
            ha='center', fontsize=12, color='yellow')
    
    # Temperature
    ax.text(x, y_position - 4.2, f'{data["temp"]:,}K', 
            ha='center', fontsize=10, color='lightblue')
    
    # Mass range
    ax.text(x, y_position - 5, f'{data["mass"]} M☉', 
            ha='center', fontsize=9, color='lightgreen')

# Formatting
ax.set_xlim(0, 18)
ax.set_ylim(-1, 10)
ax.set_xticks([])
ax.set_yticks([])
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['bottom'].set_visible(False)
ax.spines['left'].set_visible(False)

# Title
plt.title('Stellar Size Comparison\nMain Sequence Stars by Spectral Class', 
          fontsize=20, fontweight='bold', color='white', pad=20)

# Add legend
legend_text = 'R☉ = Solar Radius\nM☉ = Solar Mass\nK = Kelvin (Temperature)'
ax.text(0.02, 0.98, legend_text, transform=ax.transAxes, 
        fontsize=10, verticalalignment='top', color='white',
        bbox=dict(boxstyle='round', facecolor='black', 
                  edgecolor='white', alpha=0.7))

# Add note about scale
note = 'Note: Bubble sizes scaled to show relative stellar radii\nLargest (O-class) is 50× larger than smallest (M-class)'
ax.text(0.98, 0.02, note, transform=ax.transAxes, 
        fontsize=9, verticalalignment='bottom', ha='right',
        color='white', style='italic')

plt.tight_layout()
plt.savefig('stellar_size_comparison.png', dpi=300, bbox_inches='tight', 
            facecolor='black')
plt.show()

print("\nStellar Size Comparison:")
print("="*50)
for cls, data in stellar_data.items():
    print(f"{cls}-class: {data['radius']} solar radii, "
          f"{data['temp']:,}K, {data['mass']} solar masses")
print("="*50)