"""
Static visualization module for the bouncing ball analysis project.
"""

import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# Set the default style
sns.set_theme(style="whitegrid", palette="muted")

def plot_position_vs_time(df, label, ball_type):
    """
    Plot position vs. time for a single trial.
    
    Parameters:
        df (pandas.DataFrame): DataFrame with 'Time' and 'Position' columns
        label (str): Label for the trial (e.g., "11 inches")
        ball_type (str): Type of ball (e.g., "Golf")
    """
    fig, ax = plt.subplots(figsize=(10, 6))
    
    ax.plot(df["Time"], df["Position"], label=label, marker='o', markersize=3, linewidth=2)
    
    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)
    
    ax.legend(fontsize=12)
    
    plt.tight_layout()
    plt.show()


def plot_trial_with_bounces(df, peak_indices, signal_data, label, ball_type,
                           initial_height=None, cor=None, bounces=7, dt=0.001, g=386.09):
    """
    Plots raw data (Position vs. Time) for a single trial and marks the detected bounces with red 'x'.
    Optionally overlays an ideal bounce trajectory.
    
    Parameters:
        df (pandas.DataFrame): DataFrame with 'Time' and 'Position' columns
        peak_indices (numpy.ndarray): Indices of detected bounces
        signal_data (numpy.ndarray): Position data
        label (str): Label for the trial (e.g., "11 inches")
        ball_type (str): Type of ball (e.g., "Golf")
        initial_height (float): Initial height of the ball
        cor (float): Coefficient of Restitution
        bounces (int): Number of bounces to show in the ideal trajectory
        dt (float): Time step for the ideal trajectory simulation
        g (float): Acceleration due to gravity (in inches/s²)
    """
    plt.figure(figsize=(10, 6))
    
    # 1) Plot raw signal
    plt.plot(df['Time'], df['Position'], linewidth=1.5, color='blue', label='Raw Data')

    # 2) Mark the detected bounces
    if len(peak_indices) > 0:
        plt.plot(
            df['Time'].iloc[peak_indices], 
            df['Position'].iloc[peak_indices], 
            'gx', 
            markersize=10, 
            label='Detected Bounces'
        )
    
    # 3) Simulated Idealized Bounce Trajectory (Ensuring Exactly 7 Bounces)
    if initial_height is not None and cor is not None and not np.isnan(cor):
        time_points = []
        position_points = []
        
        # Initialize simulation parameters
        y = initial_height
        v_y = 0.0  # Starting from rest
        t = 0.0
        bounce_count = 0
        model_first_bounce_time = None  

        while bounce_count < bounces:  # Ensure exactly 7 bounces
            time_points.append(t)
            position_points.append(y)

            # Kinematic update
            v_y -= g * dt  
            y += v_y * dt  
            t += dt  

            # Bounce condition
            if y <= 0:
                bounce_count += 1
                y = 0  # Ground impact
                v_y = -v_y * cor  # Apply COR

                if model_first_bounce_time is None:
                    model_first_bounce_time = t  # Store first bounce time

        # Align first model bounce with data's first detected bounce
        if model_first_bounce_time is not None and len(peak_indices) > 0:
            data_first_bounce_time = df['Time'].iloc[peak_indices[0]]
            time_shift = data_first_bounce_time - model_first_bounce_time
            time_points = [t + time_shift for t in time_points]

        # Plot idealized trajectory
        plt.plot(time_points, position_points, 'r--', label='Ideal Trajectory')

    # 4) Formatting the plot
    plt.title(f'{ball_type} Trial: {label}')
    plt.xlabel('Time (s)')
    plt.ylabel('Position (inches)')
    plt.autoscale(True, 'x', True)
    plt.grid(True)
    plt.legend()
    plt.show()


def plot_cor_table(cor_df, ball_type):
    """
    Plots a line (and markers) of Average COR vs. Initial Height from the summary DataFrame.
    
    Parameters:
        cor_df (pandas.DataFrame): DataFrame with 'Initial Height' and 'Average COR' columns
        ball_type (str): Type of ball (e.g., "Golf")
    """
    plt.figure(figsize=(8, 5))
    sns.lineplot(x='Initial Height', y='Average COR', marker='o', data=cor_df)
    plt.title(f'Average COR vs. Initial Height for {ball_type} Ball')
    plt.xlabel('Initial Height (inches)')  
    plt.ylabel('Average COR')
    plt.grid(True)
    plt.tight_layout()
    plt.show()


def plot_all_cors(golf_df, lacrosse_df, metal_df):
    """
    Plot COR vs. initial height for all ball types on the same graph.
    
    Parameters:
        golf_df (pandas.DataFrame): DataFrame with golf ball data
        lacrosse_df (pandas.DataFrame): DataFrame with lacrosse ball data
        metal_df (pandas.DataFrame): DataFrame with metal ball data
    """
    plt.figure(figsize=(10, 6))
    
    # Plot each ball type with a different color and marker
    sns.lineplot(x='Initial Height', y='Average COR', marker='o', data=golf_df, label='Golf Ball')
    sns.lineplot(x='Initial Height', y='Average COR', marker='s', data=lacrosse_df, label='Lacrosse Ball')
    sns.lineplot(x='Initial Height', y='Average COR', marker='^', data=metal_df, label='Metal Ball')
    
    plt.title('Coefficient of Restitution vs. Initial Height')
    plt.xlabel('Initial Height (inches)')
    plt.ylabel('Average COR')
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.show() 