"""
Animation module for visualizing bouncing ball data.
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.patches import Circle
import pandas as pd
from IPython.display import HTML


def create_ball_animation(df, ball_radius=0.5, fps=30, duration=None, title="Bouncing Ball Animation"):
    """
    Create an animation of a bouncing ball from position data.
    
    Parameters:
        df (pandas.DataFrame): DataFrame with 'Time' and 'Position' columns
        ball_radius (float): Radius of the ball in inches
        fps (int): Frames per second for the animation
        duration (float): Duration of the animation in seconds (if None, uses the full data)
        title (str): Title for the animation
        
    Returns:
        matplotlib.animation.Animation: The animation object
    """
    # Extract time and position data
    time_data = df['Time'].values
    position_data = df['Position'].values
    
    # Determine animation duration
    if duration is None:
        duration = time_data[-1]
    
    # Calculate the number of frames based on fps and duration
    num_frames = int(fps * duration)
    
    # Create figure and axis
    fig, ax = plt.subplots(figsize=(10, 6))
    
    # Set axis limits
    max_height = np.max(position_data) + 2 * ball_radius
    ax.set_xlim(-5, 5)
    ax.set_ylim(0, max_height)
    
    # Set labels and title
    ax.set_xlabel('X Position (inches)')
    ax.set_ylabel('Y Position (inches)')
    ax.set_title(title)
    
    # Add a horizontal line at y=0 to represent the ground
    ax.axhline(y=0, color='black', linestyle='-', linewidth=2)
    
    # Create the ball (circle patch)
    ball = Circle((0, position_data[0]), ball_radius, color='red', zorder=2)
    ax.add_patch(ball)
    
    # Create a line to show the ball's path
    line, = ax.plot([], [], 'b-', alpha=0.3, linewidth=1)
    
    # Initialize empty lists for the path
    path_x = []
    path_y = []
    
    def init():
        """Initialize the animation"""
        ball.center = (0, position_data[0])
        line.set_data([], [])
        return ball, line
    
    def animate(i):
        """Update the animation for frame i"""
        # Calculate the current time based on the frame number
        current_time = i * duration / num_frames
        
        # Find the closest time index in our data
        idx = np.abs(time_data - current_time).argmin()
        
        # Update ball position
        y_pos = position_data[idx]
        ball.center = (0, y_pos)
        
        # Update path
        path_x.append(0)
        path_y.append(y_pos)
        line.set_data(path_x, path_y)
        
        return ball, line
    
    # Create the animation
    anim = animation.FuncAnimation(
        fig, animate, init_func=init, frames=num_frames,
        interval=1000/fps, blit=True
    )
    
    plt.close()  # Prevent the static plot from displaying
    
    return anim


def create_ball_animation_with_model(df, peak_indices, cor, initial_height, 
                                    ball_radius=0.5, fps=30, duration=None, 
                                    title="Bouncing Ball Animation with Model", 
                                    g=386.09):
    """
    Create an animation of a bouncing ball from position data with an ideal model comparison.
    
    Parameters:
        df (pandas.DataFrame): DataFrame with 'Time' and 'Position' columns
        peak_indices (numpy.ndarray): Indices of detected bounces
        cor (float): Coefficient of Restitution
        initial_height (float): Initial height of the ball in inches
        ball_radius (float): Radius of the ball in inches
        fps (int): Frames per second for the animation
        duration (float): Duration of the animation in seconds (if None, uses the full data)
        title (str): Title for the animation
        g (float): Acceleration due to gravity (in inches/s²)
        
    Returns:
        matplotlib.animation.Animation: The animation object
    """
    # Extract time and position data
    time_data = df['Time'].values
    position_data = df['Position'].values
    
    # Determine animation duration
    if duration is None:
        duration = time_data[-1]
    
    # Calculate the number of frames based on fps and duration
    num_frames = int(fps * duration)
    
    # Create figure and axis
    fig, ax = plt.subplots(figsize=(10, 6))
    
    # Set axis limits
    max_height = np.max(position_data) + 2 * ball_radius
    ax.set_xlim(-5, 5)
    ax.set_ylim(0, max_height)
    
    # Set labels and title
    ax.set_xlabel('X Position (inches)')
    ax.set_ylabel('Y Position (inches)')
    ax.set_title(title)
    
    # Add a horizontal line at y=0 to represent the ground
    ax.axhline(y=0, color='black', linestyle='-', linewidth=2)
    
    # Create the real data ball (circle patch)
    real_ball = Circle((0, position_data[0]), ball_radius, color='red', zorder=2, label='Actual')
    ax.add_patch(real_ball)
    
    # Create the model ball (circle patch)
    model_ball = Circle((2, initial_height), ball_radius, color='blue', zorder=2, alpha=0.7, label='Model')
    ax.add_patch(model_ball)
    
    # Create lines to show the balls' paths
    real_line, = ax.plot([], [], 'r-', alpha=0.3, linewidth=1)
    model_line, = ax.plot([], [], 'b-', alpha=0.3, linewidth=1)
    
    # Add legend
    ax.legend(handles=[
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='red', markersize=10, label='Actual'),
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='blue', markersize=10, label='Model')
    ])
    
    # Initialize empty lists for the paths
    real_path_x = []
    real_path_y = []
    model_path_x = []
    model_path_y = []
    
    # Generate model data
    dt = 0.001  # Time step for simulation
    model_time = []
    model_position = []
    
    # Initialize simulation parameters
    y = initial_height
    v_y = 0.0  # Starting from rest
    t = 0.0
    
    # Run simulation for the duration
    while t <= duration:
        model_time.append(t)
        model_position.append(y)
        
        # Kinematic update
        v_y -= g * dt  
        y += v_y * dt  
        t += dt  
        
        # Bounce condition
        if y <= 0:
            y = 0  # Ground impact
            v_y = -v_y * cor  # Apply COR
    
    # Convert to numpy arrays for faster indexing
    model_time = np.array(model_time)
    model_position = np.array(model_position)
    
    def init():
        """Initialize the animation"""
        real_ball.center = (0, position_data[0])
        model_ball.center = (2, initial_height)
        real_line.set_data([], [])
        model_line.set_data([], [])
        return real_ball, model_ball, real_line, model_line
    
    def animate(i):
        """Update the animation for frame i"""
        # Calculate the current time based on the frame number
        current_time = i * duration / num_frames
        
        # Find the closest time index in our data
        real_idx = np.abs(time_data - current_time).argmin()
        model_idx = np.abs(model_time - current_time).argmin()
        
        # Update ball positions
        real_y_pos = position_data[real_idx]
        model_y_pos = model_position[model_idx]
        
        real_ball.center = (0, real_y_pos)
        model_ball.center = (2, model_y_pos)
        
        # Update paths
        real_path_x.append(0)
        real_path_y.append(real_y_pos)
        model_path_x.append(2)
        model_path_y.append(model_y_pos)
        
        real_line.set_data(real_path_x, real_path_y)
        model_line.set_data(model_path_x, model_path_y)
        
        return real_ball, model_ball, real_line, model_line
    
    # Create the animation
    anim = animation.FuncAnimation(
        fig, animate, init_func=init, frames=num_frames,
        interval=1000/fps, blit=True
    )
    
    plt.close()  # Prevent the static plot from displaying
    
    return anim


def display_animation(anim, html_video=True):
    """
    Display an animation in a Jupyter notebook.
    
    Parameters:
        anim (matplotlib.animation.Animation): The animation to display
        html_video (bool): If True, convert to HTML5 video, otherwise use JavaScript
        
    Returns:
        IPython.display.HTML: The HTML object containing the animation
    """
    if html_video:
        # Convert to HTML5 video
        html = anim.to_html5_video()
        return HTML(html)
    else:
        # Use JavaScript animation
        return HTML(anim.to_jshtml())


def save_animation(anim, filename, fps=30, dpi=100):
    """
    Save an animation to a file.
    
    Parameters:
        anim (matplotlib.animation.Animation): The animation to save
        filename (str): The filename to save to (should end with .mp4, .gif, etc.)
        fps (int): Frames per second
        dpi (int): Resolution in dots per inch
    """
    # Determine the writer based on the file extension
    if filename.endswith('.mp4'):
        writer = animation.FFMpegWriter(fps=fps)
    elif filename.endswith('.gif'):
        writer = animation.PillowWriter(fps=fps)
    else:
        raise ValueError("Unsupported file format. Use .mp4 or .gif")
    
    # Save the animation
    anim.save(filename, writer=writer, dpi=dpi)