"""
Task 1 Sensitivity Analysis for Moon Colony Logistics

This module performs comprehensive sensitivity analysis on key model parameters:
1. Rocket payload capacity (100-150 tons, as specified in problem)
2. Elevator capacity (±20%)
3. Launch frequency (0.5-2 launches/day)
4. Engine technology (Isp variations)
5. Structural coefficient α

Author: MCM 2026 Team
"""

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib import rcParams
import pandas as pd
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
import warnings
warnings.filterwarnings('ignore')

# Font settings
rcParams['font.sans-serif'] = ['Arial Unicode MS', 'DejaVu Sans', 'SimHei']
rcParams['axes.unicode_minus'] = False

# ============== Physical Constants ==============
G0 = 9.81  # m/s²
OMEGA_EARTH = 7.27e-5  # rad/s
R_EARTH = 6.371e6  # m

# Baseline Mission parameters
TOTAL_PAYLOAD = 100e6  # 100 million metric tons

# Baseline Space Elevator parameters
BASELINE_NUM_ELEVATORS = 3
BASELINE_ELEVATOR_CAPACITY_PER_YEAR = 179000  # metric tons per elevator per year
BASELINE_ELEVATOR_SPECIFIC_ENERGY = 157.2e9  # J per metric ton

# Baseline Rocket parameters
BASELINE_PAYLOAD_PER_LAUNCH = 125  # metric tons (range: 100-150)
BASELINE_ISP = 363  # seconds (LOX/CH4)
BASELINE_SPECIFIC_FUEL_ENERGY = 11.9e6  # J/kg
BASELINE_ALPHA = 0.10  # structural coefficient
BASELINE_NUM_STAGES = 3
BASELINE_DELTA_V = 13300  # m/s

# Baseline Launch frequency
BASELINE_LAUNCHES_PER_DAY = 1  # per site


# ============== Launch Sites ==============
@dataclass
class LaunchSite:
    name: str
    short_name: str
    latitude: float
    
    @property
    def abs_latitude(self) -> float:
        return abs(self.latitude)
    
    def delta_v_loss(self, omega=OMEGA_EARTH, r=R_EARTH) -> float:
        v_equator = omega * r
        v_site = omega * r * np.cos(np.radians(self.abs_latitude))
        return v_equator - v_site


LAUNCH_SITES = [
    LaunchSite("Kourou", "Kourou", 5.2),
    LaunchSite("SDSC", "SDSC", 13.7),
    LaunchSite("Texas", "Texas", 26.0),
    LaunchSite("Florida", "Florida", 28.5),
    LaunchSite("California", "California", 34.7),
    LaunchSite("Virginia", "Virginia", 37.8),
    LaunchSite("Taiyuan", "Taiyuan", 38.8),
    LaunchSite("Mahia", "Mahia", 39.3),
    LaunchSite("Baikonur", "Baikonur", 45.6),
    LaunchSite("Alaska", "Alaska", 57.4),
]
LAUNCH_SITES = sorted(LAUNCH_SITES, key=lambda x: x.abs_latitude)
NUM_SITES = len(LAUNCH_SITES)


# ============== Core Calculation Functions ==============

def fuel_ratio_multistage(delta_v: float, isp: float, alpha: float, num_stages: int) -> float:
    """Calculate multi-stage rocket fuel/payload ratio"""
    ve = isp * G0
    delta_v_per_stage = delta_v / num_stages
    R_stage = np.exp(delta_v_per_stage / ve)
    
    denominator = 1 - alpha * (R_stage - 1)
    if denominator <= 0:
        return np.inf
    
    k_stage = (R_stage - 1) / denominator
    
    total_fuel_ratio = 0
    remaining_ratio = 1.0
    
    for _ in range(num_stages):
        fuel_this_stage = remaining_ratio * k_stage
        total_fuel_ratio += fuel_this_stage
        remaining_ratio *= (1 + k_stage * (1 + alpha))
    
    return total_fuel_ratio


