import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import linprog
import pandas as pd

# Problem Parameters
# Data types: 0=High-priority Science, 1=Medium-priority Telemetry, 2=Low-priority Housekeeping
data_types = ['High-priority Science', 'Medium-priority Telemetry', 'Low-priority Housekeeping']
num_data_types = 3

# Communication windows: 0=Morning, 1=Afternoon, 2=Evening, 3=Night
windows = ['Morning (2h)', 'Afternoon (3h)', 'Evening (2.5h)', 'Night (1.5h)']
num_windows = 4

# Data availability (MB)
data_available = np.array([500, 800, 1200])  # MB for each data type

# Priority weights (higher = more important)
priorities = np.array([10, 5, 1])

# Transmission rates (MB/hour) - varies by distance and ground station
transmission_rates = np.array([150, 180, 160, 140])  # MB/hour for each window

# Window durations (hours)
window_durations = np.array([2.0, 3.0, 2.5, 1.5])  # hours

# Energy consumption (Wh/MB) - varies by transmitter power settings
energy_per_mb = np.array([0.8, 0.7, 0.75, 0.9])  # Wh/MB

# Available power during each window (Wh)
power_available = np.array([600, 900, 700, 400])  # Wh

print("=" * 80)
print("SPACECRAFT DATA TRANSMISSION OPTIMIZATION")
print("=" * 80)
print("\n📡 Mission Parameters:")
print(f"   - Data Types: {num_data_types}")
print(f"   - Communication Windows: {num_windows}")
print(f"   - Planning Horizon: {sum(window_durations):.1f} hours")
print(f"   - Total Data to Transmit: {sum(data_available):.0f} MB")

# Setup Linear Programming Problem
# Decision variables: x[i,j] = data type i transmitted in window j
# Flattened: x[0,0], x[0,1], ..., x[0,3], x[1,0], ..., x[2,3]
num_vars = num_data_types * num_windows

# Objective function coefficients (negative because linprog minimizes)
c = np.zeros(num_vars)
for i in range(num_data_types):
    for j in range(num_windows):
        idx = i * num_windows + j
        c[idx] = -priorities[i]  # Negative for maximization

# Inequality constraints: A_ub @ x <= b_ub
A_ub = []
b_ub = []

# Constraint 1: Data availability for each data type
for i in range(num_data_types):
    constraint = np.zeros(num_vars)
    for j in range(num_windows):
        idx = i * num_windows + j
        constraint[idx] = 1
    A_ub.append(constraint)
    b_ub.append(data_available[i])

# Constraint 2: Window duration (time constraint)
for j in range(num_windows):
    constraint = np.zeros(num_vars)
    for i in range(num_data_types):
        idx = i * num_windows + j
        constraint[idx] = 1.0 / transmission_rates[j]
    A_ub.append(constraint)
    b_ub.append(window_durations[j])

# Constraint 3: Power constraint
for j in range(num_windows):
    constraint = np.zeros(num_vars)
    for i in range(num_data_types):
        idx = i * num_windows + j
        constraint[idx] = energy_per_mb[j]
    A_ub.append(constraint)
    b_ub.append(power_available[j])

A_ub = np.array(A_ub)
b_ub = np.array(b_ub)

# Variable bounds (non-negativity)
bounds = [(0, None) for _ in range(num_vars)]

# Solve the linear program
print("\n🔧 Solving optimization problem...")
result = linprog(c, A_ub=A_ub, b_ub=b_ub, bounds=bounds, method='highs')

if result.success:
    print("✅ Optimization successful!")
    
    # Reshape solution to matrix form
    x_optimal = result.x.reshape(num_data_types, num_windows)
    
    print(f"\n📊 Optimal Total Priority-Weighted Data: {-result.fun:.2f}")
    print(f"   Total Data Transmitted: {np.sum(x_optimal):.2f} MB")
    
    # Display solution
    print("\n" + "=" * 80)
    print("OPTIMAL TRANSMISSION SCHEDULE")
    print("=" * 80)
    
    df = pd.DataFrame(x_optimal, 
                      index=data_types, 
                      columns=windows)
    print("\nData Transmission (MB):")
    print(df.to_string())
    
    # Calculate utilization metrics
    print("\n" + "-" * 80)
    print("RESOURCE UTILIZATION")
    print("-" * 80)
    
    # Data type utilization
    print("\n1. Data Type Utilization:")
    for i in range(num_data_types):
        transmitted = np.sum(x_optimal[i, :])
        utilization = (transmitted / data_available[i]) * 100
        print(f"   {data_types[i]}: {transmitted:.2f} / {data_available[i]:.2f} MB ({utilization:.1f}%)")
    
    # Window utilization
    print("\n2. Communication Window Utilization:")
    for j in range(num_windows):
        data_in_window = np.sum(x_optimal[:, j])
        time_used = data_in_window / transmission_rates[j]
        time_utilization = (time_used / window_durations[j]) * 100
        
        power_used = data_in_window * energy_per_mb[j]
        power_utilization = (power_used / power_available[j]) * 100
        
        print(f"\n   {windows[j]}:")
        print(f"      Time: {time_used:.2f} / {window_durations[j]:.2f} hours ({time_utilization:.1f}%)")
        print(f"      Power: {power_used:.2f} / {power_available[j]:.2f} Wh ({power_utilization:.1f}%)")
        print(f"      Data: {data_in_window:.2f} MB")
    
else:
    print("❌ Optimization failed!")
    print(f"   Status: {result.message}")

# Visualization
print("\n📈 Generating visualizations...")

fig = plt.figure(figsize=(16, 12))

# Plot 1: Transmission schedule (stacked bar chart)
ax1 = plt.subplot(2, 3, 1)
bottom = np.zeros(num_windows)
colors = ['#FF6B6B', '#4ECDC4', '#45B7D1']

for i in range(num_data_types):
    ax1.bar(windows, x_optimal[i, :], bottom=bottom, 
            label=data_types[i], color=colors[i], alpha=0.8)
    bottom += x_optimal[i, :]

ax1.set_ylabel('Data Transmitted (MB)', fontsize=11, fontweight='bold')
ax1.set_xlabel('Communication Window', fontsize=11, fontweight='bold')
ax1.set_title('Optimal Transmission Schedule', fontsize=13, fontweight='bold', pad=15)
ax1.legend(loc='upper right', fontsize=9)
ax1.grid(axis='y', alpha=0.3, linestyle='--')
plt.setp(ax1.xaxis.get_majorticklabels(), rotation=15, ha='right')

# Plot 2: Data type utilization
ax2 = plt.subplot(2, 3, 2)
transmitted_by_type = np.sum(x_optimal, axis=1)
utilization_by_type = (transmitted_by_type / data_available) * 100

bars = ax2.barh(data_types, utilization_by_type, color=colors, alpha=0.8)
ax2.set_xlabel('Utilization (%)', fontsize=11, fontweight='bold')
ax2.set_title('Data Type Utilization', fontsize=13, fontweight='bold', pad=15)
ax2.set_xlim(0, 100)
ax2.grid(axis='x', alpha=0.3, linestyle='--')

