#!/usr/bin/env python3
"""
Inspect orbital trajectory data from binary files in a tabular format.

This script reads the binary trajectory files generated by the orbital simulator
and displays the data in a nicely formatted table for inspection and debugging.

Usage:
    python inspect_trajectories.py <trajectory_file.bin>
"""

import sys
import struct
import argparse
import numpy as np
from collections import defaultdict

class BinaryReader:
    def __init__(self, data):
        self.data = data
        self.pos = 0
    
    def read_u64(self):
        result = struct.unpack('<Q', self.data[self.pos:self.pos+8])[0]
        self.pos += 8
        return result
    
    def read_f64(self):
        result = struct.unpack('<d', self.data[self.pos:self.pos+8])[0]
        self.pos += 8
        return result
    
    def read_string(self):
        # Read length (u64) then string bytes
        length = self.read_u64()
        result = self.data[self.pos:self.pos+length].decode('utf-8')
        self.pos += length
        return result
    
    def read_vec3(self):
        # Read 3 f64 values for position/velocity/acceleration
        x = self.read_f64()
        y = self.read_f64()
        z = self.read_f64()
        return np.array([x, y, z])

def read_trajectory_file(filename):
    """Read the binary trajectory file and return parsed data."""
    with open(filename, 'rb') as f:
        data = f.read()
    
    reader = BinaryReader(data)
    snapshots = []
    
    try:
        while reader.pos < len(data):
            # Read snapshot
            time = reader.read_f64()
            
            # Read number of bodies (u64)
            num_bodies = reader.read_u64()
            
            bodies = []
            for _ in range(num_bodies):
                # Read Body struct
                name = reader.read_string()
                mass = reader.read_f64()
                position = reader.read_vec3()
                velocity = reader.read_vec3()
                acceleration = reader.read_vec3()
                
                bodies.append({
                    'name': name,
                    'mass': mass,
                    'position': position,
                    'velocity': velocity,
                    'acceleration': acceleration
                })
            
            snapshots.append({
                'time': time,
                'bodies': bodies
            })
            
    except struct.error:
        # End of file or corrupted data
        pass
    
    return snapshots

def format_scientific(value, precision=3):
    """Format number in scientific notation with specified precision."""
    if abs(value) < 1e-6 or abs(value) >= 1e6:
        return f"{value:.{precision}e}"
    else:
        return f"{value:.{precision}f}"

def format_vector(vec, precision=3):
    """Format a 3D vector in a compact form."""
    x, y, z = vec
    return f"({format_scientific(x, precision)}, {format_scientific(y, precision)}, {format_scientific(z, precision)})"

def format_time(seconds):
    """Format time in a human-readable format."""
    if seconds == 0:
        return "0.0s"
    elif seconds < 60:
        return f"{seconds:.1f}s"
    elif seconds < 3600:
        return f"{seconds/60:.1f}m"
    elif seconds < 86400:
        return f"{seconds/3600:.1f}h"
    else:
        return f"{seconds/86400:.1f}d"

def print_summary_table(snapshots, max_rows=15):
    """Print a summary table of all snapshots."""
    print("📊 TRAJECTORY SUMMARY")
    print("=" * 100)
    
    # Header
    print(f"{'Step':<6} {'Time':<12} {'Bodies':<8} {'Sample Position (first body)':<35} {'Sample Velocity (first body)':<35}")
    print("-" * 100)
    
    # Determine which snapshots to show
    total_snapshots = len(snapshots)
    if total_snapshots <= max_rows:
        indices = list(range(total_snapshots))
    else:
        # Show first few, middle few, and last few
        start_count = max_rows // 3
        end_count = max_rows // 3
        middle_count = max_rows - start_count - end_count
        
        indices = []
        indices.extend(range(start_count))
        
        if middle_count > 0:
            middle_start = total_snapshots // 2 - middle_count // 2
            indices.extend(range(middle_start, middle_start + middle_count))
        
        indices.extend(range(total_snapshots - end_count, total_snapshots))
        indices = sorted(set(indices))  # Remove duplicates and sort
    
    prev_index = -1
    for i, idx in enumerate(indices):
        if idx > prev_index + 1 and prev_index >= 0:
            print("    ...")
        
        snapshot = snapshots[idx]
        time_str = format_time(snapshot['time'])
        body_count = len(snapshot['bodies'])
        
        if snapshot['bodies']:
            first_body = snapshot['bodies'][0]
            pos_str = format_vector(first_body['position'], 2)
            vel_str = format_vector(first_body['velocity'], 2)
        else:
            pos_str = "N/A"
            vel_str = "N/A"
        
        print(f"{idx:<6} {time_str:<12} {body_count:<8} {pos_str:<35} {vel_str:<35}")
        prev_index = idx
    
    print("-" * 100)
    print(f"Total snapshots: {total_snapshots}")