def calculate_scenario(
    completion_years: float,
    payload_per_launch: float = BASELINE_PAYLOAD_PER_LAUNCH,
    elevator_capacity: float = BASELINE_NUM_ELEVATORS * BASELINE_ELEVATOR_CAPACITY_PER_YEAR,
    launches_per_day: float = BASELINE_LAUNCHES_PER_DAY,
    isp: float = BASELINE_ISP,
    alpha: float = BASELINE_ALPHA,
    num_stages: int = BASELINE_NUM_STAGES,
    delta_v_base: float = BASELINE_DELTA_V,
    specific_fuel_energy: float = BASELINE_SPECIFIC_FUEL_ENERGY,
) -> Optional[Dict]:
    """Calculate scenario for given parameters"""
    
    # Space elevator transport
    elevator_payload = min(elevator_capacity * completion_years, TOTAL_PAYLOAD)
    elevator_energy = elevator_payload * BASELINE_ELEVATOR_SPECIFIC_ENERGY
    
    # Remaining payload for rockets
    remaining_payload = TOTAL_PAYLOAD - elevator_payload
    
    if remaining_payload <= 0:
        return {
            'years': completion_years,
            'elevator_payload': elevator_payload,
            'rocket_payload': 0,
            'total_energy_PJ': elevator_energy / 1e15,
            'rocket_launches': 0,
            'elevator_fraction': 1.0,
            'feasible': True,
        }
    
    # Rocket launches needed
    rocket_launches_needed = int(np.ceil(remaining_payload / payload_per_launch))
    
    # Check feasibility
    days_available = completion_years * 365
    max_launches_per_site = int(days_available * launches_per_day)
    total_rocket_capacity = NUM_SITES * max_launches_per_site * payload_per_launch
    
    if remaining_payload > total_rocket_capacity:
        return {
            'years': completion_years,
            'feasible': False,
            'shortage': remaining_payload - total_rocket_capacity,
        }
    
    # Calculate rocket energy (low-latitude priority)
    rocket_energy = 0
    remaining_launches = rocket_launches_needed
    
    for site in LAUNCH_SITES:
        if remaining_launches <= 0:
            break
        allocated = min(remaining_launches, max_launches_per_site)
        
        # Calculate fuel ratio for this site
        total_dv = delta_v_base + site.delta_v_loss()
        k = fuel_ratio_multistage(total_dv, isp, alpha, num_stages)
        fuel_per_ton = k * 1000  # kg fuel per metric ton payload
        energy_per_ton = fuel_per_ton * specific_fuel_energy
        
        rocket_energy += energy_per_ton * payload_per_launch * allocated
        remaining_launches -= allocated
    
    total_energy = elevator_energy + rocket_energy
    
    return {
        'years': completion_years,
        'elevator_payload': elevator_payload,
        'rocket_payload': remaining_payload,
        'total_energy_PJ': total_energy / 1e15,
        'rocket_launches': rocket_launches_needed,
        'elevator_fraction': elevator_payload / TOTAL_PAYLOAD,
        'feasible': True,
    }


def find_minimum_timeline(
    payload_per_launch: float = BASELINE_PAYLOAD_PER_LAUNCH,
    elevator_capacity: float = BASELINE_NUM_ELEVATORS * BASELINE_ELEVATOR_CAPACITY_PER_YEAR,
    launches_per_day: float = BASELINE_LAUNCHES_PER_DAY,
    **kwargs
) -> float:
    """Find minimum feasible timeline"""
    for years in np.linspace(50, 300, 500):
        result = calculate_scenario(
            years, payload_per_launch, elevator_capacity, launches_per_day, **kwargs
        )
        if result and result.get('feasible', False):
            return years
    return 300  # fallback


def generate_tradeoff_curve(
    year_min: float = None,
    year_max: float = 250,
    num_points: int = 200,
    **kwargs
) -> pd.DataFrame:
    """Generate trade-off curve for given parameters"""
    
    if year_min is None:
        year_min = find_minimum_timeline(**kwargs)
    
    years_range = np.linspace(year_min, year_max, num_points)
    
    results = []
    for years in years_range:
        scenario = calculate_scenario(years, **kwargs)
        if scenario and scenario.get('feasible', False):
            results.append({
                'years': years,
                'energy_PJ': scenario['total_energy_PJ'],
                'elevator_fraction': scenario['elevator_fraction'],
                'rocket_launches': scenario.get('rocket_launches', 0),
            })
    
    return pd.DataFrame(results)


def find_optimal_point(df: pd.DataFrame, lambda_cost: float) -> Dict:
    """Find optimal point for given λ"""
    if df.empty:
        return {'years': np.nan, 'energy_PJ': np.nan}
    
    total_cost = df['energy_PJ'].values + lambda_cost * df['years'].values
    opt_idx = np.argmin(total_cost)
    
    return {
        'years': df['years'].iloc[opt_idx],
        'energy_PJ': df['energy_PJ'].iloc[opt_idx],
        'elevator_fraction': df['elevator_fraction'].iloc[opt_idx],
    }


# ============== Sensitivity Analysis Functions ==============

