#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from collections import defaultdict
import argparse
import re

def parse_line(line):
    """Parse a line from the simulation output."""
    # Skip comments and empty lines
    if line.startswith('#') or not line.strip():
        return None
    
    # Parse the line format: body_name: X = x1m x2m x3m, V = v1m/s v2m/s v3m/s
    try:
        # Split on colon to separate name and data
        name_part, data_part = line.strip().split(':')
        name = name_part.strip()
        
        # Split data part into position and velocity
        pos_part, vel_part = data_part.split(',')
        
        # Extract position values
        pos_str = pos_part.split('=')[1].strip()
        pos_values = [float(x.replace('m', '').strip()) for x in pos_str.split() if x.strip()]
        
        if len(pos_values) != 3:
            return None
            
        return {
            'name': name,
            'position': tuple(pos_values)
        }
    except (ValueError, IndexError):
        return None

def read_output_file(filename):
    """Read the simulation output file and organize data by body and time."""
    positions = defaultdict(list)
    times = []  # We'll use frame numbers as times since actual time isn't in output
    
    with open(filename, 'r') as f:
        frame = 0
        for line in f:
            data = parse_line(line)
            if data is None:
                continue
                
            name = data['name']
            pos = data['position']
            positions[name].append((pos[0], pos[1]))
            frame += 1
            times.append(frame)
    
    return positions, times

def create_animation(positions, times, output_file=None):
    """Create an animation of the bodies' orbits."""
    # Check if we have any data
    if not positions or not times:
        print("Error: No valid data found in the input file")
        return
        
    # Set up the figure and axis
    fig, ax = plt.subplots(figsize=(10, 10))
    
    # Set equal aspect ratio
    ax.set_aspect('equal')
    
    # Create scatter plots for each body
    scatters = {}
    for name in positions.keys():
        scatters[name], = ax.plot([], [], 'o-', label=name, alpha=0.7)
    
    # Set up the plot
    ax.set_xlabel('X (m)')
    ax.set_ylabel('Y (m)')
    ax.set_title('Orbital Simulation')
    ax.legend()
    
    # Find the bounds of the plot
    all_x = []
    all_y = []
    for pos_list in positions.values():
        if pos_list:  # Only process if we have positions
            x, y = zip(*pos_list)
            all_x.extend(x)
            all_y.extend(y)
    
    if not all_x or not all_y:
        print("Error: No valid position data found")
        return
        
    max_range = max(max(abs(min(all_x)), abs(max(all_x))),
                   max(abs(min(all_y)), abs(max(all_y))))
    ax.set_xlim(-max_range, max_range)
    ax.set_ylim(-max_range, max_range)
    
    def init():
        """Initialize the animation."""
        for scatter in scatters.values():
            scatter.set_data([], [])
        return list(scatters.values())
    
    def update(frame):
        """Update the animation for each frame."""
        for name, scatter in scatters.items():
            if positions[name]:  # Only update if we have positions
                x, y = zip(*positions[name][:frame+1])
                scatter.set_data(x, y)
        return list(scatters.values())
    
    # Create the animation
    anim = FuncAnimation(fig, update, frames=len(times),
                        init_func=init, blit=True,
                        interval=50)  # 50ms between frames
    
    if output_file:
        anim.save(output_file, writer='ffmpeg', fps=30)
    else:
        plt.show()

def main():
    parser = argparse.ArgumentParser(description='Animate orbital simulation output')
    parser.add_argument('input_file', help='Input file from simulation')
    parser.add_argument('--output', '-o', help='Output video file (optional)')
    args = parser.parse_args()
    
    positions, times = read_output_file(args.input_file)
    create_animation(positions, times, args.output)

if __name__ == '__main__':
    main()