#!/usr/bin/env python3
"""
Bouncing Ball Animation Launcher

This script provides a simple command-line interface to run different
animated visualizations of bouncing ball data.
"""

import sys
import argparse
import importlib.util
import os

def check_dependencies():
    """Check if all required dependencies are installed."""
    required_packages = ['pandas', 'numpy', 'matplotlib', 'seaborn', 'scipy']
    missing_packages = []
    
    for package in required_packages:
        try:
            importlib.import_module(package)
        except ImportError:
            missing_packages.append(package)
    
    if missing_packages:
        print("Missing required packages:")
        for package in missing_packages:
            print(f"  - {package}")
        print("\nPlease install them using:")
        print(f"pip install {' '.join(missing_packages)}")
        return False
    
    return True

def check_files():
    """Check if the animation script files exist."""
    required_files = ['lab2_part1_animated.py', 'lab2_part2_animated.py']
    missing_files = []
    
    for file in required_files:
        if not os.path.exists(file):
            missing_files.append(file)
    
    if missing_files:
        print("Missing required script files:")
        for file in missing_files:
            print(f"  - {file}")
        return False
    
    return True

def check_data_files(ball_type=None):
    """Check if the required data files exist."""
    if ball_type is None or ball_type.lower() == 'golf':
        golf_files = ["golf_11.csv", "golf_12.csv", "golf_13.csv", "golf_14.csv", "golf_15.csv"]
        for file in golf_files:
            if not os.path.exists(file):
                print(f"Warning: Golf ball data file '{file}' not found.")
                return False
    
    if ball_type is None or ball_type.lower() == 'lacrosse':
        lacrosse_files = ["l_18.csv", "l_19.csv", "l_20.csv", "l_21.csv", "l_22.csv"]
        for file in lacrosse_files:
            if not os.path.exists(file):
                print(f"Warning: Lacrosse ball data file '{file}' not found.")
                return False
    
    if ball_type is None or ball_type.lower() == 'metal':
        metal_files = ["metal_2.csv", "metal_4.csv", "metal_6.csv", "metal_8.csv", "metal_10.csv"]
        for file in metal_files:
            if not os.path.exists(file):
                print(f"Warning: Metal ball data file '{file}' not found.")
                return False
    
    return True