def sensitivity_payload_capacity():
    """Analyze sensitivity to rocket payload capacity (100-150 tons)"""
    print("\n[1] Analyzing Rocket Payload Capacity Sensitivity (100-150 tons)...")
    
    payload_range = np.linspace(100, 150, 11)
    results = []
    
    for payload in payload_range:
        elevator_cap = BASELINE_NUM_ELEVATORS * BASELINE_ELEVATOR_CAPACITY_PER_YEAR
        t_min = find_minimum_timeline(payload_per_launch=payload, elevator_capacity=elevator_cap)
        t_elev = TOTAL_PAYLOAD / elevator_cap
        
        df = generate_tradeoff_curve(
            year_min=t_min, year_max=250,
            payload_per_launch=payload, elevator_capacity=elevator_cap
        )
        
        if not df.empty:
            opt_504 = find_optimal_point(df, 504)
            energy_at_139 = df[df['years'] >= 139]['energy_PJ'].iloc[0] if any(df['years'] >= 139) else np.nan
            
            results.append({
                'payload_tons': payload,
                'min_timeline': t_min,
                'elevator_only_timeline': t_elev,
                'optimal_T_lambda504': opt_504['years'],
                'optimal_E_lambda504': opt_504['energy_PJ'],
                'energy_at_139y': energy_at_139,
            })
    
    return pd.DataFrame(results)


def sensitivity_elevator_capacity():
    """Analyze sensitivity to elevator capacity (±20%)"""
    print("\n[2] Analyzing Elevator Capacity Sensitivity (±20%)...")
    
    baseline_cap = BASELINE_NUM_ELEVATORS * BASELINE_ELEVATOR_CAPACITY_PER_YEAR
    capacity_factors = np.linspace(0.8, 1.2, 9)
    results = []
    
    for factor in capacity_factors:
        elevator_cap = baseline_cap * factor
        t_min = find_minimum_timeline(elevator_capacity=elevator_cap)
        t_elev = TOTAL_PAYLOAD / elevator_cap
        
        df = generate_tradeoff_curve(
            year_min=t_min, year_max=300,
            elevator_capacity=elevator_cap
        )
        
        if not df.empty:
            # Find critical λ
            lambda_crit = None
            for lam in np.linspace(300, 700, 200):
                opt = find_optimal_point(df, lam)
                if opt['years'] < t_elev - 5:
                    lambda_crit = lam
                    break
            
            opt_504 = find_optimal_point(df, 504)
            
            results.append({
                'capacity_factor': factor,
                'capacity_tons_per_year': elevator_cap,
                'min_timeline': t_min,
                'elevator_only_timeline': t_elev,
                'critical_lambda': lambda_crit,
                'optimal_T_lambda504': opt_504['years'],
                'optimal_E_lambda504': opt_504['energy_PJ'],
            })
    
    return pd.DataFrame(results)


def sensitivity_launch_frequency():
    """Analyze sensitivity to launch frequency"""
    print("\n[3] Analyzing Launch Frequency Sensitivity (0.5-2 launches/day)...")
    
    frequency_range = np.linspace(0.5, 2.0, 7)
    results = []
    
    for freq in frequency_range:
        elevator_cap = BASELINE_NUM_ELEVATORS * BASELINE_ELEVATOR_CAPACITY_PER_YEAR
        t_min = find_minimum_timeline(launches_per_day=freq, elevator_capacity=elevator_cap)
        t_elev = TOTAL_PAYLOAD / elevator_cap
        
        df = generate_tradeoff_curve(
            year_min=t_min, year_max=250,
            launches_per_day=freq, elevator_capacity=elevator_cap
        )
        
        if not df.empty:
            opt_504 = find_optimal_point(df, 504)
            
            results.append({
                'launches_per_day': freq,
                'min_timeline': t_min,
                'elevator_only_timeline': t_elev,
                'optimal_T_lambda504': opt_504['years'],
                'optimal_E_lambda504': opt_504['energy_PJ'],
                'rocket_capacity_tons_per_year': freq * 365 * NUM_SITES * BASELINE_PAYLOAD_PER_LAUNCH,
            })
    
    return pd.DataFrame(results)


