import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.signal import find_peaks

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

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 plot_trial(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'.
    Ensures:
      - Displays exactly 7 bounces.
      - X-axis starts just before the first data point and ends slightly after the last peak or data point.
      - Y-axis is a tight fit to cover the full range of data.
    """
    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 process_and_plot_all(file_paths, 
                         ball_type, 
                         plot_trials=False, 
                         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):
    """
    Iterates over the given 'file_paths' dict:
      - extracts the numeric initial height from the key label (e.g. "11 inches"),
      - computes effective prominence,
      - processes each trial,
      - optionally plots each trial (with or without an ideal bounce overlay).
    
    Returns a summary DataFrame with:
      'Trial', 'Initial Height', 'Average COR', 'Bounce Times', 'Bounce Heights', and 'COR Values'.
    """
    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
        })
        
        
        if plot_trials:
            plot_trial(
                df, 
                peak_indices, 
                signal_data, 
                label=init_label, 
                ball_type=ball_type,
                initial_height=initial_height,   
                cor=avg_cor,
                bounces=5, 
                dt=0.001, 
                g=386.09     #using inches for position
            )
    
    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__':

    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"
    }
    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"
    }
    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

    # 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

    # Process & plot Golf
    golf_results = process_and_plot_all(
        golf_file_paths,
        ball_type='Golf',
        plot_trials=True,
        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
    )

    # Process & plot Lacrosse
    lacrosse_results = process_and_plot_all(
        lacrosse_file_paths,
        ball_type='Lacrosse',
        plot_trials=True,
        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
    )

    # Process & plot Metal
    metal_results = process_and_plot_all(
        metal_file_paths,
        ball_type='Metal',
        plot_trials=True,
        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
    )

    # Print summary tables
    print("Golf 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')