import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import seaborn as sns
from scipy.signal import find_peaks
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"
}

# General detection/plotting parameters
min_distance = 5
min_time_diff = 0.05
fs = 50
save_animations = False  # Set to True to save animations as GIFs

# Prominence adjustments
golf_relative_prominence = 0.15
golf_low_threshold = 12.5
golf_low_adjustment = 1.1
golf_high_threshold = 14.5
golf_high_adjustment = 1.3

lacrosse_relative_prominence = 0.05
lacrosse_low_threshold = 18.5
lacrosse_low_adjustment = 1.1
lacrosse_high_threshold = 21.5
lacrosse_high_adjustment = 1.3

metal_relative_prominence = 0.05
metal_low_threshold = 5
metal_low_adjustment = 0.8
metal_high_threshold = 8
metal_high_adjustment = 1.2

def read_trial(path):
    """
    Reads a CSV file containing time and position data (with no header)
    and returns a DataFrame with columns 'Time' and 'Position'.
    """
    return pd.read_csv(path, header=None, names=['Time', 'Position'])


def detect_bounces(df, 
                   prominence=0, 
                   min_distance=10, 
                   min_time_diff=0.2,
                   max_bounces=7, 
                   fs=None):
    """
    Detects bounce events by locating peaks in the signal.
    
    Parameters:
      - df: DataFrame with 'Time' and 'Position'
      - prominence: required prominence for peaks (in the signal)
      - min_distance: minimum number of data points between peaks
      - min_time_diff: minimum time difference (seconds) between bounces
      - max_bounces: maximum bounces to report
      - fs: override sampling frequency (Hz). If None, computed from 'Time'.
    
    Returns:
      - peak_indices (array of indices),
      - bounce_heights (array of position values at these bounces),
      - bounce_times (array of times of these bounces),
      - signal_data (the data used for detection, here the raw Position values)
    """
    if fs is None:
        dt = np.median(np.diff(df['Time']))
        if dt <= 0:
            raise ValueError(
                "Time differences are non-positive. "
                "Ensure your 'Time' column is in seconds."
            )
        fs = 1 / dt
    
    signal_data = df['Position'].values
    
    peaks, _ = find_peaks(signal_data, prominence=prominence, distance=min_distance)
    
    filtered_peaks = []
    filtered_times = []
    group_start_time = None
    group_best_idx = None
    group_best_value = None
    
    for idx in peaks:
        current_time = df['Time'].iloc[idx]
        current_value = signal_data[idx]
        
        if group_start_time is None:
            group_start_time = current_time
            group_best_idx = idx
            group_best_value = current_value
            continue
        
        if (current_time - group_start_time) < min_time_diff:
            if current_value > group_best_value:
                group_best_idx = idx
                group_best_value = current_value
        else:
            filtered_peaks.append(group_best_idx)
            filtered_times.append(df['Time'].iloc[group_best_idx])
            
            group_start_time = current_time
            group_best_idx = idx
            group_best_value = current_value
    
    if group_best_idx is not None:
        filtered_peaks.append(group_best_idx)
        filtered_times.append(df['Time'].iloc[group_best_idx])
    
    filtered_peaks = filtered_peaks[:max_bounces]
    filtered_times = filtered_times[:max_bounces]
    
    bounce_heights = df['Position'].iloc[filtered_peaks].values
    bounce_times = np.array(filtered_times)

    return np.array(filtered_peaks), bounce_heights, bounce_times, signal_data


def compute_cor(bounce_heights):
    """
    Computes the Coefficient of Restitution (COR) for consecutive bounces:
      e = sqrt( h_{n+1} / h_n )
    
    Returns an array of COR values.
    """
    cor_values = []
    for i in range(len(bounce_heights) - 1):
        if bounce_heights[i] > 0:
            cor_values.append(np.sqrt(bounce_heights[i+1] / bounce_heights[i]))
        else:
            cor_values.append(np.nan)
    return np.array(cor_values)


def process_trial(path, 
                  prominence, 
                  min_distance=10, 
                  min_time_diff=0.2,
                  max_bounces=7, 
                  fs=None):
    """
    Reads trial data, detects bounces, computes COR values, and returns:
      - df (the DataFrame)
      - peak_indices
      - bounce_heights
      - bounce_times
      - cor_values
      - avg_cor
      - signal_data
    """
    df = read_trial(path)
    peak_indices, bounce_heights, bounce_times, signal_data = detect_bounces(
        df, prominence, min_distance, min_time_diff, max_bounces, fs
    )
    cor_values = compute_cor(bounce_heights)
    avg_cor = np.mean(cor_values) if cor_values.size > 0 else np.nan
    
    return df, peak_indices, bounce_heights, bounce_times, cor_values, avg_cor, signal_data