def sensitivity_engine_technology():
    """Analyze sensitivity to engine technology (Isp)"""
    print("\n[4] Analyzing Engine Technology Sensitivity (Isp)...")
    
    engines = [
        ('Solid', 280, 5.0e6),
        ('LOX/Kerosene', 350, 10.3e6),
        ('LOX/Methane', 363, 11.9e6),
        ('LOX/Hydrogen', 450, 15.5e6),
    ]
    
    results = []
    
    for name, isp, specific_energy in engines:
        elevator_cap = BASELINE_NUM_ELEVATORS * BASELINE_ELEVATOR_CAPACITY_PER_YEAR
        t_min = find_minimum_timeline(
            isp=isp, specific_fuel_energy=specific_energy, elevator_capacity=elevator_cap
        )
        
        df = generate_tradeoff_curve(
            year_min=t_min, year_max=250,
            isp=isp, specific_fuel_energy=specific_energy, elevator_capacity=elevator_cap
        )
        
        if not df.empty:
            opt_504 = find_optimal_point(df, 504)
            
            # Calculate fuel ratio at equator
            k = fuel_ratio_multistage(BASELINE_DELTA_V, isp, BASELINE_ALPHA, BASELINE_NUM_STAGES)
            
            results.append({
                'engine_type': name,
                'isp_seconds': isp,
                'specific_energy_MJ_kg': specific_energy / 1e6,
                'fuel_ratio': k,
                'min_timeline': t_min,
                'optimal_T_lambda504': opt_504['years'],
                'optimal_E_lambda504': opt_504['energy_PJ'],
            })
    
    return pd.DataFrame(results)


def sensitivity_structural_coefficient():
    """Analyze sensitivity to structural coefficient α"""
    print("\n[5] Analyzing Structural Coefficient Sensitivity (α)...")
    
    alpha_range = np.linspace(0.06, 0.14, 9)
    results = []
    
    for alpha in alpha_range:
        elevator_cap = BASELINE_NUM_ELEVATORS * BASELINE_ELEVATOR_CAPACITY_PER_YEAR
        t_min = find_minimum_timeline(alpha=alpha, elevator_capacity=elevator_cap)
        
        df = generate_tradeoff_curve(
            year_min=t_min, year_max=250,
            alpha=alpha, elevator_capacity=elevator_cap
        )
        
        if not df.empty:
            opt_504 = find_optimal_point(df, 504)
            k = fuel_ratio_multistage(BASELINE_DELTA_V, BASELINE_ISP, alpha, BASELINE_NUM_STAGES)
            
            results.append({
                'alpha': alpha,
                'fuel_ratio': k,
                'min_timeline': t_min,
                'optimal_T_lambda504': opt_504['years'],
                'optimal_E_lambda504': opt_504['energy_PJ'],
            })
    
    return pd.DataFrame(results)


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

