"""
Script to create a Jupyter notebook example for the bouncing ball analysis project.
"""

import json
import os

def create_notebook():
    """
    Create a Jupyter notebook with examples of using the bouncing ball analysis modules.
    """
    notebook = {
        "cells": [
            # Title and introduction
            {
                "cell_type": "markdown",
                "metadata": {},
                "source": [
                    "# Bouncing Ball Analysis Example\n",
                    "\n",
                    "This notebook demonstrates how to use the modular structure of the bouncing ball analysis project to:\n",
                    "1. Load data for different ball types\n",
                    "2. Detect bounces and calculate Coefficient of Restitution (COR)\n",
                    "3. Create static and animated visualizations"
                ]
            },
            
            # Setup section
            {
                "cell_type": "markdown",
                "metadata": {},
                "source": [
                    "## Setup\n",
                    "\n",
                    "First, let's import the necessary modules and set up the environment."
                ]
            },
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "import os\n",
                    "import sys\n",
                    "import numpy as np\n",
                    "import pandas as pd\n",
                    "import matplotlib.pyplot as plt\n",
                    "from IPython.display import display, HTML\n",
                    "\n",
                    "# Add the src directory to the Python path\n",
                    "sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname('__file__'), '..')))\n",
                    "\n",
                    "# Import modules from the project\n",
                    "from data.loader import load_trial, GOLF_BALL_PATHS, LACROSSE_BALL_PATHS, METAL_BALL_PATHS\n",
                    "from analysis.bounce_detection import process_trial, process_trials\n",
                    "from visualization.static_plots import plot_position_vs_time, plot_trial_with_bounces, plot_cor_table, plot_all_cors\n",
                    "from visualization.animation import create_ball_animation, create_ball_animation_with_model, display_animation\n",
                    "\n",
                    "# Create output directory\n",
                    "output_dir = \"../../output\"\n",
                    "os.makedirs(output_dir, exist_ok=True)"
                ]
            },
            
            # Section 1: Loading and Visualizing Raw Data
            {
                "cell_type": "markdown",
                "metadata": {},
                "source": [
                    "## 1. Loading and Visualizing Raw Data\n",
                    "\n",
                    "Let's start by loading data for a golf ball dropped from 11 inches and visualizing the raw position vs. time data."
                ]
            },
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "# Select a specific ball and height\n",
                    "ball_type = \"Golf\"\n",
                    "height = \"11 inches\"\n",
                    "file_path = GOLF_BALL_PATHS[height]\n",
                    "\n",
                    "# Load the trial data\n",
                    "df = load_trial(file_path)\n",
                    "print(f\"Loaded data with {len(df)} rows\")\n",
                    "\n",
                    "# Display the first few rows of the data\n",
                    "df.head()"
                ]
            },
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "# Plot the raw position vs. time data\n",
                    "plot_position_vs_time(df, height, ball_type)"
                ]
            },
            
            # Section 2: Detecting Bounces and Calculating COR
            {
                "cell_type": "markdown",
                "metadata": {},
                "source": [
                    "## 2. Detecting Bounces and Calculating COR\n",
                    "\n",
                    "Now, let's detect the bounces in the data and calculate the Coefficient of Restitution (COR)."
                ]
            },
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "# Process the trial to detect bounces and calculate COR\n",
                    "result = process_trial(\n",
                    "    df, \n",
                    "    initial_height=float(height.split()[0]),\n",
                    "    ball_type=ball_type\n",
                    ")\n",
                    "\n",
                    "# Print the results\n",
                    "print(f\"Number of bounces detected: {len(result['peak_indices'])}\")\n",
                    "print(f\"Average COR: {result['Average COR']:.4f}\")\n",
                    "print(f\"Initial height: {result['Initial Height']} inches\")\n",
                    "\n",
                    "# Display the bounce heights\n",
                    "if 'bounce_heights' in result:\n",
                    "    print(\"\\nBounce Heights (inches):\")\n",
                    "    for i, height in enumerate(result['bounce_heights']):\n",
                    "        print(f\"Bounce {i+1}: {height:.2f}\")"
                ]
            },
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "# Plot the trial with detected bounces\n",
                    "plot_trial_with_bounces(\n",
                    "    df, \n",
                    "    result['peak_indices'], \n",
                    "    df['Position'].values,\n",
                    "    height, \n",
                    "    ball_type,\n",
                    "    initial_height=float(height.split()[0]),\n",
                    "    cor=result['Average COR']\n",
                    ")"
                ]
            },
            
            # Section 3: Creating Animations
            {
                "cell_type": "markdown",
                "metadata": {},
                "source": [
                    "## 3. Creating Animations\n",
                    "\n",
                    "Let's create an animation of the bouncing ball."
                ]
            },
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "# Create a simple animation of the bouncing ball\n",
                    "anim = create_ball_animation(\n",
                    "    df,\n",
                    "    ball_radius=0.5,\n",
                    "    fps=30,\n",
                    "    duration=2.0,  # Only show the first 2 seconds\n",
                    "    title=f\"{ball_type} Ball - {height}\"\n",
                    ")\n",
                    "\n",
                    "# Display the animation\n",
                    "display_animation(anim)"
                ]
            },
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "# Create an animation with both the actual data and an ideal model\n",
                    "anim_with_model = create_ball_animation_with_model(\n",
                    "    df, \n",
                    "    result['peak_indices'], \n",
                    "    result['Average COR'], \n",
                    "    float(height.split()[0]),\n",
                    "    ball_radius=0.5,\n",
                    "    fps=30,\n",
                    "    duration=2.0,  # Only show the first 2 seconds\n",
                    "    title=f\"{ball_type} Ball - {height} (with Model)\"\n",
                    ")\n",
                    "\n",
                    "# Display the animation\n",
                    "display_animation(anim_with_model)"
                ]
            },
            
            # Section 4: Analyzing Multiple Trials
            {
                "cell_type": "markdown",
                "metadata": {},
                "source": [
                    "## 4. Analyzing Multiple Trials\n",
                    "\n",
                    "Now, let's analyze all the golf ball trials and compare the COR values."
                ]
            },
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "# Process all golf ball trials\n",
                    "golf_trials = []\n",
                    "\n",
                    "for height, path in GOLF_BALL_PATHS.items():\n",
                    "    print(f\"Processing golf ball from {height}...\")\n",
                    "    \n",
                    "    # Load the trial data\n",
                    "    df = load_trial(path)\n",
                    "    \n",
                    "    # Process the trial\n",
                    "    result = process_trial(\n",
                    "        df, \n",
                    "        initial_height=float(height.split()[0]),\n",
                    "        ball_type=\"Golf\"\n",
                    "    )\n",
                    "    \n",
                    "    # Store the result\n",
                    "    result['Initial Height'] = float(height.split()[0])\n",
                    "    result['Ball Type'] = \"Golf\"\n",
                    "    result['Path'] = path\n",
                    "    golf_trials.append(result)\n",
                    "\n",
                    "# Create a summary DataFrame\n",
                    "golf_summary = pd.DataFrame(golf_trials)\n",
                    "\n",
                    "# Display the summary\n",
                    "golf_summary[['Initial Height', 'Average COR', 'Num Bounces']]"
                ]
            },
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "# Plot COR vs. initial height for golf balls\n",
                    "plot_cor_table(golf_summary, \"Golf\")"
                ]
            },
            
            # Section 5: Comparing Different Ball Types
            {
                "cell_type": "markdown",
                "metadata": {},
                "source": [
                    "## 5. Comparing Different Ball Types\n",
                    "\n",
                    "Finally, let's compare the COR values for different ball types."
                ]
            },
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "# Process lacrosse ball trials\n",
                    "lacrosse_trials = []\n",
                    "\n",
                    "for height, path in LACROSSE_BALL_PATHS.items():\n",
                    "    print(f\"Processing lacrosse ball from {height}...\")\n",
                    "    \n",
                    "    # Load the trial data\n",
                    "    df = load_trial(path)\n",
                    "    \n",
                    "    # Process the trial\n",
                    "    result = process_trial(\n",
                    "        df, \n",
                    "        initial_height=float(height.split()[0]),\n",
                    "        ball_type=\"Lacrosse\"\n",
                    "    )\n",
                    "    \n",
                    "    # Store the result\n",
                    "    result['Initial Height'] = float(height.split()[0])\n",
                    "    result['Ball Type'] = \"Lacrosse\"\n",
                    "    result['Path'] = path\n",
                    "    lacrosse_trials.append(result)\n",
                    "\n",
                    "# Create a summary DataFrame\n",
                    "lacrosse_summary = pd.DataFrame(lacrosse_trials)\n",
                    "\n",
                    "# Process metal ball trials\n",
                    "metal_trials = []\n",
                    "\n",
                    "for height, path in METAL_BALL_PATHS.items():\n",
                    "    print(f\"Processing metal ball from {height}...\")\n",
                    "    \n",
                    "    # Load the trial data\n",
                    "    df = load_trial(path)\n",
                    "    \n",
                    "    # Process the trial\n",
                    "    result = process_trial(\n",
                    "        df, \n",
                    "        initial_height=float(height.split()[0]),\n",
                    "        ball_type=\"Metal\"\n",
                    "    )\n",
                    "    \n",
                    "    # Store the result\n",
                    "    result['Initial Height'] = float(height.split()[0])\n",
                    "    result['Ball Type'] = \"Metal\"\n",
                    "    result['Path'] = path\n",
                    "    metal_trials.append(result)\n",
                    "\n",
                    "# Create a summary DataFrame\n",
                    "metal_summary = pd.DataFrame(metal_trials)"
                ]
            },
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "# Plot all CORs together\n",
                    "plot_all_cors(golf_summary, lacrosse_summary, metal_summary)\n",
                    "\n",
                    "# Save the combined plot\n",
                    "plt.savefig(os.path.join(output_dir, \"all_cors_comparison.png\"), dpi=300)"
                ]
            },
            
            # Conclusion
            {
                "cell_type": "markdown",
                "metadata": {},
                "source": [
                    "## Conclusion\n",
                    "\n",
                    "In this notebook, we've demonstrated how to use the modular structure of the bouncing ball analysis project to:\n",
                    "1. Load data for different ball types\n",
                    "2. Detect bounces and calculate Coefficient of Restitution (COR)\n",
                    "3. Create static and animated visualizations\n",
                    "4. Compare COR values for different ball types\n",
                    "\n",
                    "The modular structure makes it easy to perform these analyses and visualizations with clean, reusable code."
                ]
            }
        ],
        "metadata": {
            "kernelspec": {
                "display_name": "Python 3",
                "language": "python",
                "name": "python3"
            },
            "language_info": {
                "codemirror_mode": {
                    "name": "ipython",
                    "version": 3
                },
                "file_extension": ".py",
                "mimetype": "text/x-python",
                "name": "python",
                "nbconvert_exporter": "python",
                "pygments_lexer": "ipython3",
                "version": "3.8.10"
            }
        },
        "nbformat": 4,
        "nbformat_minor": 4
    }
    
    # Write the notebook to a file
    notebook_path = "src/examples/bouncing_ball_analysis.ipynb"
    os.makedirs(os.path.dirname(notebook_path), exist_ok=True)
    with open(notebook_path, "w") as f:
        json.dump(notebook, f, indent=1)
    
    print(f"Jupyter notebook created: {notebook_path}")


if __name__ == "__main__":
    create_notebook() 