import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial import Voronoi, voronoi_plot_2d

# Load data
df = pd.read_csv('24000stars.csv')

# Remove rows with missing critical data
df = df.dropna(subset=['ra', 'dec', 'parallax'])

# Filter out invalid parallax (must be positive for distance calculation)
df = df[df['parallax'] > 0]

# Convert to distance in light-years
df['distance_ly'] = 3262 / df['parallax']  # parallax in mas to light-years

# Convert RA/Dec/Distance to Cartesian coordinates (X, Y, Z)
# RA and Dec are in degrees, convert to radians
ra_rad = np.radians(df['ra'])
dec_rad = np.radians(df['dec'])
distance = df['distance_ly']

# Cartesian coordinates
df['x'] = distance * np.cos(dec_rad) * np.cos(ra_rad)
df['y'] = distance * np.cos(dec_rad) * np.sin(ra_rad)
df['z'] = distance * np.sin(dec_rad)

# For visualization, we'll project onto XY plane (galactic view from "above")
points = df[['x', 'y']].values

# Sample subset for clearer visualization (too many points make it cluttered)
# Try 500 stars first
np.random.seed(42)
sample_size = 500
sample_indices = np.random.choice(len(points), size=sample_size, replace=False)
sample_points = points[sample_indices]
sample_distances = df.iloc[sample_indices]['distance_ly'].values

# Create Voronoi diagram
vor = Voronoi(sample_points)

# Create figure
fig, ax = plt.subplots(figsize=(14, 14))

# Plot Voronoi diagram with color-coded cells
voronoi_plot_2d(vor, ax=ax, show_vertices=False, line_colors='black', 
                line_width=0.5, line_alpha=0.6, point_size=0)

# Color the regions by distance
from matplotlib.collections import PolyCollection
from matplotlib.colors import Normalize
from matplotlib.cm import plasma

norm = Normalize(vmin=sample_distances.min(), vmax=sample_distances.max())
cmap = plasma

# Plot each finite region with color based on distance
for point_idx, region_idx in enumerate(vor.point_region):
    region = vor.regions[region_idx]
    if not -1 in region and len(region) > 0:
        polygon = [vor.vertices[i] for i in region]
        color = cmap(norm(sample_distances[point_idx]))
        ax.fill(*zip(*polygon), color=color, alpha=0.7, edgecolor='black', linewidth=0.5)

# Plot star positions
scatter = ax.scatter(sample_points[:, 0], sample_points[:, 1], 
                    c=sample_distances, cmap='plasma', s=10, 
                    edgecolors='white', linewidth=0.5, zorder=5)

# Add Earth at origin
ax.scatter(0, 0, c='cyan', s=200, marker='*', edgecolors='white', 
          linewidth=2, zorder=10, label='Earth')

# Formatting
ax.set_xlabel('X (light-years)', fontsize=12, fontweight='bold')
ax.set_ylabel('Y (light-years)', fontsize=12, fontweight='bold')
ax.set_title('Voronoi Diagram: Stellar Territories by Distance\n500 Stars Projected onto XY Plane', 
             fontsize=16, fontweight='bold', pad=20)
ax.set_aspect('equal')
ax.grid(True, alpha=0.3, linestyle='--')
ax.legend(fontsize=10)

# Add colorbar
cbar = plt.colorbar(scatter, ax=ax, pad=0.02, fraction=0.046)
cbar.set_label('Distance from Earth (light-years)', fontsize=11, fontweight='bold')

plt.tight_layout()
plt.savefig('53_voronoi_2d_distance.png', dpi=300, bbox_inches='tight')
plt.show()

print(f"Voronoi diagram created with {sample_size} stars")
print(f"Distance range: {sample_distances.min():.1f} - {sample_distances.max():.1f} light-years")
print(f"Each cell shows the 'territory' dominated by that star")