def plot_sensitivity_analysis(
    payload_df: pd.DataFrame,
    elevator_df: pd.DataFrame,
    frequency_df: pd.DataFrame,
    engine_df: pd.DataFrame,
    alpha_df: pd.DataFrame,
    save_path: str = '/Volumes/Files/code/mm/20260130_b/p1/sensitivity_analysis.png'
):
    """Create comprehensive sensitivity analysis visualization"""
    
    fig, axes = plt.subplots(2, 3, figsize=(18, 12))
    
    # ========== Plot 1: Rocket Payload Capacity ==========
    ax1 = axes[0, 0]
    
    ax1.plot(payload_df['payload_tons'], payload_df['min_timeline'], 
             'b-o', linewidth=2, markersize=8, label='Minimum Timeline')
    ax1.plot(payload_df['payload_tons'], payload_df['optimal_T_lambda504'], 
             'r-s', linewidth=2, markersize=8, label='Optimal T* (λ=504)')
    
    ax1.axhline(y=186.2, color='green', linestyle='--', alpha=0.7, label='Elevator-only')
    ax1.axvline(x=125, color='gray', linestyle=':', alpha=0.7, label='Baseline (125t)')
    
    ax1.fill_betweenx([80, 220], 100, 150, alpha=0.1, color='blue', label='Problem range')
    
    ax1.set_xlabel('Rocket Payload Capacity (metric tons)', fontsize=11)
    ax1.set_ylabel('Timeline (years)', fontsize=11)
    ax1.set_title('(a) Sensitivity to Rocket Payload Capacity\n(Problem specifies 100-150 tons)', fontsize=12)
    ax1.legend(loc='upper right', fontsize=9)
    ax1.grid(True, alpha=0.3)
    ax1.set_xlim(95, 155)
    ax1.set_ylim(80, 200)
    
    # ========== Plot 2: Elevator Capacity ==========
    ax2 = axes[0, 1]
    
    ax2.plot(elevator_df['capacity_factor'] * 100, elevator_df['elevator_only_timeline'], 
             'g-o', linewidth=2, markersize=8, label='Elevator-only Timeline')
    ax2.plot(elevator_df['capacity_factor'] * 100, elevator_df['optimal_T_lambda504'], 
             'r-s', linewidth=2, markersize=8, label='Optimal T* (λ=504)')
    ax2.plot(elevator_df['capacity_factor'] * 100, elevator_df['min_timeline'], 
             'b-^', linewidth=2, markersize=8, label='Minimum Timeline')
    
    ax2.axvline(x=100, color='gray', linestyle=':', alpha=0.7, label='Baseline')
    
    ax2.set_xlabel('Elevator Capacity (% of baseline)', fontsize=11)
    ax2.set_ylabel('Timeline (years)', fontsize=11)
    ax2.set_title('(b) Sensitivity to Elevator Capacity (±20%)', fontsize=12)
    ax2.legend(loc='upper right', fontsize=9)
    ax2.grid(True, alpha=0.3)
    ax2.set_xlim(75, 125)
    
    # ========== Plot 3: Launch Frequency ==========
    ax3 = axes[0, 2]
    
    ax3.plot(frequency_df['launches_per_day'], frequency_df['min_timeline'], 
             'b-o', linewidth=2, markersize=8, label='Minimum Timeline')
    ax3.plot(frequency_df['launches_per_day'], frequency_df['optimal_T_lambda504'], 
             'r-s', linewidth=2, markersize=8, label='Optimal T* (λ=504)')
    
    ax3.axhline(y=186.2, color='green', linestyle='--', alpha=0.7, label='Elevator-only')
    ax3.axvline(x=1.0, color='gray', linestyle=':', alpha=0.7, label='Baseline (1/day)')
    
    ax3.set_xlabel('Launch Frequency (launches/site/day)', fontsize=11)
    ax3.set_ylabel('Timeline (years)', fontsize=11)
    ax3.set_title('(c) Sensitivity to Launch Frequency', fontsize=12)
    ax3.legend(loc='upper right', fontsize=9)
    ax3.grid(True, alpha=0.3)
    
    # ========== Plot 4: Engine Technology ==========
    ax4 = axes[1, 0]
    
    x_pos = np.arange(len(engine_df))
    width = 0.35
    
    bars1 = ax4.bar(x_pos - width/2, engine_df['optimal_E_lambda504'] / 1000, width, 
                    label='Energy at T* (×10³ PJ)', color='steelblue', alpha=0.8)
    
    ax4_twin = ax4.twinx()
    bars2 = ax4_twin.bar(x_pos + width/2, engine_df['fuel_ratio'], width, 
                         label='Fuel Ratio k', color='coral', alpha=0.8)
    
    ax4.set_xlabel('Engine Type', fontsize=11)
    ax4.set_ylabel('Energy (×10³ PJ)', fontsize=11, color='steelblue')
    ax4_twin.set_ylabel('Fuel Ratio k', fontsize=11, color='coral')
    ax4.set_title('(d) Sensitivity to Engine Technology', fontsize=12)
    ax4.set_xticks(x_pos)
    ax4.set_xticklabels(engine_df['engine_type'], rotation=15)
    ax4.legend(loc='upper left', fontsize=9)
    ax4_twin.legend(loc='upper right', fontsize=9)
    ax4.grid(True, alpha=0.3, axis='y')
    
    # ========== Plot 5: Structural Coefficient ==========
    ax5 = axes[1, 1]
    
    ax5.plot(alpha_df['alpha'], alpha_df['fuel_ratio'], 
             'g-o', linewidth=2, markersize=8, label='Fuel Ratio k')
    
    ax5_twin = ax5.twinx()
    ax5_twin.plot(alpha_df['alpha'], alpha_df['optimal_E_lambda504'] / 1000, 
                  'r-s', linewidth=2, markersize=8, label='Energy at T* (×10³ PJ)')
    
    ax5.axvline(x=0.10, color='gray', linestyle=':', alpha=0.7, label='Baseline')
    
    ax5.set_xlabel('Structural Coefficient α', fontsize=11)
    ax5.set_ylabel('Fuel Ratio k', fontsize=11, color='green')
    ax5_twin.set_ylabel('Energy (×10³ PJ)', fontsize=11, color='red')
    ax5.set_title('(e) Sensitivity to Structural Coefficient α', fontsize=12)
    ax5.legend(loc='upper left', fontsize=9)
    ax5_twin.legend(loc='upper right', fontsize=9)
    ax5.grid(True, alpha=0.3)
    
    # ========== Plot 6: Summary Tornado Chart ==========
    ax6 = axes[1, 2]
    
    # Calculate sensitivity indices (% change in T* for parameter variation)
    baseline_T = 139  # baseline optimal timeline
    
    sensitivities = [
        ('Payload\n(100→150t)', 
         (payload_df['optimal_T_lambda504'].iloc[0] - baseline_T) / baseline_T * 100,
         (payload_df['optimal_T_lambda504'].iloc[-1] - baseline_T) / baseline_T * 100),
        ('Elevator Cap.\n(80%→120%)', 
         (elevator_df['optimal_T_lambda504'].iloc[0] - baseline_T) / baseline_T * 100,
         (elevator_df['optimal_T_lambda504'].iloc[-1] - baseline_T) / baseline_T * 100),
        ('Launch Freq.\n(0.5→2/day)', 
         (frequency_df['optimal_T_lambda504'].iloc[0] - baseline_T) / baseline_T * 100,
         (frequency_df['optimal_T_lambda504'].iloc[-1] - baseline_T) / baseline_T * 100),
        ('Struct. Coef.\n(0.06→0.14)', 
         (alpha_df['optimal_T_lambda504'].iloc[0] - baseline_T) / baseline_T * 100,
         (alpha_df['optimal_T_lambda504'].iloc[-1] - baseline_T) / baseline_T * 100),
    ]
    
    y_pos = np.arange(len(sensitivities))
    
    for i, (name, low, high) in enumerate(sensitivities):
        ax6.barh(i, high - low, left=min(low, high), height=0.6, 
                color='steelblue' if high > low else 'coral', alpha=0.7)
        ax6.plot([low, high], [i, i], 'k-', linewidth=2)
        ax6.plot(low, i, 'ko', markersize=8)
        ax6.plot(high, i, 'k^', markersize=8)
    
    ax6.axvline(x=0, color='black', linestyle='-', linewidth=1)
    ax6.set_yticks(y_pos)
    ax6.set_yticklabels([s[0] for s in sensitivities])
    ax6.set_xlabel('Change in Optimal Timeline T* (%)', fontsize=11)
    ax6.set_title('(f) Sensitivity Summary (Tornado Chart)', fontsize=12)
    ax6.grid(True, alpha=0.3, axis='x')
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"\nSensitivity analysis plot saved to: {save_path}")
    
    return fig