for i, bar in enumerate(bars):
    width = bar.get_width()
    ax2.text(width + 2, bar.get_y() + bar.get_height()/2, 
             f'{width:.1f}%', va='center', fontsize=10, fontweight='bold')

# Plot 3: Time utilization by window
ax3 = plt.subplot(2, 3, 3)
time_used_per_window = np.sum(x_optimal, axis=0) / transmission_rates
time_utilization = (time_used_per_window / window_durations) * 100

bars = ax3.bar(windows, time_utilization, color='#95E1D3', alpha=0.8, edgecolor='black', linewidth=1.5)
ax3.set_ylabel('Time Utilization (%)', fontsize=11, fontweight='bold')
ax3.set_xlabel('Communication Window', fontsize=11, fontweight='bold')
ax3.set_title('Time Resource Utilization', fontsize=13, fontweight='bold', pad=15)
ax3.set_ylim(0, 100)
ax3.axhline(y=100, color='red', linestyle='--', linewidth=2, alpha=0.7, label='Maximum')
ax3.grid(axis='y', alpha=0.3, linestyle='--')
ax3.legend(fontsize=9)
plt.setp(ax3.xaxis.get_majorticklabels(), rotation=15, ha='right')

for i, bar in enumerate(bars):
    height = bar.get_height()
    ax3.text(bar.get_x() + bar.get_width()/2, height + 2, 
             f'{height:.1f}%', ha='center', fontsize=9, fontweight='bold')

# Plot 4: Power utilization by window
ax4 = plt.subplot(2, 3, 4)
power_used_per_window = np.sum(x_optimal, axis=0) * energy_per_mb
power_utilization = (power_used_per_window / power_available) * 100

bars = ax4.bar(windows, power_utilization, color='#F38181', alpha=0.8, edgecolor='black', linewidth=1.5)
ax4.set_ylabel('Power Utilization (%)', fontsize=11, fontweight='bold')
ax4.set_xlabel('Communication Window', fontsize=11, fontweight='bold')
ax4.set_title('Power Resource Utilization', fontsize=13, fontweight='bold', pad=15)
ax4.set_ylim(0, 100)
ax4.axhline(y=100, color='red', linestyle='--', linewidth=2, alpha=0.7, label='Maximum')
ax4.grid(axis='y', alpha=0.3, linestyle='--')
ax4.legend(fontsize=9)
plt.setp(ax4.xaxis.get_majorticklabels(), rotation=15, ha='right')

for i, bar in enumerate(bars):
    height = bar.get_height()
    ax4.text(bar.get_x() + bar.get_width()/2, height + 2, 
             f'{height:.1f}%', ha='center', fontsize=9, fontweight='bold')

# Plot 5: Priority-weighted contribution by data type
ax5 = plt.subplot(2, 3, 5)
priority_contribution = transmitted_by_type * priorities
total_priority = np.sum(priority_contribution)

wedges, texts, autotexts = ax5.pie(priority_contribution, labels=data_types, autopct='%1.1f%%',
                                     colors=colors, startangle=90, textprops={'fontsize': 10, 'fontweight': 'bold'})
ax5.set_title(f'Priority-Weighted Contribution\n(Total: {total_priority:.0f})', 
              fontsize=13, fontweight='bold', pad=15)

# Plot 6: Transmission rate analysis
ax6 = plt.subplot(2, 3, 6)
data_per_window = np.sum(x_optimal, axis=0)
effective_rate = data_per_window / window_durations

x_pos = np.arange(len(windows))
width = 0.35

bars1 = ax6.bar(x_pos - width/2, transmission_rates, width, 
                label='Max Rate', color='#A8E6CF', alpha=0.8, edgecolor='black')
bars2 = ax6.bar(x_pos + width/2, effective_rate, width, 
                label='Achieved Rate', color='#FFD3B6', alpha=0.8, edgecolor='black')

ax6.set_ylabel('Transmission Rate (MB/hour)', fontsize=11, fontweight='bold')
ax6.set_xlabel('Communication Window', fontsize=11, fontweight='bold')
ax6.set_title('Transmission Rate: Maximum vs Achieved', fontsize=13, fontweight='bold', pad=15)
ax6.set_xticks(x_pos)
ax6.set_xticklabels(windows)
ax6.legend(fontsize=9)
ax6.grid(axis='y', alpha=0.3, linestyle='--')
plt.setp(ax6.xaxis.get_majorticklabels(), rotation=15, ha='right')

plt.tight_layout()
plt.savefig('spacecraft_transmission_optimization.png', dpi=300, bbox_inches='tight')
plt.show()

print("\n✅ Visualization complete!")
print("=" * 80)

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import differential_evolution, minimize
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.patches import Circle
import matplotlib.patches as mpatches

# Physical constants
AU = 1.496e11  # Astronomical Unit in meters
G = 6.67430e-11  # Gravitational constant
M_sun = 1.989e30  # Solar mass in kg
M_jupiter = 1.898e27  # Jupiter mass in kg
M_saturn = 5.683e26  # Saturn mass in kg

# Orbital parameters (simplified circular orbits)
class CelestialBody:
    def __init__(self, name, semi_major_axis, period, mu_parent):
        self.name = name
        self.a = semi_major_axis  # meters
        self.T = period  # seconds
        self.mu = mu_parent  # gravitational parameter
        self.n = 2 * np.pi / period  # mean motion
    
    def position_at_time(self, t):
        """Calculate position at time t (circular orbit approximation)"""
        theta = self.n * t
        x = self.a * np.cos(theta)
        y = self.a * np.sin(theta)
        z = 0
        return np.array([x, y, z])
    
    def velocity_at_time(self, t):
        """Calculate velocity at time t"""
        theta = self.n * t
        v = np.sqrt(self.mu / self.a)
        vx = -v * np.sin(theta)
        vy = v * np.cos(theta)
        vz = 0
        return np.array([vx, vy, vz])

# Define celestial bodies
# Jupiter (simplified as orbiting the Sun)
jupiter = CelestialBody(
    "Jupiter",
    5.2 * AU,
    11.86 * 365.25 * 24 * 3600,  # ~11.86 years
    G * M_sun
)

# Saturn
saturn = CelestialBody(
    "Saturn",
    9.5 * AU,
    29.46 * 365.25 * 24 * 3600,  # ~29.46 years
    G * M_sun
)

# Europa (orbiting Jupiter)
europa_orbit = 671100e3  # meters from Jupiter
europa_period = 3.551 * 24 * 3600  # seconds

# Enceladus (orbiting Saturn)
enceladus_orbit = 238000e3  # meters from Saturn
enceladus_period = 1.370 * 24 * 3600  # seconds