def print_detailed_table(snapshots, body_name=None, max_rows=15):
    """Print detailed information for a specific body."""
    if not snapshots:
        print("No data to display!")
        return
    
    # Get all body names
    all_bodies = set()
    for snapshot in snapshots:
        for body in snapshot['bodies']:
            all_bodies.add(body['name'])
    
    if body_name and body_name not in all_bodies:
        print(f"❌ Body '{body_name}' not found. Available bodies: {', '.join(sorted(all_bodies))}")
        return
    
    # If no body specified, show the first one
    if not body_name:
        body_name = sorted(all_bodies)[0]
    
    print(f"🌍 DETAILED VIEW: {body_name}")
    print("=" * 120)
    
    # Header
    print(f"{'Step':<6} {'Time':<12} {'Position (x, y, z)':<40} {'Velocity (x, y, z)':<40} {'|v|':<12} {'|a|':<12}")
    print("-" * 120)
    
    # Determine which snapshots to show
    total_snapshots = len(snapshots)
    if total_snapshots <= max_rows:
        indices = list(range(total_snapshots))
    else:
        # Show first few, middle few, and last few
        start_count = max_rows // 3
        end_count = max_rows // 3
        middle_count = max_rows - start_count - end_count
        
        indices = []
        indices.extend(range(start_count))
        
        if middle_count > 0:
            middle_start = total_snapshots // 2 - middle_count // 2
            indices.extend(range(middle_start, middle_start + middle_count))
        
        indices.extend(range(total_snapshots - end_count, total_snapshots))
        indices = sorted(set(indices))
    
    prev_index = -1
    for idx in indices:
        if idx > prev_index + 1 and prev_index >= 0:
            print("    ...")
        
        snapshot = snapshots[idx]
        time_str = format_time(snapshot['time'])
        
        # Find the body in this snapshot
        body_data = None
        for body in snapshot['bodies']:
            if body['name'] == body_name:
                body_data = body
                break
        
        if body_data:
            pos_str = format_vector(body_data['position'], 3)
            vel_str = format_vector(body_data['velocity'], 3)
            vel_mag = np.linalg.norm(body_data['velocity'])
            acc_mag = np.linalg.norm(body_data['acceleration'])
            vel_mag_str = format_scientific(vel_mag, 2)
            acc_mag_str = format_scientific(acc_mag, 2)
            
            print(f"{idx:<6} {time_str:<12} {pos_str:<40} {vel_str:<40} {vel_mag_str:<12} {acc_mag_str:<12}")
        else:
            print(f"{idx:<6} {time_str:<12} {'BODY NOT FOUND':<40}")
        
        prev_index = idx
    
    print("-" * 120)

def print_statistics(snapshots):
    """Print statistical information about the trajectory."""
    if not snapshots:
        return
    
    print("\n📈 TRAJECTORY STATISTICS")
    print("=" * 80)
    
    # Time statistics
    times = [s['time'] for s in snapshots]
    time_start, time_end = times[0], times[-1]
    duration = time_end - time_start
    
    print(f"Time range: {format_time(time_start)} to {format_time(time_end)}")
    print(f"Total duration: {format_time(duration)}")
    print(f"Number of snapshots: {len(snapshots)}")
    
    if len(times) > 1:
        time_steps = [times[i+1] - times[i] for i in range(len(times)-1)]
        avg_step = np.mean(time_steps)
        print(f"Average time step: {format_time(avg_step)}")
    
    # Body statistics
    body_names = set()
    for snapshot in snapshots:
        for body in snapshot['bodies']:
            body_names.add(body['name'])
    
    print(f"Bodies tracked: {', '.join(sorted(body_names))}")
    
    # Position and velocity ranges for each body
    for body_name in sorted(body_names):
        positions = []
        velocities = []
        
        for snapshot in snapshots:
            for body in snapshot['bodies']:
                if body['name'] == body_name:
                    positions.append(body['position'])
                    velocities.append(body['velocity'])
        
        if positions:
            positions = np.array(positions)
            velocities = np.array(velocities)
            
            pos_min = np.min(positions, axis=0)
            pos_max = np.max(positions, axis=0)
            vel_min = np.min(np.linalg.norm(velocities, axis=1))
            vel_max = np.max(np.linalg.norm(velocities, axis=1))
            
            print(f"\n{body_name}:")
            print(f"  Position range: X[{format_scientific(pos_min[0])}, {format_scientific(pos_max[0])}]")
            print(f"                  Y[{format_scientific(pos_min[1])}, {format_scientific(pos_max[1])}]")
            print(f"                  Z[{format_scientific(pos_min[2])}, {format_scientific(pos_max[2])}]")
            print(f"  Speed range: {format_scientific(vel_min)} to {format_scientific(vel_max)} m/s")

def main():
    parser = argparse.ArgumentParser(description='Inspect orbital trajectory data in tabular format')
    parser.add_argument('trajectory_file', help='Binary trajectory file to inspect')
    parser.add_argument('--rows', '-r', type=int, default=15, help='Maximum number of rows to display (default: 15)')
    parser.add_argument('--body', '-b', type=str, help='Show detailed view for specific body')
    parser.add_argument('--summary', '-s', action='store_true', help='Show summary of all snapshots')
    parser.add_argument('--stats', action='store_true', help='Show trajectory statistics')
    parser.add_argument('--all', '-a', action='store_true', help='Show all available information')
    
    args = parser.parse_args()
    
    print(f"🔍 Inspecting trajectory file: {args.trajectory_file}")
    print()
    
    try:
        snapshots = read_trajectory_file(args.trajectory_file)
        
        if not snapshots:
            print("❌ No data found in file!")
            return
        
        # Determine what to show
        show_summary = args.summary or args.all or (not args.body and not args.stats)
        show_detailed = args.body or args.all
        show_stats = args.stats or args.all
        
        if show_summary:
            print_summary_table(snapshots, args.rows)
            print()
        
        if show_detailed:
            print_detailed_table(snapshots, args.body, args.rows)
            print()
        
        if show_stats:
            print_statistics(snapshots)
        
    except FileNotFoundError:
        print(f"❌ Error: File '{args.trajectory_file}' not found!")
    except Exception as e:
        print(f"❌ Error reading file: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()