def run_animation(animation_type, ball_type=None, save=False):
    """Run the selected animation."""
    if animation_type == 'position':
        # Import the position vs. time animation script
        spec = importlib.util.spec_from_file_location("lab2_part1_animated", "lab2_part1_animated.py")
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        
        # Run the animation
        if ball_type:
            if ball_type.lower() == 'golf':
                print("Running position vs. time animation for Golf balls...")
                module.animate_position_vs_time(module.golf_file_paths, "Golf Ball", save)
            elif ball_type.lower() == 'lacrosse':
                print("Running position vs. time animation for Lacrosse balls...")
                module.animate_position_vs_time(module.lacrosse_file_paths, "Lacrosse Ball", save)
            elif ball_type.lower() == 'metal':
                print("Running position vs. time animation for Metal balls...")
                module.animate_position_vs_time(module.metal_file_paths, "Metal Ball", save)
            else:
                print(f"Unknown ball type: {ball_type}")
                print("Available ball types: golf, lacrosse, metal")
        else:
            print("Running position vs. time animations for all ball types...")
            module.animate_all_ball_types(save)
    
    elif animation_type == 'bounce':
        # Import the bouncing ball animation script
        spec = importlib.util.spec_from_file_location("lab2_part2_animated", "lab2_part2_animated.py")
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        
        # Set save_animations flag
        module.save_animations = save
        
        # Run the animation
        if ball_type:
            if ball_type.lower() == 'golf':
                print("Running bouncing ball animation for Golf balls...")
                module.golf_results = module.process_and_animate_all(
                    module.golf_file_paths,
                    ball_type='Golf',
                    min_distance=module.min_distance,
                    min_time_diff=module.min_time_diff,
                    relative_prominence=module.golf_relative_prominence,
                    low_height_threshold=module.golf_low_threshold,
                    low_height_adjustment=module.golf_low_adjustment,
                    high_height_threshold=module.golf_high_threshold,
                    high_height_adjustment=module.golf_high_adjustment,
                    max_bounces=7,
                    fs=module.fs,
                    save_animations=save
                )
                module.plot_cor_table(module.golf_results, 'Golf')
            elif ball_type.lower() == 'lacrosse':
                print("Running bouncing ball animation for Lacrosse balls...")
                module.lacrosse_results = module.process_and_animate_all(
                    module.lacrosse_file_paths,
                    ball_type='Lacrosse',
                    min_distance=module.min_distance,
                    min_time_diff=module.min_time_diff,
                    relative_prominence=module.lacrosse_relative_prominence,
                    low_height_threshold=module.lacrosse_low_threshold,
                    low_height_adjustment=module.lacrosse_low_adjustment,
                    high_height_threshold=module.lacrosse_high_threshold,
                    high_height_adjustment=module.lacrosse_high_adjustment,
                    max_bounces=7,
                    fs=module.fs,
                    save_animations=save
                )
                module.plot_cor_table(module.lacrosse_results, 'Lacrosse')
            elif ball_type.lower() == 'metal':
                print("Running bouncing ball animation for Metal balls...")
                module.metal_results = module.process_and_animate_all(
                    module.metal_file_paths,
                    ball_type='Metal',
                    min_distance=module.min_distance,
                    min_time_diff=module.min_time_diff,
                    relative_prominence=module.metal_relative_prominence,
                    low_height_threshold=module.metal_low_threshold,
                    low_height_adjustment=module.metal_low_adjustment,
                    high_height_threshold=module.metal_high_threshold,
                    high_height_adjustment=module.metal_high_adjustment,
                    max_bounces=7,
                    fs=module.fs,
                    save_animations=save
                )
                module.plot_cor_table(module.metal_results, 'Metal')
            else:
                print(f"Unknown ball type: {ball_type}")
                print("Available ball types: golf, lacrosse, metal")
        else:
            # Run the animations for all ball types one by one
            print("Running bouncing ball animations for all ball types...")
            
            # Golf balls
            print("\nProcessing and animating Golf Ball trials...")
            module.golf_results = module.process_and_animate_all(
                module.golf_file_paths,
                ball_type='Golf',
                min_distance=module.min_distance,
                min_time_diff=module.min_time_diff,
                relative_prominence=module.golf_relative_prominence,
                low_height_threshold=module.golf_low_threshold,
                low_height_adjustment=module.golf_low_adjustment,
                high_height_threshold=module.golf_high_threshold,
                high_height_adjustment=module.golf_high_adjustment,
                max_bounces=7,
                fs=module.fs,
                save_animations=save
            )
            
            # Lacrosse balls
            print("\nProcessing and animating Lacrosse Ball trials...")
            module.lacrosse_results = module.process_and_animate_all(
                module.lacrosse_file_paths,
                ball_type='Lacrosse',
                min_distance=module.min_distance,
                min_time_diff=module.min_time_diff,
                relative_prominence=module.lacrosse_relative_prominence,
                low_height_threshold=module.lacrosse_low_threshold,
                low_height_adjustment=module.lacrosse_low_adjustment,
                high_height_threshold=module.lacrosse_high_threshold,
                high_height_adjustment=module.lacrosse_high_adjustment,
                max_bounces=7,
                fs=module.fs,
                save_animations=save
            )
            
            # Metal balls
            print("\nProcessing and animating Metal Ball trials...")
            module.metal_results = module.process_and_animate_all(
                module.metal_file_paths,
                ball_type='Metal',
                min_distance=module.min_distance,
                min_time_diff=module.min_time_diff,
                relative_prominence=module.metal_relative_prominence,
                low_height_threshold=module.metal_low_threshold,
                low_height_adjustment=module.metal_low_adjustment,
                high_height_threshold=module.metal_high_threshold,
                high_height_adjustment=module.metal_high_adjustment,
                max_bounces=7,
                fs=module.fs,
                save_animations=save
            )
            
            # Print summary tables
            print("\nGolf Ball COR Table:")
            print(module.golf_results[['Trial', 'Initial Height', 'Average COR']])
            print("\nLacrosse Ball COR Table:")
            print(module.lacrosse_results[['Trial', 'Initial Height', 'Average COR']])
            print("\nMetal Ball COR Table:")
            print(module.metal_results[['Trial', 'Initial Height', 'Average COR']])
            
            # Plot COR vs. initial height for each ball type
            module.plot_cor_table(module.golf_results, 'Golf')
            module.plot_cor_table(module.lacrosse_results, 'Lacrosse')
            module.plot_cor_table(module.metal_results, 'Metal')
    
    else:
        print(f"Unknown animation type: {animation_type}")
        print("Available animation types: position, bounce")

def main():
    """Main function to parse arguments and run the selected animation."""
    parser = argparse.ArgumentParser(description='Run bouncing ball animations.')
    parser.add_argument('animation_type', choices=['position', 'bounce'], 
                        help='Type of animation to run: position (for position vs. time) or bounce (for bouncing ball physics)')
    parser.add_argument('--ball', choices=['golf', 'lacrosse', 'metal'], 
                        help='Type of ball to animate (default: all)')
    parser.add_argument('--save', action='store_true', 
                        help='Save animations as GIF files')
    
    args = parser.parse_args()
    
    # Check dependencies and files
    if not check_dependencies():
        print("Please install the required dependencies and try again.")
        sys.exit(1)
    
    if not check_files():
        print("Please ensure the animation script files are in the current directory and try again.")
        sys.exit(1)
    
    if not check_data_files(args.ball):
        print("Warning: Some data files are missing. The animations may not work correctly.")
        response = input("Do you want to continue anyway? (y/n): ")
        if response.lower() != 'y':
            sys.exit(1)
    
    # Run the selected animation
    run_animation(args.animation_type, args.ball, args.save)

if __name__ == '__main__':
    main()