def lambert_solver_simplified(r1, r2, tof, mu):
    """
    Simplified Lambert solver using the universal variable formulation.
    Returns initial and final velocity vectors.
    """
    r1_mag = np.linalg.norm(r1)
    r2_mag = np.linalg.norm(r2)
    
    cos_dtheta = np.dot(r1, r2) / (r1_mag * r2_mag)
    cos_dtheta = np.clip(cos_dtheta, -1, 1)
    
    # Determine short or long way
    cross_prod = np.cross(r1, r2)
    if cross_prod[2] >= 0:
        sin_dtheta = np.linalg.norm(cross_prod) / (r1_mag * r2_mag)
    else:
        sin_dtheta = -np.linalg.norm(cross_prod) / (r1_mag * r2_mag)
    
    A = np.sqrt(r1_mag * r2_mag * (1 + cos_dtheta))
    
    if A == 0:
        return None, None
    
    # Iterative solution for z
    c2 = 0.5
    c3 = 1.0 / 6.0
    z = 0.0
    
    for _ in range(100):
        y = r1_mag + r2_mag + A * (z * c3 - 1) / np.sqrt(c2)
        
        if y < 0:
            z = z - 0.1
            continue
        
        chi = np.sqrt(y / c2)
        tof_new = (chi**3 * c3 + A * np.sqrt(y)) / np.sqrt(mu)
        
        if abs(tof_new - tof) < 1e-6:
            break
        
        if tof_new <= tof:
            z = z + 0.1
        else:
            z = z - 0.1
        
        # Update Stumpff functions
        if z > 1e-6:
            c2 = (1 - np.cos(np.sqrt(z))) / z
            c3 = (np.sqrt(z) - np.sin(np.sqrt(z))) / np.sqrt(z**3)
        elif z < -1e-6:
            c2 = (1 - np.cosh(np.sqrt(-z))) / z
            c3 = (np.sinh(np.sqrt(-z)) - np.sqrt(-z)) / np.sqrt((-z)**3)
        else:
            c2 = 0.5
            c3 = 1.0 / 6.0
    
    # Calculate velocities
    f = 1 - y / r1_mag
    g = A * np.sqrt(y / mu)
    g_dot = 1 - y / r2_mag
    
    v1 = (r2 - f * r1) / g
    v2 = (g_dot * r2 - r1) / g
    
    return v1, v2

def calculate_mission_deltav(params, bodies, return_details=False):
    """
    Calculate total delta-v for a mission.
    params: [t_departure, t_jupiter_arrival, t_saturn_arrival]
    """
    t_dep, t_jup, t_sat = params
    
    # Ensure proper time ordering
    if t_jup <= t_dep or t_sat <= t_jup:
        return 1e10 if not return_details else (1e10, None)
    
    # Earth starting position (simplified at 1 AU)
    r_earth = np.array([AU, 0, 0])
    v_earth = np.array([0, np.sqrt(G * M_sun / AU), 0])
    
    # Jupiter position at arrival
    r_jupiter = bodies['jupiter'].position_at_time(t_jup)
    v_jupiter = bodies['jupiter'].velocity_at_time(t_jup)
    
    # Saturn position at arrival
    r_saturn = bodies['saturn'].position_at_time(t_sat)
    v_saturn = bodies['saturn'].velocity_at_time(t_sat)
    
    # Leg 1: Earth to Jupiter
    tof1 = t_jup - t_dep
    v1_dep, v1_arr = lambert_solver_simplified(r_earth, r_jupiter, tof1, G * M_sun)
    
    if v1_dep is None:
        return 1e10 if not return_details else (1e10, None)
    
    dv_departure = np.linalg.norm(v1_dep - v_earth)
    dv_jupiter = np.linalg.norm(v1_arr - v_jupiter)
    
    # Leg 2: Jupiter to Saturn (with gravity assist)
    tof2 = t_sat - t_jup
    v2_dep, v2_arr = lambert_solver_simplified(r_jupiter, r_saturn, tof2, G * M_sun)
    
    if v2_dep is None:
        return 1e10 if not return_details else (1e10, None)
    
    # Gravity assist reduces required delta-v
    v_inf_in = v1_arr - v_jupiter
    v_inf_out = v2_dep - v_jupiter
    
    # Simplified gravity assist calculation
    ga_factor = 0.7  # Realistic gravity assist efficiency
    dv_jupiter_corrected = np.linalg.norm(v_inf_out - v_inf_in) * (1 - ga_factor)
    
    dv_saturn = np.linalg.norm(v2_arr - v_saturn)
    
    total_dv = dv_departure + dv_jupiter_corrected + dv_saturn
    
    if return_details:
        details = {
            'dv_departure': dv_departure,
            'dv_jupiter': dv_jupiter_corrected,
            'dv_saturn': dv_saturn,
            'positions': {
                'earth': r_earth,
                'jupiter': r_jupiter,
                'saturn': r_saturn
            },
            'times': {
                't_dep': t_dep,
                't_jup': t_jup,
                't_sat': t_sat
            },
            'trajectories': {
                'leg1': (r_earth, r_jupiter, v1_dep, v1_arr),
                'leg2': (r_jupiter, r_saturn, v2_dep, v2_arr)
            }
        }
        return total_dv, details
    
    return total_dv

# Set up optimization
bodies = {
    'jupiter': jupiter,
    'saturn': saturn
}

# Time bounds (in seconds from mission start)
year = 365.25 * 24 * 3600
bounds = [
    (0, 0.1 * year),  # Departure time
    (1.5 * year, 3 * year),  # Jupiter arrival
    (5 * year, 8 * year)  # Saturn arrival
]

print("="*70)
print("MULTI-BODY TRAJECTORY OPTIMIZATION")
print("Mission: Earth → Europa (Jupiter) → Enceladus (Saturn)")
print("="*70)
print("\nStarting optimization...")
print("This may take a minute...\n")

# Optimize using differential evolution
result = differential_evolution(
    lambda x: calculate_mission_deltav(x, bodies),
    bounds,
    maxiter=100,
    popsize=15,
    seed=42,
    atol=1e-6,
    tol=1e-6
)

optimal_params = result.x
optimal_dv, details = calculate_mission_deltav(optimal_params, bodies, return_details=True)

# Convert to human-readable units
print("\n" + "="*70)
print("OPTIMIZATION RESULTS")
print("="*70)
print(f"\nOptimal Total Δv: {optimal_dv/1000:.2f} km/s")
print(f"\nΔv Breakdown:")
print(f"  • Earth Departure:     {details['dv_departure']/1000:.2f} km/s")
print(f"  • Jupiter Gravity Assist: {details['dv_jupiter']/1000:.2f} km/s")
print(f"  • Saturn Arrival:      {details['dv_saturn']/1000:.2f} km/s")

print(f"\nMission Timeline:")
print(f"  • Departure:           t = 0 days")
print(f"  • Jupiter Arrival:     t = {details['times']['t_jup']/(24*3600):.1f} days ({details['times']['t_jup']/(365.25*24*3600):.2f} years)")
print(f"  • Saturn Arrival:      t = {details['times']['t_sat']/(24*3600):.1f} days ({details['times']['t_sat']/(365.25*24*3600):.2f} years)")
print(f"  • Total Mission Duration: {details['times']['t_sat']/(365.25*24*3600):.2f} years")
print("="*70)