def animate_bouncing_ball(df, peak_indices, bounce_heights, bounce_times, cor, 
                          label, ball_type, initial_height=None, g=386.09, 
                          save_animation=False):
    """
    Creates an animated visualization of a bouncing ball using the detected bounce data.
    
    Parameters:
        df: DataFrame with Time and Position data
        peak_indices: Indices of detected bounces
        bounce_heights: Heights of detected bounces
        bounce_times: Times of detected bounces
        cor: Coefficient of Restitution
        label: Trial label (e.g., "11 inches")
        ball_type: Type of ball (e.g., "Golf")
        initial_height: Initial drop height
        g: Acceleration due to gravity (in inches/s²)
        save_animation: Whether to save the animation as a GIF
    """
    # Convert DataFrame to numpy arrays for better performance
    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 (s)", fontsize=14)
    ax.set_ylabel("Position (inches)", fontsize=14)
    ax.set_title(f"{ball_type} Ball Bounce - Initial Height: {label}\nAvg. COR: {cor:.3f}", 
                 fontsize=16, pad=15)
    
    # Set axis limits
    time_range = time_data.max() - time_data.min()
    ax.set_xlim(time_data.min() - 0.05 * time_range, time_data.max() + 0.05 * time_range)
    ax.set_ylim(-0.5, max(position_data.max(), initial_height if initial_height else 0) * 1.1)
    
    # Plot the raw data
    ax.plot(time_data, position_data, 'b-', alpha=0.3, linewidth=1, label='Raw Data')
    
    # Mark the detected bounces
    if len(peak_indices) > 0:
        bounce_times_array = df['Time'].iloc[peak_indices].to_numpy()
        bounce_heights_array = df['Position'].iloc[peak_indices].to_numpy()
        ax.plot(
            bounce_times_array, 
            bounce_heights_array, 
            'gx', 
            markersize=10, 
            label='Detected Bounces'
        )
    
    # Create a ball object (circle)
    ball_radius = initial_height/20 if initial_height else 0.5
    ball = plt.Circle((time_data[0], position_data[0]), radius=ball_radius, 
                     color='red', fill=True, alpha=0.7)
    ax.add_patch(ball)
    
    # Create a ground line
    ax.axhline(y=0, color='k', linestyle='-', linewidth=2)
    
    # Add legend
    ax.legend(fontsize=12, loc='upper right')
    
    # Add grid
    ax.grid(True, which='both', linestyle='--', linewidth=0.5, alpha=0.7)
    
    # Generate simulated trajectory data
    sim_line = None
    sim_time = None
    sim_pos = None
    
    if initial_height is not None and cor is not None and not np.isnan(cor):
        # Create time points for simulation
        sim_dt = 0.001  # Simulation time step
        sim_time = np.arange(time_data.min(), time_data.max(), sim_dt)
        sim_pos = np.zeros_like(sim_time)
        
        # Initialize simulation parameters
        y = initial_height
        v_y = 0.0  # Starting from rest
        bounce_count = 0
        
        # Align first bounce with data
        time_offset = 0
        if len(bounce_times) > 0:
            # Calculate time to first bounce
            t_to_first_bounce = np.sqrt(2 * initial_height / g)
            time_offset = bounce_times[0] - t_to_first_bounce
        
        # Simulate the bouncing motion
        for i in range(len(sim_time)):
            t = sim_time[i] - time_offset
            
            # Reset if before simulation start
            if t < 0:
                sim_pos[i] = initial_height
                continue
                
            # Apply physics
            y += v_y * sim_dt
            v_y -= g * sim_dt
            
            # Check for bounce
            if y <= 0:
                y = 0
                v_y = -v_y * cor
                bounce_count += 1
            
            sim_pos[i] = y
        
        # Plot the simulated trajectory
        sim_line, = ax.plot([], [], 'r--', linewidth=1.5, label='Simulated Trajectory')
    
    # Animation initialization function
    def init():
        ball.center = (time_data[0], position_data[0])
        if sim_line is not None:
            sim_line.set_data([], [])
        return [ball] if sim_line is None else [ball, sim_line]
    
    # Animation update function
    def update(frame):
        # Scale frame to data length
        idx = min(int(frame * len(time_data) / 200), len(time_data) - 1)
        
        # Update ball position
        t = time_data[idx]
        y = position_data[idx]
        ball.center = (t, max(0, y))  # Ensure ball doesn't go below ground
        
        # Update simulated trajectory line if available
        if sim_line is not None and sim_time is not None:
            # Show trajectory up to current time
            mask = sim_time <= t
            if np.any(mask):  # Make sure we have at least one point
                sim_line.set_data(sim_time[mask], sim_pos[mask])
        
        return [ball] if sim_line is None else [ball, sim_line]
    
    # Create animation
    ani = animation.FuncAnimation(
        fig, update, frames=200, init_func=init, 
        blit=True, interval=30, repeat=True
    )
    
    # Save animation if requested
    if save_animation:
        writer = PillowWriter(fps=30)
        filename = f"{ball_type.lower()}_{label.replace(' ', '_')}_bounce.gif"
        ani.save(filename, writer=writer)
        print(f"Animation saved as {filename}")
    
    plt.tight_layout()
    plt.show()


