"""
Main script for the bouncing ball analysis project.

This script serves as the entry point for analyzing bouncing ball data,
detecting bounces, calculating COR values, and generating visualizations.
"""

import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from IPython.display import display, HTML

# Import modules from the project
from data.loader import load_trial, GOLF_BALL_PATHS, LACROSSE_BALL_PATHS, METAL_BALL_PATHS
from analysis.bounce_detection import process_trial, process_trials, BALL_PARAMS
from visualization.static_plots import (
    plot_position_vs_time, 
    plot_trial_with_bounces, 
    plot_cor_table,
    plot_all_cors
)
from visualization.animation import (
    create_ball_animation, 
    create_ball_animation_with_model,
    display_animation,
    save_animation
)


def analyze_ball_type(ball_type, paths_dict, output_dir="output"):
    """
    Analyze a specific type of ball using the provided paths.
    
    Parameters:
        ball_type (str): Type of ball (e.g., "Golf", "Lacrosse", "Metal")
        paths_dict (dict): Dictionary mapping initial heights to file paths
        output_dir (str): Directory to save output files
        
    Returns:
        pandas.DataFrame: Summary DataFrame with trial information
    """
    print(f"\n{'='*50}")
    print(f"Analyzing {ball_type} Ball Data")
    print(f"{'='*50}")
    
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
    
    # Process all trials for this ball type
    trials_data = []
    
    for height, path in paths_dict.items():
        print(f"\nProcessing {ball_type} ball from {height} inches...")
        
        # Load the trial data
        df = load_trial(path)
        
        # Process the trial
        result = process_trial(
            df, 
            initial_height=float(height.split()[0]),
            ball_type=ball_type
        )
        
        # Store the result
        result['Initial Height'] = float(height.split()[0])
        result['Ball Type'] = ball_type
        result['Path'] = path
        trials_data.append(result)
        
        # Plot the trial with detected bounces
        plot_trial_with_bounces(
            df, 
            result['peak_indices'], 
            df['Position'].values,
            height, 
            ball_type,
            initial_height=float(height.split()[0]),
            cor=result['Average COR']
        )
        
        # Create and save animation
        anim = create_ball_animation_with_model(
            df, 
            result['peak_indices'], 
            result['Average COR'], 
            float(height.split()[0]),
            title=f"{ball_type} Ball - {height}"
        )
        
        # Save the animation
        animation_path = os.path.join(output_dir, f"{ball_type.lower()}_{height.split()[0]}_animation.mp4")
        save_animation(anim, animation_path)
        print(f"Animation saved to {animation_path}")
    
    # Create a summary DataFrame
    summary_df = pd.DataFrame(trials_data)
    
    # Plot COR vs. initial height
    plot_cor_table(summary_df, ball_type)
    
    # Save the summary DataFrame
    summary_path = os.path.join(output_dir, f"{ball_type.lower()}_summary.csv")
    summary_df.to_csv(summary_path, index=False)
    print(f"Summary data saved to {summary_path}")
    
    return summary_df


def main():
    """
    Main function to run the bouncing ball analysis.
    """
    print("Starting Bouncing Ball Analysis")
    
    # Create output directory
    output_dir = "output"
    os.makedirs(output_dir, exist_ok=True)
    
    # Analyze each ball type
    golf_summary = analyze_ball_type("Golf", GOLF_BALL_PATHS, output_dir)
    lacrosse_summary = analyze_ball_type("Lacrosse", LACROSSE_BALL_PATHS, output_dir)
    metal_summary = analyze_ball_type("Metal", METAL_BALL_PATHS, output_dir)
    
    # Plot all CORs together
    plot_all_cors(golf_summary, lacrosse_summary, metal_summary)
    
    # Save the combined plot
    plt.savefig(os.path.join(output_dir, "all_cors_comparison.png"), dpi=300)
    
    print("\nAnalysis complete! Results saved to the 'output' directory.")


if __name__ == "__main__":
    main() 