# Generate trajectory for visualization
def generate_trajectory_points(r1, r2, v1, tof, mu, n_points=100):
    """Generate points along the trajectory"""
    points = []
    for i in range(n_points):
        t = tof * i / (n_points - 1)
        # Simple propagation
        r = r1 + v1 * t
        points.append(r)
    return np.array(points)

# Create visualizations
fig = plt.figure(figsize=(18, 12))

# 3D trajectory plot
ax1 = fig.add_subplot(2, 3, 1, projection='3d')
ax1.set_title('3D Trajectory Overview', fontsize=14, fontweight='bold')

# Plot orbits
theta = np.linspace(0, 2*np.pi, 100)
jupiter_orbit_x = jupiter.a * np.cos(theta) / AU
jupiter_orbit_y = jupiter.a * np.sin(theta) / AU
saturn_orbit_x = saturn.a * np.cos(theta) / AU
saturn_orbit_y = saturn.a * np.sin(theta) / AU

ax1.plot(jupiter_orbit_x, jupiter_orbit_y, 0, 'orange', alpha=0.3, linestyle='--', label='Jupiter Orbit')
ax1.plot(saturn_orbit_x, saturn_orbit_y, 0, 'gold', alpha=0.3, linestyle='--', label='Saturn Orbit')

# Plot trajectories
leg1_r1, leg1_r2, leg1_v1, leg1_v2 = details['trajectories']['leg1']
leg2_r1, leg2_r2, leg2_v1, leg2_v2 = details['trajectories']['leg2']

traj1 = generate_trajectory_points(leg1_r1, leg1_r2, leg1_v1, 
                                   details['times']['t_jup'] - details['times']['t_dep'], 
                                   G * M_sun)
traj2 = generate_trajectory_points(leg2_r1, leg2_r2, leg2_v1,
                                   details['times']['t_sat'] - details['times']['t_jup'],
                                   G * M_sun)

ax1.plot(traj1[:, 0]/AU, traj1[:, 1]/AU, traj1[:, 2]/AU, 'b-', linewidth=2, label='Earth to Jupiter')
ax1.plot(traj2[:, 0]/AU, traj2[:, 1]/AU, traj2[:, 2]/AU, 'r-', linewidth=2, label='Jupiter to Saturn')

# Plot bodies
ax1.scatter([1], [0], [0], c='blue', s=100, marker='o', label='Earth (Start)', edgecolors='black', linewidths=1.5)
ax1.scatter([details['positions']['jupiter'][0]/AU], 
           [details['positions']['jupiter'][1]/AU], 
           [0], c='orange', s=200, marker='o', label='Jupiter', edgecolors='black', linewidths=1.5)
ax1.scatter([details['positions']['saturn'][0]/AU], 
           [details['positions']['saturn'][1]/AU], 
           [0], c='gold', s=200, marker='o', label='Saturn', edgecolors='black', linewidths=1.5)
ax1.scatter([0], [0], [0], c='yellow', s=300, marker='*', label='Sun', edgecolors='orange', linewidths=2)

ax1.set_xlabel('X (AU)', fontsize=10)
ax1.set_ylabel('Y (AU)', fontsize=10)
ax1.set_zlabel('Z (AU)', fontsize=10)
ax1.legend(fontsize=8, loc='upper right')
ax1.grid(True, alpha=0.3)

# 2D Top view
ax2 = fig.add_subplot(2, 3, 2)
ax2.set_title('Top View (X-Y Plane)', fontsize=14, fontweight='bold')
ax2.plot(jupiter_orbit_x, jupiter_orbit_y, 'orange', alpha=0.3, linestyle='--', linewidth=1)
ax2.plot(saturn_orbit_x, saturn_orbit_y, 'gold', alpha=0.3, linestyle='--', linewidth=1)
ax2.plot(traj1[:, 0]/AU, traj1[:, 1]/AU, 'b-', linewidth=2, label='Leg 1')
ax2.plot(traj2[:, 0]/AU, traj2[:, 1]/AU, 'r-', linewidth=2, label='Leg 2')
ax2.scatter([1], [0], c='blue', s=100, marker='o', edgecolors='black', linewidths=1.5, zorder=5)
ax2.scatter([details['positions']['jupiter'][0]/AU], 
           [details['positions']['jupiter'][1]/AU], 
           c='orange', s=200, marker='o', edgecolors='black', linewidths=1.5, zorder=5)
ax2.scatter([details['positions']['saturn'][0]/AU], 
           [details['positions']['saturn'][1]/AU], 
           c='gold', s=200, marker='o', edgecolors='black', linewidths=1.5, zorder=5)
ax2.scatter([0], [0], c='yellow', s=300, marker='*', edgecolors='orange', linewidths=2, zorder=5)
ax2.set_xlabel('X (AU)', fontsize=10)
ax2.set_ylabel('Y (AU)', fontsize=10)
ax2.axis('equal')
ax2.grid(True, alpha=0.3)
ax2.legend(fontsize=9)

# Delta-v breakdown pie chart
ax3 = fig.add_subplot(2, 3, 3)
ax3.set_title('Δv Budget Breakdown', fontsize=14, fontweight='bold')
dv_values = [details['dv_departure']/1000, details['dv_jupiter']/1000, details['dv_saturn']/1000]
labels = ['Earth Departure\n{:.2f} km/s'.format(dv_values[0]),
          'Jupiter GA\n{:.2f} km/s'.format(dv_values[1]),
          'Saturn Arrival\n{:.2f} km/s'.format(dv_values[2])]
colors = ['#3498db', '#e74c3c', '#f39c12']
explode = (0.05, 0.05, 0.05)
ax3.pie(dv_values, labels=labels, colors=colors, autopct='%1.1f%%', 
        explode=explode, startangle=90, textprops={'fontsize': 10})

# Mission timeline
ax4 = fig.add_subplot(2, 3, 4)
ax4.set_title('Mission Timeline', fontsize=14, fontweight='bold')
times = [0, details['times']['t_jup']/(365.25*24*3600), details['times']['t_sat']/(365.25*24*3600)]
events = ['Earth\nDeparture', 'Jupiter\nFlyby', 'Saturn\nArrival']
colors_timeline = ['blue', 'orange', 'gold']

ax4.barh(events, times, color=colors_timeline, edgecolor='black', linewidth=1.5)
ax4.set_xlabel('Time (years)', fontsize=10)
ax4.set_xlim(0, max(times) * 1.1)
for i, (event, time) in enumerate(zip(events, times)):
    ax4.text(time + 0.1, i, f'{time:.2f} yr', va='center', fontsize=10, fontweight='bold')
ax4.grid(True, alpha=0.3, axis='x')

# Velocity profile along trajectory
ax5 = fig.add_subplot(2, 3, 5)
ax5.set_title('Velocity Profile Along Trajectory', fontsize=14, fontweight='bold')

