import pandas as pd
import numpy as np
from scipy.ndimage import gaussian_filter
import pyvista as pv

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

# Convert parallax to distance (light years)
df['distance_ly'] = 3262 / df['parallax']

# Remove rows with NaN values
df = df.dropna(subset=['ra', 'dec', 'distance_ly', 'parallax'])
df = df[np.isfinite(df['distance_ly'])]
df = df.reset_index(drop=True)

# Convert RA/Dec to Cartesian coordinates
ra_rad = np.radians(df['ra'])
dec_rad = np.radians(df['dec'])
x = df['distance_ly'] * np.cos(dec_rad) * np.cos(ra_rad)
y = df['distance_ly'] * np.cos(dec_rad) * np.sin(ra_rad)
z = df['distance_ly'] * np.sin(dec_rad)

# Create 3D grid for density calculation
grid_size = 60  # REDUCED from 80 for faster rendering
x_range = np.linspace(-130, 130, grid_size)
y_range = np.linspace(-130, 130, grid_size)
z_range = np.linspace(-130, 130, grid_size)

X, Y, Z = np.meshgrid(x_range, y_range, z_range, indexing='ij')

# Calculate density using histogram in 3D
H, edges = np.histogramdd(
    np.column_stack([x, y, z]),
    bins=(x_range, y_range, z_range)
)

# Smooth the density field
H_smooth = gaussian_filter(H, sigma=2)  # REDUCED sigma from 3

# Print diagnostics
print(f"Density stats: min={H_smooth.min():.4f}, max={H_smooth.max():.4f}, mean={H_smooth.mean():.4f}")
print(f"Non-zero cells: {np.sum(H_smooth > 0)}/{H_smooth.size}")

# Create PyVista uniform grid
grid = pv.ImageData()
grid.dimensions = H_smooth.shape
grid.spacing = (
    (x_range.max() - x_range.min()) / (grid_size - 1),
    (y_range.max() - y_range.min()) / (grid_size - 1),
    (z_range.max() - z_range.min()) / (grid_size - 1)
)
grid.origin = (x_range.min(), y_range.min(), z_range.min())

# Add density as POINT data for contouring
grid.point_data['density'] = H_smooth.flatten(order='F')

# ===== BUILD SCENE FUNCTION =====
def build_scene(plotter):
    """Add all meshes and settings to plotter"""
    
    plotter.set_background('black')
    
    # Define density thresholds - REDUCED to 3 surfaces for speed
    thresholds = [
        (0.05, 'cyan', 0.4, 'Sparse'),
        (0.15, 'yellow', 0.6, 'Medium'),
        (0.30, 'red', 0.8, 'Dense')
    ]
    
    print(f"Thresholds: {[t[0] for t in thresholds]}")
    
    # Add isosurfaces for each density level
    surfaces_added = 0
    for threshold, color, opacity, label in thresholds:
        try:
            contour = grid.contour([threshold], scalars='density')
            if contour.n_points > 0:
                plotter.add_mesh(
                    contour,
                    color=color,
                    opacity=opacity,
                    show_edges=False,
                    smooth_shading=True,
                    label=label
                )
                surfaces_added += 1
                print(f"Added surface at threshold {threshold:.3f}: {contour.n_points} points")
            else:
                print(f"No surface at threshold {threshold:.3f}")
        except Exception as e:
            print(f"Failed at threshold {threshold:.3f}: {e}")
            continue
    
    print(f"Total surfaces added: {surfaces_added}")
    
    # Add title
    plotter.add_text(
        '3D Stellar Density Isosurfaces\n24,071 Stars within 131 Light Years',
        position='upper_edge',
        font_size=12,
        color='white'
    )
    
    plotter.add_text(
        'Thin shells near Earth = Stellar Desert\nThick shells at distance = Dense regions',
        position='lower_left',
        font_size=9,
        color='white'
    )
    
    # Add axes
    plotter.show_axes()
    
    # Add Earth marker
    origin_sphere = pv.Sphere(radius=5, center=(0, 0, 0))
    plotter.add_mesh(origin_sphere, color='white', opacity=0.8, label='Earth')
    
    # Set camera
    plotter.camera_position = 'iso'
    plotter.camera.zoom(1.2)

# ===== LIVE INTERACTIVE WINDOW =====
print("Opening live interactive window...")
print("Close the window when ready.")

plotter_live = pv.Plotter()
build_scene(plotter_live)
plotter_live.show()

print("Live window closed.")

# Set camera
plotter.camera_position = 'iso'
plotter.camera.zoom(1.2)

"""
print("Generating video...")
plotter_video = pv.Plotter(off_screen=True, window_size=[1920, 1080])
build_scene(plotter_video)

plotter_video.open_movie('isosurface_rotation.mp4', framerate=30)
plotter_video.render()

n_frames = 360
path = plotter_video.generate_orbital_path(n_points=n_frames, shift=0, viewup=[0, 0, 1], factor=2.5)

for i in range(n_frames):
    plotter_video.camera.position = path.points[i]
    plotter_video.write_frame()
    
plotter_video.close()
print("Video saved as isosurface_rotation.mp4")
"""

print(f"\nTotal stars: {len(df)}")
print(f"Grid resolution: {grid_size}³ = {grid_size**3} cells")
print(f"Density range: {H_smooth.min():.2f} to {H_smooth.max():.2f}")