def plot_tornado_chart_standalone(
    payload_df: pd.DataFrame,
    elevator_df: pd.DataFrame, 
    frequency_df: pd.DataFrame,
    alpha_df: pd.DataFrame,
    save_path: str = '/Volumes/Files/code/mm/20260130_b/p1/tornado_chart.png'
):
    """
    单独绘制 Tornado Chart（敏感性汇总图）
    - XY轴反转（垂直柱状图）
    - 长宽比 0.618
    - 无标题
    - 有图例
    - 低饱和度色系
    """
    # 使用0.618黄金比例
    fig_width = 8
    fig_height = fig_width * 0.618
    fig, ax = plt.subplots(figsize=(fig_width, fig_height))
    
    # 计算敏感性指标
    baseline_T = 139
    
    sensitivities = [
        ('Payload\n(100→150t)', 
         (payload_df['optimal_T_lambda504'].iloc[0] - baseline_T) / baseline_T * 100,
         (payload_df['optimal_T_lambda504'].iloc[-1] - baseline_T) / baseline_T * 100),
        ('Elevator Cap.\n(80%→120%)', 
         (elevator_df['optimal_T_lambda504'].iloc[0] - baseline_T) / baseline_T * 100,
         (elevator_df['optimal_T_lambda504'].iloc[-1] - baseline_T) / baseline_T * 100),
        ('Launch Freq.\n(0.5→2/day)', 
         (frequency_df['optimal_T_lambda504'].iloc[0] - baseline_T) / baseline_T * 100,
         (frequency_df['optimal_T_lambda504'].iloc[-1] - baseline_T) / baseline_T * 100),
        ('Struct. Coef.\n(0.06→0.14)', 
         (alpha_df['optimal_T_lambda504'].iloc[0] - baseline_T) / baseline_T * 100,
         (alpha_df['optimal_T_lambda504'].iloc[-1] - baseline_T) / baseline_T * 100),
    ]
    
    # 低饱和度配色
    color_positive = '#7A9CBF'  # 灰蓝色（参数增加导致T*增加）
    color_negative = '#C4816B'  # 灰橙红色（参数增加导致T*减少）
    
    x_pos = np.arange(len(sensitivities))
    bar_width = 0.6
    
    # 绘制柱状图（XY反转，使用垂直柱状图）
    for i, (name, low, high) in enumerate(sensitivities):
        bar_color = color_positive if high > low else color_negative
        bar_height = high - low
        bar_bottom = min(low, high)
        
        ax.bar(i, bar_height, bottom=bar_bottom, width=bar_width, 
               color=bar_color, alpha=0.8, edgecolor='#555555', linewidth=0.8)
        
        # 标记端点
        ax.plot(i, low, 'o', color='#444444', markersize=7, zorder=5)
        ax.plot(i, high, '^', color='#444444', markersize=7, zorder=5)
        ax.plot([i, i], [low, high], '-', color='#444444', linewidth=1.5, zorder=4)
    
    # 基准线
    ax.axhline(y=0, color='#333333', linestyle='-', linewidth=1.2)
    
    # 设置轴
    ax.set_xticks(x_pos)
    ax.set_xticklabels([s[0] for s in sensitivities], fontsize=10)
    ax.set_ylabel('Change in Optimal Timeline T* (%)', fontsize=11)
    ax.grid(True, alpha=0.3, axis='y')
    
    # 添加图例
    from matplotlib.patches import Patch
    from matplotlib.lines import Line2D
    legend_elements = [
        Patch(facecolor=color_negative, edgecolor='#555555', alpha=0.8, 
              label='Parameter ↑ → T* ↓'),
        Patch(facecolor=color_positive, edgecolor='#555555', alpha=0.8, 
              label='Parameter ↑ → T* ↑'),
        Line2D([0], [0], marker='o', color='#444444', linestyle='None', 
               markersize=6, label='Low value'),
        Line2D([0], [0], marker='^', color='#444444', linestyle='None', 
               markersize=6, label='High value'),
    ]
    ax.legend(handles=legend_elements, loc='upper right', fontsize=9, framealpha=0.9)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"Tornado chart saved to: {save_path}")
    
    return fig