# Calculate velocity magnitudes
t1_points = np.linspace(0, details['times']['t_jup'] - details['times']['t_dep'], 50)
t2_points = np.linspace(details['times']['t_jup'], details['times']['t_sat'], 50)

v1_mag = np.linalg.norm(leg1_v1)
v2_mag = np.linalg.norm(leg2_v1)

t_total = np.concatenate([t1_points, t2_points]) / (365.25 * 24 * 3600)
v_profile = np.concatenate([
    np.full(len(t1_points), v1_mag/1000),
    np.full(len(t2_points), v2_mag/1000)
])

ax5.plot(t_total, v_profile, linewidth=2, color='purple')
ax5.axvline(details['times']['t_jup']/(365.25*24*3600), color='orange', 
           linestyle='--', linewidth=2, label='Jupiter Flyby')
ax5.set_xlabel('Mission Time (years)', fontsize=10)
ax5.set_ylabel('Velocity (km/s)', fontsize=10)
ax5.grid(True, alpha=0.3)
ax5.legend(fontsize=9)

# Distance from Sun over time
ax6 = fig.add_subplot(2, 3, 6)
ax6.set_title('Distance from Sun vs Time', fontsize=14, fontweight='bold')

r1_mag = np.linalg.norm(traj1, axis=1) / AU
r2_mag = np.linalg.norm(traj2, axis=1) / AU
t1 = np.linspace(0, details['times']['t_jup']/(365.25*24*3600), len(r1_mag))
t2 = np.linspace(details['times']['t_jup']/(365.25*24*3600), 
                details['times']['t_sat']/(365.25*24*3600), len(r2_mag))

ax6.plot(t1, r1_mag, 'b-', linewidth=2, label='Leg 1')
ax6.plot(t2, r2_mag, 'r-', linewidth=2, label='Leg 2')
ax6.axhline(jupiter.a/AU, color='orange', linestyle='--', alpha=0.5, label='Jupiter Orbit')
ax6.axhline(saturn.a/AU, color='gold', linestyle='--', alpha=0.5, label='Saturn Orbit')
ax6.set_xlabel('Mission Time (years)', fontsize=10)
ax6.set_ylabel('Distance from Sun (AU)', fontsize=10)
ax6.grid(True, alpha=0.3)
ax6.legend(fontsize=9)

plt.tight_layout()
plt.savefig('trajectory_optimization.png', dpi=300, bbox_inches='tight')
plt.show()

print("\n✓ Visualization saved as 'trajectory_optimization.png'")
print("\n" + "="*70)
print("ANALYSIS COMPLETE")
print("="*70)

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize
from matplotlib import rcParams

# Configure matplotlib for better-looking plots
rcParams['figure.figsize'] = (14, 10)
rcParams['font.size'] = 11

# Define biosignature molecules and their properties
molecules = {
    'O₂': {'weight': 1.0, 'alpha': 10.0, 'beta': 2.0, 'wavelength': 0.76, 'color': '#FF6B6B'},
    'CH₄': {'weight': 0.9, 'alpha': 8.5, 'beta': 1.8, 'wavelength': 3.3, 'color': '#4ECDC4'},
    'H₂O': {'weight': 0.85, 'alpha': 9.0, 'beta': 2.2, 'wavelength': 2.7, 'color': '#45B7D1'},
    'CO₂': {'weight': 0.7, 'alpha': 7.0, 'beta': 1.5, 'wavelength': 4.3, 'color': '#96CEB4'},
    'O₃': {'weight': 0.95, 'alpha': 9.5, 'beta': 1.9, 'wavelength': 9.6, 'color': '#FFEAA7'}
}

# Total available bandwidth (in micrometers)
B_total = 5.0

# Sensitivity function: S_i(b) = alpha * log(1 + beta * b)
def sensitivity(bandwidth, alpha, beta):
    """Calculate detection sensitivity for given bandwidth allocation"""
    return alpha * np.log(1 + beta * bandwidth)

# Objective function to MINIMIZE (negative of total weighted sensitivity)
def objective(bandwidths):
    """
    Calculate negative total weighted sensitivity (for minimization)
    bandwidths: array of bandwidth allocations for each molecule
    """
    total_sensitivity = 0
    for i, (mol_name, props) in enumerate(molecules.items()):
        s = sensitivity(bandwidths[i], props['alpha'], props['beta'])
        total_sensitivity += props['weight'] * s
    return -total_sensitivity  # Negative because we minimize

# Constraint: sum of bandwidths <= B_total
def constraint_total_bandwidth(bandwidths):
    """Ensure total bandwidth doesn't exceed available budget"""
    return B_total - np.sum(bandwidths)

# Initial guess: equal distribution
n_molecules = len(molecules)
initial_guess = np.ones(n_molecules) * (B_total / n_molecules)

# Bounds: each bandwidth must be non-negative
bounds = [(0, B_total) for _ in range(n_molecules)]

# Constraint dictionary for scipy.optimize
constraints = {'type': 'ineq', 'fun': constraint_total_bandwidth}

# Solve the optimization problem
print("="*70)
print("OPTIMIZING WAVELENGTH BAND ALLOCATION FOR BIOSIGNATURE DETECTION")
print("="*70)
print(f"\nTotal Available Bandwidth: {B_total} μm")
print(f"Number of Target Molecules: {n_molecules}")
print("\nMolecule Properties:")
print(f"{'Molecule':<10} {'Weight':<10} {'Alpha':<10} {'Beta':<10} {'λ_center (μm)':<15}")
print("-"*70)
for mol_name, props in molecules.items():
    print(f"{mol_name:<10} {props['weight']:<10.2f} {props['alpha']:<10.1f} "
          f"{props['beta']:<10.1f} {props['wavelength']:<15.2f}")

print("\n" + "="*70)
print("OPTIMIZATION IN PROGRESS...")
print("="*70)

result = minimize(
    objective,
    initial_guess,
    method='SLSQP',
    bounds=bounds,
    constraints=constraints,
    options={'ftol': 1e-9, 'disp': False}
)

# Extract optimal bandwidths
optimal_bandwidths = result.x
optimal_total_sensitivity = -result.fun

print("\n✓ Optimization completed successfully!\n")
print("="*70)
print("OPTIMAL BANDWIDTH ALLOCATION")
print("="*70)
print(f"{'Molecule':<10} {'Allocated BW (μm)':<20} {'Sensitivity':<15} {'Contribution':<15}")
print("-"*70)

sensitivities = []
contributions = []
for i, (mol_name, props) in enumerate(molecules.items()):
    bw = optimal_bandwidths[i]
    sens = sensitivity(bw, props['alpha'], props['beta'])
    contrib = props['weight'] * sens
    sensitivities.append(sens)
    contributions.append(contrib)
    print(f"{mol_name:<10} {bw:<20.4f} {sens:<15.2f} {contrib:<15.2f}")

print("-"*70)
print(f"{'TOTAL':<10} {np.sum(optimal_bandwidths):<20.4f} {'':<15} {optimal_total_sensitivity:<15.2f}")
print("="*70)

