#!/usr/bin/env python3
"""
Flight Log Analyzer — Analyze PX4/ArduPilot .bin/.ulg flight logs.

This tool reads and plots key flight data for post-flight analysis.

Usage:
    python3 log_analyzer.py <logfile.ulg>
    python3 log_analyzer.py <logfile.bin>
"""

import sys
import argparse
import numpy as np
import matplotlib.pyplot as plt
from collections import defaultdict

try:
    from pyulog import ULog  # pip install pyulog
    HAS_ULOG = True
except ImportError:
    HAS_ULOG = False
    print("ULog parser not found. Install: pip install pyulog")


def analyze_ulog(filename):
    """Analyze a .ulg (PX4 binary log) file."""
    if not HAS_ULOG:
        print("Cannot analyze .ulg: pyulog not installed")
        return
    
    ulog = ULog(filename)
    data = ulog.data_list
    
    print(f"Log file: {filename}")
    print(f"Duration: {ulog.last_timestamp - ulog.start_timestamp:.1f}s")
    print(f"Messages: {len(data)} topics")
    
    # Extract key topics
    topics = {}
    for d in data:
        topics[d.name] = d
    
    # Create plots
    fig, axes = plt.subplots(4, 1, figsize=(12, 10), sharex=True)
    
    # 1. Attitude
    if 'vehicle_attitude' in topics:
        att = topics['vehicle_attitude']
        t = att.data['timestamp'] - att.data['timestamp'][0]
        # Extract Euler angles from quaternion
        q = [att.data['q[0]'], att.data['q[1]'], 
             att.data['q[2]'], att.data['q[3]']]
        roll = [np.arctan2(2*(q0*q1 + q2*q3), 1-2*(q1**2+q2**2)) 
                for q0,q1,q2,q3 in zip(*q)]
        pitch = [np.arcsin(2*(q0*q2 - q3*q1)) 
                 for q0,q1,q2,q3 in zip(*q)]
        yaw = [np.arctan2(2*(q0*q3 + q1*q2), 1-2*(q2**2+q3**2)) 
               for q0,q1,q2,q3 in zip(*q)]
        
        axes[0].plot(t, np.degrees(roll), label='Roll')
        axes[0].plot(t, np.degrees(pitch), label='Pitch')
        axes[0].plot(t, np.degrees(yaw), label='Yaw')
        axes[0].set_ylabel('Angle (deg)')
        axes[0].legend()
        axes[0].grid(True)
        axes[0].set_title('Attitude')
    
    # 2. Gyro rates
    if 'sensor_gyro' in topics:
        gyro = topics['sensor_gyro']
        t = gyro.data['timestamp'] - gyro.data['timestamp'][0]
        axes[1].plot(t, np.degrees(gyro.data['x']), label='Gyro X')
        axes[1].plot(t, np.degrees(gyro.data['y']), label='Gyro Y')
        axes[1].plot(t, np.degrees(gyro.data['z']), label='Gyro Z')
        axes[1].set_ylabel('Rate (deg/s)')
        axes[1].legend()
        axes[1].grid(True)
        axes[1].set_title('Gyroscope')
    
    # 3. Battery
    if 'battery_status' in topics:
        bat = topics['battery_status']
        t = bat.data['timestamp'] - bat.data['timestamp'][0]
        axes[2].plot(t, bat.data['voltage_v'], label='Voltage')
        axes[2].set_ylabel('Voltage (V)')
        axes[2].legend()
        axes[2].grid(True)
        axes[2].set_title('Battery')
        axes[2].axhline(y=3.5, color='r', linestyle='--', label='Low')
        axes[2].axhline(y=3.3, color='r', linestyle=':', label='Critical')
    
    # 4. Actuator outputs
    if 'actuator_outputs' in topics:
        act = topics['actuator_outputs']
        t = act.data['timestamp'] - act.data['timestamp'][0]
        # Plot first 4 motor outputs
        for i in range(min(4, len(act.data.get('output[0]', [])))):
            key = f'output[{i}]'
            if key in act.data:
                axes[3].plot(t, act.data[key], label=f'Motor {i+1}')
        axes[3].set_ylabel('PWM (us)')
        axes[3].set_xlabel('Time (s)')
        axes[3].legend()
        axes[3].grid(True)
        axes[3].set_title('Actuator Outputs')
    
    plt.tight_layout()
    plt.show()
    
    # Print summary stats
    print("\n--- Summary ---")
    if 'vehicle_attitude' in topics:
        pitch_deg = np.degrees(pitch)
        print(f"Max pitch: {max(pitch_deg):.1f}°")
        print(f"Max roll: {max(np.degrees(roll)):.1f}°")
    
    if 'battery_status' in topics:
        v = bat.data['voltage_v']
        print(f"Battery: {np.mean(v):.2f}V avg, {np.min(v):.2f}V min, {np.max(v):.2f}V max")
    
    if 'sensor_gyro' in topics:
        for axis, name in [(gyro.data['x'], 'X'), (gyro.data['y'], 'Y'), (gyro.data['z'], 'Z')]:
            rms = np.sqrt(np.mean(np.degrees(axis)**2))
            print(f"Gyro {name} RMS noise: {rms:.2f} °/s")


def main():
    parser = argparse.ArgumentParser(description="Drone Flight Log Analyzer")
    parser.add_argument("file", help="Flight log file (.ulg or .bin)")
    parser.add_argument("--plot", action="store_true", default=True,
                       help="Show plots")
    args = parser.parse_args()
    
    if args.file.endswith('.ulg'):
        analyze_ulog(args.file)
    elif args.file.endswith('.bin'):
        print(".bin logs require pymavlink (pip install pymavlink)")
        print("Convert to .ulg via: python3 Tools/ulog2csv.py log.bin")
    else:
        print(f"Unknown format: {args.file}")
        print("Supported: .ulg (PX4), .bin (ArduPilot)")


if __name__ == "__main__":
    main()