"""
Script to organize data files into the new directory structure.
"""

import os
import shutil
import glob

def create_data_directories():
    """
    Create the necessary directories for data organization.
    """
    # Create data directories
    os.makedirs("src/data/golf", exist_ok=True)
    os.makedirs("src/data/lacrosse", exist_ok=True)
    os.makedirs("src/data/metal", exist_ok=True)
    os.makedirs("src/data/images", exist_ok=True)
    
    # Create output directory
    os.makedirs("output", exist_ok=True)
    
    print("Created directory structure for data organization.")


def move_data_files():
    """
    Move data files to their appropriate directories.
    """
    # Move golf ball data
    for file in glob.glob("golf_*.csv"):
        shutil.copy(file, os.path.join("src/data/golf", file))
        print(f"Copied {file} to src/data/golf/")
    
    # Move lacrosse ball data
    for file in glob.glob("l_*.csv"):
        shutil.copy(file, os.path.join("src/data/lacrosse", file))
        print(f"Copied {file} to src/data/lacrosse/")
    
    # Move metal ball data
    for file in glob.glob("metal_*.csv"):
        shutil.copy(file, os.path.join("src/data/metal", file))
        print(f"Copied {file} to src/data/metal/")
    
    # Move image files
    for file in glob.glob("*.png"):
        shutil.copy(file, os.path.join("src/data/images", file))
        print(f"Copied {file} to src/data/images/")


def move_script_files():
    """
    Move script files to their appropriate directories.
    """
    # Move original scripts to src directory
    if os.path.exists("lab2_part1.py"):
        shutil.copy("lab2_part1.py", os.path.join("src", "lab2_part1.py"))
        print(f"Copied lab2_part1.py to src/")
    
    if os.path.exists("lab2_part2.py"):
        shutil.copy("lab2_part2.py", os.path.join("src", "lab2_part2.py"))
        print(f"Copied lab2_part2.py to src/")
    
    # Move animated scripts to visualization directory
    if os.path.exists("lab2_part1_animated.py"):
        shutil.copy("lab2_part1_animated.py", os.path.join("src/visualization", "lab2_part1_animated.py"))
        print(f"Copied lab2_part1_animated.py to src/visualization/")
    
    if os.path.exists("lab2_part2_animated.py"):
        shutil.copy("lab2_part2_animated.py", os.path.join("src/visualization", "lab2_part2_animated.py"))
        print(f"Copied lab2_part2_animated.py to src/visualization/")
    
    if os.path.exists("animate_bouncing_balls.py"):
        shutil.copy("animate_bouncing_balls.py", os.path.join("src/visualization", "animate_bouncing_balls.py"))
        print(f"Copied animate_bouncing_balls.py to src/visualization/")


def main():
    """
    Main function to organize files.
    """
    print("Starting file organization...")
    
    # Create directories
    create_data_directories()
    
    # Move data files
    move_data_files()
    
    # Move script files
    move_script_files()
    
    print("File organization complete!")


if __name__ == "__main__":
    main() 