# Compare with naive equal allocation
equal_bandwidths = np.ones(n_molecules) * (B_total / n_molecules)
equal_sensitivity = -objective(equal_bandwidths)

print(f"\nComparison:")
print(f"  Equal Allocation Strategy:   Total Sensitivity = {equal_sensitivity:.2f}")
print(f"  Optimized Allocation:        Total Sensitivity = {optimal_total_sensitivity:.2f}")
print(f"  Improvement:                 {((optimal_total_sensitivity/equal_sensitivity - 1)*100):.2f}%")
print("="*70)

# Visualization
fig = plt.figure(figsize=(16, 12))

# Plot 1: Bandwidth Allocation Comparison
ax1 = plt.subplot(2, 3, 1)
x = np.arange(n_molecules)
width = 0.35
mol_names = list(molecules.keys())
colors = [molecules[m]['color'] for m in mol_names]

bars1 = ax1.bar(x - width/2, equal_bandwidths, width, label='Equal Allocation', 
                alpha=0.7, color='gray', edgecolor='black', linewidth=1.5)
bars2 = ax1.bar(x + width/2, optimal_bandwidths, width, label='Optimal Allocation',
                alpha=0.8, color=colors, edgecolor='black', linewidth=1.5)

ax1.set_xlabel('Molecule', fontweight='bold', fontsize=12)
ax1.set_ylabel('Allocated Bandwidth (μm)', fontweight='bold', fontsize=12)
ax1.set_title('Bandwidth Allocation: Equal vs Optimal', fontweight='bold', fontsize=13)
ax1.set_xticks(x)
ax1.set_xticklabels(mol_names)
ax1.legend(fontsize=10)
ax1.grid(axis='y', alpha=0.3)
ax1.axhline(y=B_total/n_molecules, color='red', linestyle='--', alpha=0.5, label='Average')

# Plot 2: Sensitivity Curves
ax2 = plt.subplot(2, 3, 2)
bw_range = np.linspace(0, B_total, 200)
for i, (mol_name, props) in enumerate(molecules.items()):
    sens_curve = sensitivity(bw_range, props['alpha'], props['beta'])
    ax2.plot(bw_range, sens_curve, label=mol_name, linewidth=2.5, 
             color=props['color'], alpha=0.8)
    # Mark optimal point
    opt_sens = sensitivity(optimal_bandwidths[i], props['alpha'], props['beta'])
    ax2.scatter(optimal_bandwidths[i], opt_sens, s=150, color=props['color'], 
                edgecolor='black', linewidth=2, zorder=5, marker='*')

ax2.set_xlabel('Allocated Bandwidth (μm)', fontweight='bold', fontsize=12)
ax2.set_ylabel('Detection Sensitivity', fontweight='bold', fontsize=12)
ax2.set_title('Sensitivity Functions with Optimal Points', fontweight='bold', fontsize=13)
ax2.legend(fontsize=9, loc='lower right')
ax2.grid(True, alpha=0.3)

# Plot 3: Contribution to Total Detection Score
ax3 = plt.subplot(2, 3, 3)
ax3.barh(mol_names, contributions, color=colors, alpha=0.8, edgecolor='black', linewidth=1.5)
ax3.set_xlabel('Weighted Contribution', fontweight='bold', fontsize=12)
ax3.set_title('Contribution to Total Detection Score', fontweight='bold', fontsize=13)
ax3.grid(axis='x', alpha=0.3)
for i, v in enumerate(contributions):
    ax3.text(v + 0.1, i, f'{v:.2f}', va='center', fontweight='bold')

# Plot 4: Pie Chart - Bandwidth Distribution
ax4 = plt.subplot(2, 3, 4)
wedges, texts, autotexts = ax4.pie(optimal_bandwidths, labels=mol_names, autopct='%1.1f%%',
                                     colors=colors, startangle=90, 
                                     wedgeprops={'edgecolor': 'black', 'linewidth': 1.5})
for autotext in autotexts:
    autotext.set_color('white')
    autotext.set_fontweight('bold')
    autotext.set_fontsize(10)
ax4.set_title('Optimal Bandwidth Distribution', fontweight='bold', fontsize=13)

# Plot 5: Weight vs Allocated Bandwidth
ax5 = plt.subplot(2, 3, 5)
weights = [molecules[m]['weight'] for m in mol_names]
ax5.scatter(weights, optimal_bandwidths, s=300, c=colors, alpha=0.7, 
            edgecolor='black', linewidth=2)
for i, mol in enumerate(mol_names):
    ax5.annotate(mol, (weights[i], optimal_bandwidths[i]), 
                xytext=(5, 5), textcoords='offset points', fontweight='bold')
ax5.set_xlabel('Importance Weight', fontweight='bold', fontsize=12)
ax5.set_ylabel('Allocated Bandwidth (μm)', fontweight='bold', fontsize=12)
ax5.set_title('Weight vs Allocation Relationship', fontweight='bold', fontsize=13)
ax5.grid(True, alpha=0.3)

# Plot 6: Sensitivity per Unit Bandwidth (Efficiency)
ax6 = plt.subplot(2, 3, 6)
efficiency = np.array(sensitivities) / (optimal_bandwidths + 1e-10)
ax6.bar(mol_names, efficiency, color=colors, alpha=0.8, edgecolor='black', linewidth=1.5)
ax6.set_ylabel('Sensitivity / Bandwidth', fontweight='bold', fontsize=12)
ax6.set_title('Detection Efficiency (Sensitivity per μm)', fontweight='bold', fontsize=13)
ax6.grid(axis='y', alpha=0.3)
for i, v in enumerate(efficiency):
    ax6.text(i, v + 0.1, f'{v:.2f}', ha='center', fontweight='bold')

plt.tight_layout()
plt.savefig('biosignature_optimization.png', dpi=300, bbox_inches='tight')
plt.show()

print("\n📊 Visualization saved as 'biosignature_optimization.png'")
print("\n" + "="*70)
print("ANALYSIS COMPLETE")
print("="*70)

molecules = {
    'O₂': {'weight': 1.0, 'alpha': 10.0, 'beta': 2.0, 'wavelength': 0.76, 'color': '#FF6B6B'},
    ...
}

def sensitivity(bandwidth, alpha, beta):
    return alpha * np.log(1 + beta * bandwidth)

def objective(bandwidths):
    total_sensitivity = 0
    for i, (mol_name, props) in enumerate(molecules.items()):
        s = sensitivity(bandwidths[i], props['alpha'], props['beta'])
        total_sensitivity += props['weight'] * s
    return -total_sensitivity

def constraint_total_bandwidth(bandwidths):
    return B_total - np.sum(bandwidths)

result = minimize(
    objective,
    initial_guess,
    method='SLSQP',  # Sequential Least Squares Programming
    bounds=bounds,
    constraints=constraints,
    options={'ftol': 1e-9}
)

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from scipy.spatial.distance import euclidean
import random
from typing import List, Tuple
import copy

# Set random seed for reproducibility
np.random.seed(42)
random.seed(42)