def process_and_animate_all(file_paths, 
                           ball_type, 
                           min_distance=10, 
                           min_time_diff=0.2,
                           relative_prominence=None,
                           low_height_threshold=8, 
                           low_height_adjustment=0.5,
                           high_height_threshold=15, 
                           high_height_adjustment=1.2, 
                           max_bounces=7, 
                           fs=None,
                           save_animations=False):
    """
    Processes each trial and creates an animated visualization for each one.
    """
    results = []
    for init_label, path in file_paths.items():
        try:
            initial_height = float(init_label.split()[0])
        except ValueError:
            initial_height = np.nan
        
        # Compute effective prominence
        effective_prominence = relative_prominence * initial_height
        if initial_height < low_height_threshold:
            effective_prominence *= low_height_adjustment
        elif initial_height > high_height_threshold:
            effective_prominence *= high_height_adjustment
        
        print(f"Processing trial '{init_label}' with effective prominence: {effective_prominence}")
        
        # Process the trial
        df, peak_indices, bounce_heights, bounce_times, cor_values, avg_cor, signal_data = process_trial(
            path, effective_prominence,
            min_distance=min_distance, 
            min_time_diff=min_time_diff, 
            max_bounces=max_bounces, 
            fs=fs
        )
        
        results.append({
            'Trial': init_label,
            'Initial Height': initial_height,
            'Average COR': avg_cor,
            'Bounce Times': bounce_times,
            'Bounce Heights': bounce_heights,
            'COR Values': cor_values
        })
        
        # Create animation
        animate_bouncing_ball(
            df, peak_indices, bounce_heights, bounce_times, avg_cor,
            label=init_label, ball_type=ball_type, initial_height=initial_height,
            g=386.09, save_animation=save_animations
        )
    
    summary_df = pd.DataFrame(results)
    summary_df.sort_values(by='Initial Height', inplace=True)
    return summary_df


def plot_cor_table(cor_df, ball_type):
    """
    Plots a line (and markers) of Average COR vs. Initial Height from the summary DataFrame.
    """
    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()


if __name__ == '__main__':
    # Process & animate Golf Ball trials
    print("\nProcessing and animating Golf Ball trials...")
    golf_results = process_and_animate_all(
        golf_file_paths,
        ball_type='Golf',
        min_distance=min_distance,
        min_time_diff=min_time_diff,
        relative_prominence=golf_relative_prominence,
        low_height_threshold=golf_low_threshold,
        low_height_adjustment=golf_low_adjustment,
        high_height_threshold=golf_high_threshold,
        high_height_adjustment=golf_high_adjustment,
        max_bounces=7,
        fs=fs,
        save_animations=save_animations
    )

    # Process & animate Lacrosse Ball trials
    print("\nProcessing and animating Lacrosse Ball trials...")
    lacrosse_results = process_and_animate_all(
        lacrosse_file_paths,
        ball_type='Lacrosse',
        min_distance=min_distance,
        min_time_diff=min_time_diff,
        relative_prominence=lacrosse_relative_prominence,
        low_height_threshold=lacrosse_low_threshold,
        low_height_adjustment=lacrosse_low_adjustment,
        high_height_threshold=lacrosse_high_threshold,
        high_height_adjustment=lacrosse_high_adjustment,
        max_bounces=7,
        fs=fs,
        save_animations=save_animations
    )

    # Process & animate Metal Ball trials
    print("\nProcessing and animating Metal Ball trials...")
    metal_results = process_and_animate_all(
        metal_file_paths,
        ball_type='Metal',
        min_distance=min_distance,
        min_time_diff=min_time_diff,
        relative_prominence=metal_relative_prominence,
        low_height_threshold=metal_low_threshold,
        low_height_adjustment=metal_low_adjustment,
        high_height_threshold=metal_high_threshold,
        high_height_adjustment=metal_high_adjustment,
        max_bounces=7,
        fs=fs,
        save_animations=save_animations
    )

    # Print summary tables
    print("\nGolf Ball COR Table:")
    print(golf_results[['Trial', 'Initial Height', 'Average COR']])
    print("\nLacrosse Ball COR Table:")
    print(lacrosse_results[['Trial', 'Initial Height', 'Average COR']])
    print("\nMetal Ball COR Table:")
    print(metal_results[['Trial', 'Initial Height', 'Average COR']])
    
    # Plot COR vs. initial height for each ball type
    plot_cor_table(golf_results, 'Golf')
    plot_cor_table(lacrosse_results, 'Lacrosse')
    plot_cor_table(metal_results, 'Metal') 