import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import seaborn as sns
from matplotlib.animation import PillowWriter

sns.set_theme(style="whitegrid", palette="muted")

# File paths for golf ball data
golf_file_paths = {
    "11 inches": "golf_11.csv",
    "12 inches": "golf_12.csv",
    "13 inches": "golf_13.csv",
    "14 inches": "golf_14.csv",
    "15 inches": "golf_15.csv"
}

# File paths for lacrosse ball data
lacrosse_file_paths = {
    "18 inches": "l_18.csv",
    "19 inches": "l_19.csv",
    "20 inches": "l_20.csv",
    "21 inches": "l_21.csv",
    "22 inches": "l_22.csv"
}

# File paths for metal ball data
metal_file_paths = {
    "2 inches": "metal_2.csv",
    "4 inches": "metal_4.csv",
    "6 inches": "metal_6.csv",
    "8 inches": "metal_8.csv",
    "10 inches": "metal_10.csv"
}

def animate_position_vs_time(file_paths, ball_type, save_animation=False):
    """
    Creates an animated plot showing the position vs time data for each trial.
    The animation shows the data being drawn progressively over time.
    """
    for label, path in file_paths.items():
        # Read CSV with no header, assuming columns: Time, Position
        df = pd.read_csv(path, header=None, names=["Time", "Position"])
        
        # Convert to numpy arrays for better performance and to avoid indexing issues
        time_data = df["Time"].to_numpy()
        position_data = df["Position"].to_numpy()
        
        # Create figure and axis
        fig, ax = plt.subplots(figsize=(10, 6))
        
        # Set up the plot
        ax.set_xlabel("Time (seconds)", fontsize=14)
        ax.set_ylabel("Position (inches)", fontsize=14)
        ax.set_title(f"{ball_type} - Position vs. Time\nInitial Height: {label}", fontsize=16, pad=15)
        ax.grid(True, which='both', linestyle='--', linewidth=0.5, alpha=0.7)
        
        # Set axis limits
        ax.set_xlim(time_data.min(), time_data.max())
        ax.set_ylim(0, position_data.max() * 1.1)
        
        # Create a line object with no data
        line, = ax.plot([], [], 'o-', markersize=3, linewidth=2, label=label)
        
        # Create a point that will move along the line
        point, = ax.plot([], [], 'ro', markersize=8)
        
        # Add legend
        ax.legend(fontsize=12)
        
        # Animation initialization function
        def init():
            line.set_data([], [])
            point.set_data([], [])
            return line, point
        
        # Animation update function
        def update(frame):
            # For smooth animation, use a percentage of the data
            idx = int(len(time_data) * frame / 100) if frame < 100 else len(time_data) - 1
            
            # Update line data
            line.set_data(time_data[:idx+1], position_data[:idx+1])
            
            # Update point position - ensure we're passing arrays/lists, not single values
            if idx > 0:
                point.set_data([time_data[idx]], [position_data[idx]])
            
            return line, point
        
        # Create animation
        ani = animation.FuncAnimation(
            fig, update, frames=120, init_func=init, 
            blit=True, interval=50, repeat=True
        )
        
        # Save animation if requested
        if save_animation:
            writer = PillowWriter(fps=30)
            filename = f"{ball_type.lower().replace(' ', '_')}_{label.replace(' ', '_')}.gif"
            ani.save(filename, writer=writer)
            print(f"Animation saved as {filename}")
        
        plt.tight_layout()
        plt.show()

# Function to animate all ball types
def animate_all_ball_types(save_animations=False):
    print("Animating Golf Ball data...")
    animate_position_vs_time(golf_file_paths, "Golf Ball", save_animations)
    
    print("Animating Lacrosse Ball data...")
    animate_position_vs_time(lacrosse_file_paths, "Lacrosse Ball", save_animations)
    
    print("Animating Metal Ball data...")
    animate_position_vs_time(metal_file_paths, "Metal Ball", save_animations)

if __name__ == "__main__":
    # Set to True if you want to save the animations as GIF files
    save_animations = False
    animate_all_ball_types(save_animations) 