class RoverPathOptimizer:
    """
    Optimizes rover exploration routes using genetic algorithms
    to maximize scientific value while minimizing fuel and time.
    """
    
    def __init__(self, start_pos: Tuple[float, float], 
                 poi_positions: List[Tuple[float, float]],
                 poi_values: List[float],
                 max_distance: float,
                 max_time: float,
                 sampling_time: float = 2.0):
        """
        Initialize the optimizer.
        
        Parameters:
        - start_pos: Starting position (x, y)
        - poi_positions: List of points of interest coordinates
        - poi_values: Scientific value of each POI (0-100)
        - max_distance: Maximum travel distance (fuel constraint)
        - max_time: Maximum mission time
        - sampling_time: Time required to sample each POI
        """
        self.start = start_pos
        self.pois = poi_positions
        self.values = np.array(poi_values)
        self.max_dist = max_distance
        self.max_time = max_time
        self.sampling_time = sampling_time
        self.n_pois = len(poi_positions)
        self.rover_speed = 1.0  # units per time unit
        
    def calculate_distance(self, route: List[int]) -> float:
        """Calculate total distance for a given route."""
        if len(route) == 0:
            return 0.0
        
        total_dist = euclidean(self.start, self.pois[route[0]])
        for i in range(len(route) - 1):
            total_dist += euclidean(self.pois[route[i]], self.pois[route[i+1]])
        return total_dist
    
    def calculate_time(self, route: List[int]) -> float:
        """Calculate total mission time."""
        travel_time = self.calculate_distance(route) / self.rover_speed
        sampling_time = len(route) * self.sampling_time
        return travel_time + sampling_time
    
    def calculate_value(self, route: List[int]) -> float:
        """Calculate total scientific value."""
        if len(route) == 0:
            return 0.0
        return np.sum(self.values[route])
    
    def is_feasible(self, route: List[int]) -> bool:
        """Check if route satisfies constraints."""
        distance = self.calculate_distance(route)
        time = self.calculate_time(route)
        return distance <= self.max_dist and time <= self.max_time
    
    def fitness(self, route: List[int]) -> float:
        """
        Calculate fitness score (higher is better).
        Combines scientific value with penalties for constraint violations.
        """
        if len(route) == 0:
            return 0.0
        
        value = self.calculate_value(route)
        distance = self.calculate_distance(route)
        time = self.calculate_time(route)
        
        # Penalty for constraint violations
        distance_penalty = max(0, distance - self.max_dist) * 100
        time_penalty = max(0, time - self.max_time) * 100
        
        # Fitness = value - normalized_distance - penalties
        fitness = value - (distance / self.max_dist) * 20 - distance_penalty - time_penalty
        
        return fitness
    
    def generate_random_route(self) -> List[int]:
        """Generate a random feasible route."""
        all_indices = list(range(self.n_pois))
        random.shuffle(all_indices)
        
        # Start with empty route and add POIs until constraints violated
        route = []
        for idx in all_indices:
            test_route = route + [idx]
            if self.is_feasible(test_route):
                route = test_route
            else:
                break
        
        return route
    
    def mutate(self, route: List[int]) -> List[int]:
        """Apply mutation to a route."""
        if len(route) < 2:
            return route
        
        new_route = route.copy()
        mutation_type = random.choice(['swap', 'reverse', 'insert', 'remove'])
        
        if mutation_type == 'swap' and len(new_route) >= 2:
            i, j = random.sample(range(len(new_route)), 2)
            new_route[i], new_route[j] = new_route[j], new_route[i]
        
        elif mutation_type == 'reverse' and len(new_route) >= 2:
            i, j = sorted(random.sample(range(len(new_route)), 2))
            new_route[i:j+1] = reversed(new_route[i:j+1])
        
        elif mutation_type == 'insert':
            # Try to insert a new POI
            available = [i for i in range(self.n_pois) if i not in new_route]
            if available:
                new_poi = random.choice(available)
                pos = random.randint(0, len(new_route))
                new_route.insert(pos, new_poi)
        
        elif mutation_type == 'remove' and len(new_route) > 0:
            pos = random.randint(0, len(new_route) - 1)
            new_route.pop(pos)
        
        return new_route if self.is_feasible(new_route) else route
    
    def crossover(self, parent1: List[int], parent2: List[int]) -> List[int]:
        """Perform ordered crossover."""
        if len(parent1) == 0 or len(parent2) == 0:
            return parent1 if len(parent1) > len(parent2) else parent2
        
        # Take a subset from parent1
        size = min(len(parent1), len(parent2))
        start = random.randint(0, size - 1)
        end = random.randint(start + 1, size)
        
        child = parent1[start:end]
        
        # Add elements from parent2 that are not in child
        for gene in parent2:
            if gene not in child:
                child.append(gene)
        
        return child if self.is_feasible(child) else parent1
    
    def genetic_algorithm(self, population_size: int = 100, 
                         generations: int = 200,
                         mutation_rate: float = 0.2) -> Tuple[List[int], List[float]]:
        """
        Run genetic algorithm to find optimal route.
        
        Returns:
        - best_route: Optimal route found
        - history: Fitness history over generations
        """
        # Initialize population
        population = [self.generate_random_route() for _ in range(population_size)]
        best_fitness_history = []
        
        for gen in range(generations):
            # Evaluate fitness
            fitness_scores = [self.fitness(route) for route in population]
            
            # Track best solution
            best_idx = np.argmax(fitness_scores)
            best_fitness_history.append(fitness_scores[best_idx])
            
            # Selection (tournament selection)
            new_population = []
            for _ in range(population_size):
                # Tournament
                tournament_size = 5
                tournament = random.sample(list(zip(population, fitness_scores)), 
                                         min(tournament_size, len(population)))
                winner = max(tournament, key=lambda x: x[1])[0]
                new_population.append(winner.copy())
            
            # Crossover and mutation
            next_generation = [population[best_idx].copy()]  # Elitism
            
            for i in range(1, population_size):
                parent1, parent2 = random.sample(new_population, 2)
                
                # Crossover
                if random.random() < 0.7:
                    child = self.crossover(parent1, parent2)
                else:
                    child = parent1.copy()
                
                # Mutation
                if random.random() < mutation_rate:
                    child = self.mutate(child)
                
                next_generation.append(child)
            
            population = next_generation
        
        # Return best solution
        final_fitness = [self.fitness(route) for route in population]
        best_idx = np.argmax(final_fitness)
        
        return population[best_idx], best_fitness_history


# ============================================================================
# Example Problem Setup
# ============================================================================

print("=" * 70)
print("ROVER EXPLORATION ROUTE OPTIMIZATION")
print("=" * 70)
print()

# Define the problem
start_position = (0, 0)

# 15 Points of Interest with varying scientific values
poi_positions = [
    (5, 8), (12, 15), (18, 6), (25, 20), (8, 3),
    (15, 10), (22, 14), (10, 18), (28, 8), (6, 12),
    (20, 4), (14, 22), (30, 16), (4, 15), (26, 11)
]

