#!/usr/bin/env python3
"""
Swarm Localization Demo

This example demonstrates the cooperative acoustic localization
capabilities of the HADEX swarm navigation system.
"""

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from src.navigation.swarm_algorithm import SwarmNavigationEngine
from src.navigation.position_estimator import PositionEstimator
from src.navigation.adaptive_ranging import AdaptiveRanging
from src.core.acoustic_model import AcousticModel


def main():
    """Run swarm localization demo"""
    print("=" * 60)
    print("HADEX - Swarm Localization Demo")
    print("=" * 60)
    
    # 1. Initialize swarm
    print("\n1. Initializing swarm...")
    num_vehicles = 8
    swarm = SwarmNavigationEngine(
        num_vehicles=num_vehicles,
        acoustic_frequency=12000,
        max_range=500
    )
    print(f"   Swarm initialized with {num_vehicles} vehicles")
    
    # 2. Initialize acoustic model
    acoustic = AcousticModel(temperature=2.0, salinity=34.7, depth=5000)
    sound_velocity = acoustic.sound_speed()
    print(f"   Sound velocity: {sound_velocity:.1f} m/s")
    
    # 3. Run localization simulation
    print("\n2. Running localization simulation...")
    
    # Simulate for different durations
    durations = [1, 6, 12, 24, 48, 72]  # hours
    errors = []
    improvements = []
    
    for hours in durations:
        results = swarm.simulate_mission(
            duration_hours=hours,
            depth=5000,
            survey_area=500
        )
        errors.append(results['rms_error'])
        improvements.append(results['improvement_factor'])
        print(f"   {hours:2d} hours: RMS Error = {results['rms_error']:.2f} m, "
              f"Improvement = {results['improvement_factor']:.1f}x")
    
    # 4. Visualize trajectory
    print("\n3. Generating trajectory visualization...")
    
    # Run a short mission for visualization
    vis_results = swarm.simulate_mission(
        duration_hours=2,
        depth=5000,
        survey_area=500
    )
    
    # Extract trajectory
    if vis_results['history']:
        positions = np.array([h['positions'] for h in vis_results['history']])
        
        # Plot 3D trajectory
        fig = plt.figure(figsize=(14, 10))
        
        ax1 = fig.add_subplot(221, projection='3d')
        for i in range(min(4, positions.shape[1])):  # Plot first 4 vehicles
            traj = positions[:, i, :]
            ax1.plot(traj[:, 0], traj[:, 1], -traj[:, 2], 
                    linewidth=2, label=f'Vehicle {i}')
        
        ax1.set_xlabel('X (m)')
        ax1.set_ylabel('Y (m)')
        ax1.set_zlabel('Depth (m)')
        ax1.set_title('3D Swarm Trajectories')
        ax1.legend()
        
        # Plot XY projection
        ax2 = fig.add_subplot(222)
        for i in range(min(4, positions.shape[1])):
            traj = positions[:, i, :]
            ax2.plot(traj[:, 0], traj[:, 1], linewidth=2, label=f'Vehicle {i}')
        
        ax2.set_xlabel('X (m)')
        ax2.set_ylabel('Y (m)')
        ax2.set_title('Horizontal Trajectories')
        ax2.legend()
        ax2.grid(True, alpha=0.3)
        ax2.set_aspect('equal')
        
        # Plot error over time
        ax3 = fig.add_subplot(223)
        ax3.plot(durations, errors, 'bo-', linewidth=2, markersize=8)
        ax3.set_xlabel('Mission Duration (hours)')
        ax3.set_ylabel('RMS Error (m)')
        ax3.set_title('Position Error vs Duration')
        ax3.grid(True, alpha=0.3)
        
        # Plot improvement factor
        ax4 = fig.add_subplot(224)
        ax4.plot(durations, improvements, 'ro-', linewidth=2, markersize=8)
        ax4.axhline(y=40, color='g', linestyle='--', label='Target (40x)')
        ax4.set_xlabel('Mission Duration (hours)')
        ax4.set_ylabel('Improvement Factor')
        ax4.set_title('Improvement vs Dead Reckoning')
        ax4.legend()
        ax4.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig('swarm_localization.png', dpi=150)
        print("   Trajectory plot saved to: swarm_localization.png")
    
    # 5. Demonstrate cooperative positioning
    print("\n4. Cooperative Positioning Example")
    print("-" * 40)
    
    estimator = PositionEstimator()
    ranging = AdaptiveRanging()
    
    # Simulate a vehicle lost in the swarm
    print("\n   Vehicle 0 lost - using neighbors for positioning...")
    
    # Generate true position
    true_pos = np.array([125.3, 142.7, 5000])
    
    # Neighbor positions
    neighbors = {
        1: np.array([100, 100, 5000]),
        2: np.array([150, 100, 5000]),
        3: np.array([100, 150, 5000]),
        4: np.array([150, 150, 5000])
    }
    
    # Calculate true ranges
    true_ranges = {}
    for nid, npos in neighbors.items():
        dist = np.linalg.norm(true_pos - npos)
        true_ranges[nid] = dist
    
    # Add measurement noise
    measured_ranges = {}
    for nid, dist in true_ranges.items():
        # Add noise and multipath
        noisy_dist = dist + np.random.normal(0, 0.5)
        if np.random.random() < 0.1:  # 10% multipath
            noisy_dist += np.random.exponential(2.0)
        measured_ranges[nid] = noisy_dist
    
    print(f"\n   True position: ({true_pos[0]:.1f}, {true_pos[1]:.1f}, {true_pos[2]:.1f})")
    
    # Estimate position
    ranges_list = list(measured_ranges.values())
    anchors_list = list(neighbors.values())
    
    estimated_pos = estimator.solve(ranges_list, anchors_list)
    print(f"   Estimated position: ({estimated_pos[0]:.1f}, {estimated_pos[1]:.1f}, {estimated_pos[2]:.1f})")
    
    error = np.linalg.norm(true_pos - estimated_pos)
    print(f"   Positioning error: {error:.2f} m")
    
    # Apply Kalman filter
    filtered_pos, _ = estimator.kalman_update(
        measurement=estimated_pos,
        uncertainty=error,
        prior_position=true_pos + np.random.randn(3) * 2,
        prior_covariance=np.eye(3) * 4,
        dt=30.0
    )
    
    filter_error = np.linalg.norm(true_pos - filtered_pos)
    print(f"   After Kalman filter: {filter_error:.2f} m")
    print(f"   Improvement: {(1 - filter_error/error)*100:.1f}%")
    
    # 6. Show statistics
    print("\n5. Swarm Statistics")
    print("-" * 40)
    
    # Get final positions
    final_positions = swarm.get_positions()
    
    # Calculate inter-vehicle distances
    distances = []
    for i in range(num_vehicles):
        for j in range(i+1, num_vehicles):
            if i in final_positions and j in final_positions:
                dist = np.linalg.norm(final_positions[i] - final_positions[j])
                distances.append(dist)
    
    print(f"\n   Number of vehicles: {num_vehicles}")
    print(f"   Mean inter-vehicle distance: {np.mean(distances):.1f} m")
    print(f"   Min distance: {np.min(distances):.1f} m")
    print(f"   Max distance: {np.max(distances):.1f} m")
    print(f"   Current accuracy: {swarm.position_accuracy():.2f} m")
    
    print("\n" + "=" * 60)
    print("Swarm localization demo completed!")
    print("=" * 60)


if __name__ == "__main__":
    main()