def plot_tradeoff_comparison(
    save_path: str = '/Volumes/Files/code/mm/20260130_b/p1/sensitivity_tradeoff.png'
):
    """Plot trade-off curves for different parameter values"""
    
    fig, axes = plt.subplots(2, 2, figsize=(14, 12))
    
    elevator_cap = BASELINE_NUM_ELEVATORS * BASELINE_ELEVATOR_CAPACITY_PER_YEAR
    
    # ========== Plot 1: Different Payload Capacities ==========
    ax1 = axes[0, 0]
    
    for payload, color in [(100, 'red'), (125, 'blue'), (150, 'green')]:
        df = generate_tradeoff_curve(payload_per_launch=payload, elevator_capacity=elevator_cap)
        if not df.empty:
            ax1.plot(df['years'], df['energy_PJ'] / 1000, color=color, linewidth=2, 
                    label=f'{payload} tons')
    
    ax1.set_xlabel('Timeline (years)', fontsize=11)
    ax1.set_ylabel('Energy (×10³ PJ)', fontsize=11)
    ax1.set_title('Trade-off Curves: Rocket Payload Capacity', fontsize=12)
    ax1.legend(title='Payload')
    ax1.grid(True, alpha=0.3)
    ax1.set_xlim(95, 200)
    
    # ========== Plot 2: Different Elevator Capacities ==========
    ax2 = axes[0, 1]
    
    for factor, color in [(0.8, 'red'), (1.0, 'blue'), (1.2, 'green')]:
        df = generate_tradeoff_curve(elevator_capacity=elevator_cap * factor)
        if not df.empty:
            ax2.plot(df['years'], df['energy_PJ'] / 1000, color=color, linewidth=2, 
                    label=f'{factor*100:.0f}%')
    
    ax2.set_xlabel('Timeline (years)', fontsize=11)
    ax2.set_ylabel('Energy (×10³ PJ)', fontsize=11)
    ax2.set_title('Trade-off Curves: Elevator Capacity', fontsize=12)
    ax2.legend(title='Capacity')
    ax2.grid(True, alpha=0.3)
    ax2.set_xlim(95, 250)
    
    # ========== Plot 3: Different Launch Frequencies ==========
    ax3 = axes[1, 0]
    
    for freq, color in [(0.5, 'red'), (1.0, 'blue'), (2.0, 'green')]:
        df = generate_tradeoff_curve(launches_per_day=freq, elevator_capacity=elevator_cap)
        if not df.empty:
            ax3.plot(df['years'], df['energy_PJ'] / 1000, color=color, linewidth=2, 
                    label=f'{freq}/day')
    
    ax3.set_xlabel('Timeline (years)', fontsize=11)
    ax3.set_ylabel('Energy (×10³ PJ)', fontsize=11)
    ax3.set_title('Trade-off Curves: Launch Frequency', fontsize=12)
    ax3.legend(title='Frequency')
    ax3.grid(True, alpha=0.3)
    ax3.set_xlim(45, 200)
    
    # ========== Plot 4: Different Engine Types ==========
    ax4 = axes[1, 1]
    
    engines = [
        ('LOX/Kerosene', 350, 10.3e6, 'red'),
        ('LOX/Methane', 363, 11.9e6, 'blue'),
        ('LOX/Hydrogen', 450, 15.5e6, 'green'),
    ]
    
    for name, isp, se, color in engines:
        df = generate_tradeoff_curve(isp=isp, specific_fuel_energy=se, elevator_capacity=elevator_cap)
        if not df.empty:
            ax4.plot(df['years'], df['energy_PJ'] / 1000, color=color, linewidth=2, 
                    label=f'{name}')
    
    ax4.set_xlabel('Timeline (years)', fontsize=11)
    ax4.set_ylabel('Energy (×10³ PJ)', fontsize=11)
    ax4.set_title('Trade-off Curves: Engine Technology', fontsize=12)
    ax4.legend(title='Engine')
    ax4.grid(True, alpha=0.3)
    ax4.set_xlim(95, 200)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"Trade-off comparison plot saved to: {save_path}")
    
    return fig