# Scientific values (0-100): higher means more interesting
poi_values = [85, 92, 70, 95, 60, 88, 75, 80, 65, 78, 
              72, 90, 98, 68, 82]

# Constraints
max_distance = 120.0  # Maximum fuel allows 120 distance units
max_time = 150.0      # Maximum mission time

print("Problem Configuration:")
print(f"  Start Position: {start_position}")
print(f"  Number of POIs: {len(poi_positions)}")
print(f"  Maximum Distance: {max_distance} units")
print(f"  Maximum Mission Time: {max_time} time units")
print(f"  Sampling Time per POI: 2.0 time units")
print()

# Create optimizer
optimizer = RoverPathOptimizer(
    start_pos=start_position,
    poi_positions=poi_positions,
    poi_values=poi_values,
    max_distance=max_distance,
    max_time=max_time,
    sampling_time=2.0
)

# Run optimization
print("Running Genetic Algorithm...")
print("  Population Size: 100")
print("  Generations: 200")
print("  Mutation Rate: 0.2")
print()

best_route, fitness_history = optimizer.genetic_algorithm(
    population_size=100,
    generations=200,
    mutation_rate=0.2
)

# Calculate metrics for best route
total_distance = optimizer.calculate_distance(best_route)
total_time = optimizer.calculate_time(best_route)
total_value = optimizer.calculate_value(best_route)

print("=" * 70)
print("OPTIMIZATION RESULTS")
print("=" * 70)
print()
print(f"Best Route Found: {best_route}")
print(f"Number of POIs Visited: {len(best_route)}")
print(f"Total Scientific Value: {total_value:.2f} / {sum(poi_values):.2f}")
print(f"Total Distance: {total_distance:.2f} / {max_distance:.2f} units")
print(f"Total Mission Time: {total_time:.2f} / {max_time:.2f} units")
print(f"Distance Utilization: {(total_distance/max_distance)*100:.1f}%")
print(f"Time Utilization: {(total_time/max_time)*100:.1f}%")
print(f"Value Efficiency: {(total_value/sum(poi_values))*100:.1f}%")
print()

# ============================================================================
# Visualization
# ============================================================================

fig, axes = plt.subplots(2, 2, figsize=(16, 14))

# Plot 1: Route Visualization
ax1 = axes[0, 0]
ax1.set_title('Optimal Rover Exploration Route', fontsize=14, fontweight='bold')
ax1.set_xlabel('X Coordinate (km)', fontsize=11)
ax1.set_ylabel('Y Coordinate (km)', fontsize=11)
ax1.grid(True, alpha=0.3)
ax1.set_aspect('equal')

# Plot all POIs
for i, (pos, value) in enumerate(zip(poi_positions, poi_values)):
    if i in best_route:
        color = 'green'
        size = 200
        alpha = 0.7
    else:
        color = 'lightgray'
        size = 100
        alpha = 0.4
    
    ax1.scatter(pos[0], pos[1], c=color, s=size, alpha=alpha, edgecolors='black', linewidth=1.5)
    ax1.text(pos[0], pos[1]-1.2, f'POI{i}\n({value})', 
             ha='center', va='top', fontsize=8, fontweight='bold')

# Plot start position
ax1.scatter(start_position[0], start_position[1], c='red', s=300, marker='*', 
           edgecolors='black', linewidth=2, label='Start', zorder=5)

# Plot route
if len(best_route) > 0:
    route_coords = [start_position] + [poi_positions[i] for i in best_route]
    route_x = [coord[0] for coord in route_coords]
    route_y = [coord[1] for coord in route_coords]
    ax1.plot(route_x, route_y, 'b-', linewidth=2, alpha=0.6, label='Route')
    
    # Add arrows
    for i in range(len(route_coords) - 1):
        dx = route_coords[i+1][0] - route_coords[i][0]
        dy = route_coords[i+1][1] - route_coords[i][1]
        ax1.arrow(route_coords[i][0], route_coords[i][1], dx*0.8, dy*0.8,
                 head_width=0.8, head_length=0.5, fc='blue', ec='blue', alpha=0.5)

ax1.legend(loc='upper right', fontsize=10)

# Plot 2: Fitness Evolution
ax2 = axes[0, 1]
ax2.set_title('Genetic Algorithm Convergence', fontsize=14, fontweight='bold')
ax2.set_xlabel('Generation', fontsize=11)
ax2.set_ylabel('Fitness Score', fontsize=11)
ax2.plot(fitness_history, linewidth=2, color='darkblue')
ax2.grid(True, alpha=0.3)
ax2.fill_between(range(len(fitness_history)), fitness_history, alpha=0.3, color='lightblue')

# Plot 3: POI Values Comparison
ax3 = axes[1, 0]
ax3.set_title('Scientific Value: Visited vs Unvisited POIs', fontsize=14, fontweight='bold')
ax3.set_xlabel('POI Index', fontsize=11)
ax3.set_ylabel('Scientific Value', fontsize=11)

visited_values = [poi_values[i] if i in best_route else 0 for i in range(len(poi_values))]
unvisited_values = [poi_values[i] if i not in best_route else 0 for i in range(len(poi_values))]

x = np.arange(len(poi_values))
width = 0.8

ax3.bar(x, visited_values, width, label='Visited', color='green', alpha=0.7)
ax3.bar(x, unvisited_values, width, bottom=visited_values, label='Unvisited', 
        color='lightgray', alpha=0.7)
ax3.set_xticks(x)
ax3.set_xticklabels([f'POI{i}' for i in range(len(poi_values))], rotation=45, ha='right')
ax3.legend(fontsize=10)
ax3.grid(True, alpha=0.3, axis='y')

# Plot 4: Resource Utilization
ax4 = axes[1, 1]
ax4.set_title('Resource Utilization Analysis', fontsize=14, fontweight='bold')

categories = ['Distance\n(Fuel)', 'Time', 'Scientific\nValue']
utilization = [
    (total_distance / max_distance) * 100,
    (total_time / max_time) * 100,
    (total_value / sum(poi_values)) * 100
]
colors = ['#FF6B6B', '#4ECDC4', '#45B7D1']

bars = ax4.barh(categories, utilization, color=colors, alpha=0.7, edgecolor='black', linewidth=1.5)
ax4.set_xlabel('Utilization (%)', fontsize=11)
ax4.set_xlim(0, 110)
ax4.axvline(x=100, color='red', linestyle='--', linewidth=2, label='Maximum Capacity')
ax4.grid(True, alpha=0.3, axis='x')

# Add percentage labels
for i, (bar, val) in enumerate(zip(bars, utilization)):
    ax4.text(val + 2, bar.get_y() + bar.get_height()/2, f'{val:.1f}%', 
            va='center', fontsize=11, fontweight='bold')

ax4.legend(fontsize=10)

plt.tight_layout()
plt.savefig('rover_optimization_results.png', dpi=300, bbox_inches='tight')
plt.show()

print()
print("=" * 70)
print("Visualization saved as 'rover_optimization_results.png'")
print("=" * 70)