# ============== Main ==============

def main():
    print("=" * 70)
    print("Task 1 Sensitivity Analysis for Moon Colony Logistics")
    print("=" * 70)
    
    # Run all sensitivity analyses
    payload_df = sensitivity_payload_capacity()
    elevator_df = sensitivity_elevator_capacity()
    frequency_df = sensitivity_launch_frequency()
    engine_df = sensitivity_engine_technology()
    alpha_df = sensitivity_structural_coefficient()
    
    # Print summary tables
    print("\n" + "=" * 70)
    print("SENSITIVITY ANALYSIS RESULTS")
    print("=" * 70)
    
    print("\n[1] Rocket Payload Capacity (Problem range: 100-150 tons)")
    print(payload_df.to_string(index=False))
    
    print("\n[2] Elevator Capacity (±20%)")
    print(elevator_df.to_string(index=False))
    
    print("\n[3] Launch Frequency")
    print(frequency_df.to_string(index=False))
    
    print("\n[4] Engine Technology")
    print(engine_df.to_string(index=False))
    
    print("\n[5] Structural Coefficient α")
    print(alpha_df.to_string(index=False))
    
    # Generate plots
    print("\n" + "=" * 70)
    print("Generating Visualization...")
    print("=" * 70)
    
    plot_sensitivity_analysis(payload_df, elevator_df, frequency_df, engine_df, alpha_df)
    plot_tradeoff_comparison()
    
    # Save data
    payload_df.to_csv('/Volumes/Files/code/mm/20260130_b/p1/sensitivity_payload.csv', index=False)
    elevator_df.to_csv('/Volumes/Files/code/mm/20260130_b/p1/sensitivity_elevator.csv', index=False)
    frequency_df.to_csv('/Volumes/Files/code/mm/20260130_b/p1/sensitivity_frequency.csv', index=False)
    engine_df.to_csv('/Volumes/Files/code/mm/20260130_b/p1/sensitivity_engine.csv', index=False)
    alpha_df.to_csv('/Volumes/Files/code/mm/20260130_b/p1/sensitivity_alpha.csv', index=False)
    print("\nData saved to CSV files.")
    
    # Summary
    print("\n" + "=" * 70)
    print("KEY FINDINGS")
    print("=" * 70)
    print("""
    1. ROCKET PAYLOAD CAPACITY (100-150 tons):
       - Higher capacity reduces minimum timeline significantly
       - Limited impact on energy when elevator dominates
       - Baseline 125 tons is a reasonable middle estimate
    
    2. ELEVATOR CAPACITY (±20%):
       - Most sensitive parameter for timeline
       - ±20% capacity changes elevator-only timeline by ~37 years
       - Critical for long-term planning
    
    3. LAUNCH FREQUENCY:
       - Doubling frequency halves minimum rocket-only timeline
       - Essential for time-constrained scenarios
       - Less impact when elevator provides majority of transport
    
    4. ENGINE TECHNOLOGY:
       - Primarily affects energy, not timeline
       - LOX/Hydrogen saves ~30% energy vs LOX/Methane
       - But LOX/Methane has practical advantages (storage, cost)
    
    5. STRUCTURAL COEFFICIENT α:
       - Lower α (better technology) improves fuel efficiency
       - Moderate impact on total energy
       - Progress from 0.10 to 0.08 saves ~8% energy
    """)
    
    print("=" * 70)
    print("Analysis Complete!")
    print("=" * 70)
    
    return payload_df, elevator_df, frequency_df, engine_df, alpha_df


if __name__ == "__main__":
    results = main()
