"""
Moon Colony Construction - Combined Scenario Analysis
(Space Elevator + Traditional Rockets)

Problem: Deliver 100 million metric tons to Moon
- 3 Space Elevators: 179,000 metric tons/year each (priority)
- 10 Rocket launch sites: max 1 launch/day each, 125 tons/launch
- Optimize: Use elevators first (lower specific energy), then rockets

Key: As timeline extends:
1. More payload handled by elevators (lower energy)
2. Remaining rocket launches shift to lower-latitude sites
"""

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

# 设置字体
rcParams['font.sans-serif'] = ['Arial Unicode MS', 'DejaVu Sans', 'SimHei']
rcParams['axes.unicode_minus'] = False

# ============== 物理常数 ==============
G0 = 9.81  # m/s²
OMEGA_EARTH = 7.27e-5  # rad/s
R_EARTH = 6.371e6  # m

# ============== 任务参数 ==============
TOTAL_PAYLOAD = 100e6  # 100 million metric tons

# ============== 太空电梯参数 ==============
NUM_ELEVATORS = 3
ELEVATOR_CAPACITY_PER_YEAR = 179000  # metric tons per elevator per year
TOTAL_ELEVATOR_CAPACITY = NUM_ELEVATORS * ELEVATOR_CAPACITY_PER_YEAR  # 537,000 tons/year

# 太空电梯能量参数 (从之前分析)
# 电梯提升能量: ~88.7 MJ/kg
# 顶端推进能量: ~68.5 MJ/kg
# 总计: ~157.2 MJ/kg = 157.2 TJ/1000 tons
ELEVATOR_SPECIFIC_ENERGY = 157.2e9  # J per metric ton (157.2 MJ/kg * 1000 kg)

# ============== 火箭参数 ==============
PAYLOAD_PER_LAUNCH = 125  # metric tons per launch
ISP = 363  # 比冲 (秒) - 液氧甲烷 (LOX/CH4, Raptor-class)
SPECIFIC_FUEL_ENERGY = 11.9e6  # J/kg
ALPHA = 0.10  # 结构系数
NUM_STAGES = 3
DELTA_V_BASE = 13300  # m/s (赤道发射到月球)

# 火箭比能量 (从之前分析: ~492.7 MJ/kg = 492.7 TJ/1000 tons for equator)
# 实际会因纬度变化


# ============== 发射场定义 ==============
@dataclass
class LaunchSite:
    name: str
    short_name: str
    latitude: float
    max_launches_per_day: int = 1
    
    @property
    def abs_latitude(self) -> float:
        return abs(self.latitude)
    
    @property
    def rotation_velocity(self) -> float:
        return OMEGA_EARTH * R_EARTH * np.cos(np.radians(self.abs_latitude))
    
    @property
    def delta_v_loss(self) -> float:
        v_equator = OMEGA_EARTH * R_EARTH
        return v_equator - self.rotation_velocity
    
    @property
    def total_delta_v(self) -> float:
        return DELTA_V_BASE + self.delta_v_loss


# 10个发射场 (按纬度排序)
LAUNCH_SITES = sorted([
    LaunchSite("Kourou (French Guiana)", "Kourou", 5.2),
    LaunchSite("Satish Dhawan (India)", "SDSC", 13.7),
    LaunchSite("Boca Chica (Texas)", "Texas", 26.0),
    LaunchSite("Cape Canaveral (Florida)", "Florida", 28.5),
    LaunchSite("Vandenberg (California)", "California", 34.7),
    LaunchSite("Wallops (Virginia)", "Virginia", 37.8),
    LaunchSite("Taiyuan (China)", "Taiyuan", 38.8),
    LaunchSite("Mahia (New Zealand)", "Mahia", 39.3),
    LaunchSite("Baikonur (Kazakhstan)", "Baikonur", 45.6),
    LaunchSite("Kodiak (Alaska)", "Alaska", 57.4),
], key=lambda x: x.abs_latitude)


# ============== 核心计算函数 ==============

def fuel_ratio_multistage(delta_v: float) -> float:
    """多级火箭燃料/载荷比"""
    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 rocket_energy_per_ton(site: LaunchSite) -> float:
    """火箭发射每吨载荷的能量消耗 (J/ton)"""
    k = fuel_ratio_multistage(site.total_delta_v)
    fuel_per_ton = k * 1000  # kg fuel per metric ton payload
    return fuel_per_ton * SPECIFIC_FUEL_ENERGY


def rocket_energy_per_launch(site: LaunchSite) -> float:
    """单次火箭发射的能量消耗 (J)"""
    return rocket_energy_per_ton(site) * PAYLOAD_PER_LAUNCH


# ============== 组合方案计算 ==============

def calculate_combined_scenario(completion_years: float) -> Dict:
    """
    计算给定完成年限下的组合方案
    
    策略:
    1. 太空电梯满负荷运行
    2. 剩余载荷由火箭发射，优先低纬度站点
    
    返回:
        方案详情字典
    """
    # 太空电梯运输量
    elevator_payload = min(TOTAL_ELEVATOR_CAPACITY * completion_years, TOTAL_PAYLOAD)
    elevator_energy = elevator_payload * ELEVATOR_SPECIFIC_ENERGY
    
    # 剩余需要火箭运输的载荷
    remaining_payload = TOTAL_PAYLOAD - elevator_payload
    
    if remaining_payload <= 0:
        # 完全由电梯完成
        return {
            'completion_years': completion_years,
            'elevator_payload': elevator_payload,
            'elevator_energy_PJ': elevator_energy / 1e15,
            'rocket_payload': 0,
            'rocket_energy_PJ': 0,
            'total_energy_PJ': elevator_energy / 1e15,
            'rocket_launches': 0,
            'sites_used': 0,
            'elevator_fraction': 1.0,
            'rocket_distribution': [],
        }
    
    # 需要火箭发射
    rocket_launches_needed = int(np.ceil(remaining_payload / PAYLOAD_PER_LAUNCH))
    
    # 按纬度优先分配火箭发射
    days_available = completion_years * 365
    max_launches_per_site = int(days_available)
    
    rocket_distribution = []
    remaining_launches = rocket_launches_needed
    rocket_energy = 0
    sites_used = 0
    
    for site in LAUNCH_SITES:
        if remaining_launches <= 0:
            rocket_distribution.append((site, 0))
        else:
            allocated = min(remaining_launches, max_launches_per_site)
            rocket_distribution.append((site, allocated))
            rocket_energy += rocket_energy_per_launch(site) * allocated
            remaining_launches -= allocated
            if allocated > 0:
                sites_used += 1
    
    # 检查是否能在规定时间内完成
    total_rocket_capacity = len(LAUNCH_SITES) * max_launches_per_site * PAYLOAD_PER_LAUNCH
    if remaining_payload > total_rocket_capacity:
        # 无法完成
        return None
    
    rocket_payload = rocket_launches_needed * PAYLOAD_PER_LAUNCH
    total_energy = elevator_energy + rocket_energy
    
    return {
        'completion_years': completion_years,
        'elevator_payload': elevator_payload,
        'elevator_energy_PJ': elevator_energy / 1e15,
        'rocket_payload': rocket_payload,
        'rocket_energy_PJ': rocket_energy / 1e15,
        'total_energy_PJ': total_energy / 1e15,
        'rocket_launches': rocket_launches_needed,
        'sites_used': sites_used,
        'elevator_fraction': elevator_payload / TOTAL_PAYLOAD,
        'rocket_distribution': rocket_distribution,
    }


def analyze_combined_timeline() -> pd.DataFrame:
    """
    分析不同完成年限下的组合方案
    """
    # 计算关键时间点
    # 1. 仅电梯完成所需时间
    elevator_only_years = TOTAL_PAYLOAD / TOTAL_ELEVATOR_CAPACITY
    
    # 2. 最短时间（电梯+全部火箭站点满负荷）
    rocket_capacity_per_year = len(LAUNCH_SITES) * 365 * PAYLOAD_PER_LAUNCH
    total_capacity_per_year = TOTAL_ELEVATOR_CAPACITY + rocket_capacity_per_year
    min_years = TOTAL_PAYLOAD / total_capacity_per_year
    
    print(f"Elevator-only completion: {elevator_only_years:.1f} years")
    print(f"Minimum completion (all systems): {min_years:.1f} years")
    print(f"Elevator capacity: {TOTAL_ELEVATOR_CAPACITY:,} tons/year")
    print(f"Rocket capacity: {rocket_capacity_per_year:,} tons/year")
    
    # 分析范围
    years_range = np.linspace(min_years, elevator_only_years * 1.1, 200)
    
    results = []
    for years in years_range:
        scenario = calculate_combined_scenario(years)
        if scenario is not None:
            results.append(scenario)
    
    return pd.DataFrame(results)


# ============== 可视化 ==============

def plot_combined_analysis(
    save_path: str = '/Volumes/Files/code/mm/20260130_b/combined_scenario_analysis.png'
):
    """
    绘制组合方案分析图
    """
    df = analyze_combined_timeline()
    
    # 关键时间点
    elevator_only_years = TOTAL_PAYLOAD / TOTAL_ELEVATOR_CAPACITY
    rocket_capacity_per_year = len(LAUNCH_SITES) * 365 * PAYLOAD_PER_LAUNCH
    total_capacity_per_year = TOTAL_ELEVATOR_CAPACITY + rocket_capacity_per_year
    min_years = TOTAL_PAYLOAD / total_capacity_per_year
    
    fig, axes = plt.subplots(2, 2, figsize=(16, 14))
    
    # ========== 图1: 完成年份 vs 总能量消耗 ==========
    ax1 = axes[0, 0]
    
    ax1.plot(df['completion_years'], df['total_energy_PJ'], 'b-', linewidth=2.5, 
             label='Total Energy')
    ax1.plot(df['completion_years'], df['elevator_energy_PJ'], 'g--', linewidth=2, 
             label='Elevator Energy')
    ax1.plot(df['completion_years'], df['rocket_energy_PJ'], 'r--', linewidth=2, 
             label='Rocket Energy')
    
    # 标记关键点
    ax1.axvline(x=min_years, color='orange', linestyle=':', linewidth=1.5,
                label=f'Min time: {min_years:.0f} years')
    ax1.axvline(x=elevator_only_years, color='purple', linestyle=':', linewidth=1.5,
                label=f'Elevator-only: {elevator_only_years:.0f} years')
    
    ax1.set_xlabel('Completion Time (years)', fontsize=12)
    ax1.set_ylabel('Energy Consumption (PJ)', fontsize=12)
    ax1.set_title('Combined Scenario: Completion Time vs Energy\n'
                  '(Space Elevator + Rockets)', fontsize=14)
    ax1.legend(loc='upper right')
    ax1.grid(True, alpha=0.3)
    ax1.set_xlim(min_years * 0.95, elevator_only_years * 1.05)
    
    # ========== 图2: 电梯 vs 火箭载荷比例 ==========
    ax2 = axes[0, 1]
    
    ax2.fill_between(df['completion_years'], 0, df['elevator_fraction'] * 100, 
                     alpha=0.7, color='green', label='Space Elevator')
    ax2.fill_between(df['completion_years'], df['elevator_fraction'] * 100, 100, 
                     alpha=0.7, color='red', label='Rockets')
    
    ax2.axvline(x=min_years, color='orange', linestyle=':', linewidth=1.5)
    ax2.axvline(x=elevator_only_years, color='purple', linestyle=':', linewidth=1.5)
    
    ax2.set_xlabel('Completion Time (years)', fontsize=12)
    ax2.set_ylabel('Payload Share (%)', fontsize=12)
    ax2.set_title('Payload Distribution: Elevator vs Rockets', fontsize=14)
    ax2.legend(loc='center right')
    ax2.set_ylim(0, 100)
    ax2.set_xlim(min_years * 0.95, elevator_only_years * 1.05)
    ax2.grid(True, alpha=0.3)
    
    # ========== 图3: 使用的火箭发射场数量 ==========
    ax3 = axes[1, 0]
    
    ax3.plot(df['completion_years'], df['sites_used'], 'r-', linewidth=2)
    ax3.fill_between(df['completion_years'], 0, df['sites_used'], alpha=0.3, color='red')
    
    ax3.axvline(x=min_years, color='orange', linestyle=':', linewidth=1.5)
    ax3.axvline(x=elevator_only_years, color='purple', linestyle=':', linewidth=1.5)
    
    # 标注发射场
    for i, site in enumerate(LAUNCH_SITES):
        ax3.axhline(y=i+1, color='gray', linestyle=':', alpha=0.2)
        ax3.text(df['completion_years'].max() * 0.99, i+1 + 0.15, 
                f'{site.short_name} ({site.abs_latitude:.0f}°)', 
                fontsize=8, ha='right', va='bottom')
    
    ax3.set_xlabel('Completion Time (years)', fontsize=12)
    ax3.set_ylabel('Number of Rocket Sites Used', fontsize=12)
    ax3.set_title('Rocket Launch Sites Required\n(Longer timeline → fewer high-lat sites)', fontsize=14)
    ax3.set_ylim(0, 11)
    ax3.set_xlim(min_years * 0.95, elevator_only_years * 1.05)
    ax3.grid(True, alpha=0.3)
    
    # ========== 图4: 能量效率对比 ==========
    ax4 = axes[1, 1]
    
    # 计算平均比能量
    avg_specific_energy = df['total_energy_PJ'] * 1e15 / TOTAL_PAYLOAD / 1e9  # GJ/ton
    
    ax4.plot(df['completion_years'], avg_specific_energy, 'b-', linewidth=2)
    
    # 参考线
    elevator_specific = ELEVATOR_SPECIFIC_ENERGY / 1e9  # GJ/ton
    ax4.axhline(y=elevator_specific, color='green', linestyle='--', 
                label=f'Elevator only: {elevator_specific:.1f} GJ/ton')
    
    # 平均火箭比能量 (赤道)
    rocket_specific_equator = rocket_energy_per_ton(LAUNCH_SITES[0]) / 1e9
    ax4.axhline(y=rocket_specific_equator, color='red', linestyle='--', 
                label=f'Rocket (equator): {rocket_specific_equator:.1f} GJ/ton')
    
    ax4.axvline(x=min_years, color='orange', linestyle=':', linewidth=1.5)
    ax4.axvline(x=elevator_only_years, color='purple', linestyle=':', linewidth=1.5)
    
    ax4.set_xlabel('Completion Time (years)', fontsize=12)
    ax4.set_ylabel('Average Specific Energy (GJ/ton)', fontsize=12)
    ax4.set_title('Energy Efficiency vs Timeline\n(Longer time → more elevator → lower energy)', fontsize=14)
    ax4.legend()
    ax4.set_xlim(min_years * 0.95, elevator_only_years * 1.05)
    ax4.grid(True, alpha=0.3)
    
    plt.suptitle(f'Combined Scenario: 100M tons to Moon\n'
                 f'3 Elevators ({TOTAL_ELEVATOR_CAPACITY:,} t/yr) + 10 Rocket Sites',
                 fontsize=15, y=1.02)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"\n组合方案分析图已保存至: {save_path}")
    
    return df


def plot_scenario_comparison(
    save_path: str = '/Volumes/Files/code/mm/20260130_b/scenario_comparison.png'
):
    """
    对比三种方案：仅电梯、仅火箭、组合
    """
    # 关键时间点分析
    elevator_only_years = TOTAL_PAYLOAD / TOTAL_ELEVATOR_CAPACITY
    rocket_capacity_per_year = len(LAUNCH_SITES) * 365 * PAYLOAD_PER_LAUNCH
    rocket_only_min_years = TOTAL_PAYLOAD / rocket_capacity_per_year
    combined_min_years = TOTAL_PAYLOAD / (TOTAL_ELEVATOR_CAPACITY + rocket_capacity_per_year)
    
    # 各方案能量
    elevator_only_energy = TOTAL_PAYLOAD * ELEVATOR_SPECIFIC_ENERGY / 1e15  # PJ
    
    # 火箭方案 (平均所有站点)
    avg_rocket_energy_per_ton = sum(rocket_energy_per_ton(s) for s in LAUNCH_SITES) / len(LAUNCH_SITES)
    rocket_only_energy = TOTAL_PAYLOAD * avg_rocket_energy_per_ton / 1e15  # PJ
    
    # 组合方案 (最短时间 - 加余量)
    combined_scenario = calculate_combined_scenario(combined_min_years * 1.01)
    combined_min_energy = combined_scenario['total_energy_PJ'] if combined_scenario else elevator_only_energy
    
    # 组合方案 (仅电梯时间)
    combined_extended = calculate_combined_scenario(elevator_only_years * 0.95)
    combined_extended_energy = combined_extended['total_energy_PJ'] if combined_extended else elevator_only_energy
    
    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
    
    # ========== 图1: 完成时间对比 ==========
    ax1 = axes[0]
    
    scenarios = ['Elevator\nOnly', 'Rocket\nOnly', 'Combined\n(Min Time)', 'Combined\n(Extended)']
    times = [elevator_only_years, rocket_only_min_years, combined_min_years, elevator_only_years * 0.99]
    colors = ['green', 'red', 'blue', 'purple']
    
    bars = ax1.bar(scenarios, times, color=colors, alpha=0.7)
    
    ax1.set_ylabel('Completion Time (years)', fontsize=12)
    ax1.set_title('Completion Time Comparison', fontsize=14)
    ax1.grid(True, alpha=0.3, axis='y')
    
    for bar, t in zip(bars, times):
        ax1.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 2, 
                f'{t:.0f}y', ha='center', fontsize=11, fontweight='bold')
    
    # ========== 图2: 能量消耗对比 ==========
    ax2 = axes[1]
    
    energies = [elevator_only_energy, rocket_only_energy, combined_min_energy, combined_extended_energy]
    
    bars = ax2.bar(scenarios, energies, color=colors, alpha=0.7)
    
    ax2.set_ylabel('Total Energy (PJ)', fontsize=12)
    ax2.set_title('Energy Consumption Comparison', fontsize=14)
    ax2.grid(True, alpha=0.3, axis='y')
    
    for bar, e in zip(bars, energies):
        ax2.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 500, 
                f'{e:.0f}', ha='center', fontsize=11, fontweight='bold')
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"方案对比图已保存至: {save_path}")
    
    return {
        'elevator_only': {'years': elevator_only_years, 'energy_PJ': elevator_only_energy},
        'rocket_only': {'years': rocket_only_min_years, 'energy_PJ': rocket_only_energy},
        'combined_min': {'years': combined_min_years, 'energy_PJ': combined_min_energy},
        'combined_extended': {'years': elevator_only_years * 0.99, 'energy_PJ': combined_extended_energy},
    }


def print_summary(df: pd.DataFrame):
    """打印分析摘要"""
    elevator_only_years = TOTAL_PAYLOAD / TOTAL_ELEVATOR_CAPACITY
    rocket_capacity_per_year = len(LAUNCH_SITES) * 365 * PAYLOAD_PER_LAUNCH
    combined_min_years = TOTAL_PAYLOAD / (TOTAL_ELEVATOR_CAPACITY + rocket_capacity_per_year)
    
    print("\n" + "=" * 90)
    print("COMBINED SCENARIO (SPACE ELEVATOR + ROCKETS) ANALYSIS")
    print("=" * 90)
    
    print(f"\nSystem Capacities:")
    print(f"  - 3 Space Elevators: {TOTAL_ELEVATOR_CAPACITY:,} tons/year")
    print(f"  - 10 Rocket Sites: {rocket_capacity_per_year:,} tons/year (@ 1 launch/day each)")
    print(f"  - Combined: {TOTAL_ELEVATOR_CAPACITY + rocket_capacity_per_year:,} tons/year")
    
    print(f"\nSpecific Energy:")
    print(f"  - Elevator: {ELEVATOR_SPECIFIC_ENERGY/1e9:.1f} GJ/ton ({ELEVATOR_SPECIFIC_ENERGY/1e9/rocket_energy_per_ton(LAUNCH_SITES[0])*1e9*100:.1f}% of rocket)")
    print(f"  - Rocket (Kourou): {rocket_energy_per_ton(LAUNCH_SITES[0])/1e9:.1f} GJ/ton")
    print(f"  - Rocket (Alaska): {rocket_energy_per_ton(LAUNCH_SITES[-1])/1e9:.1f} GJ/ton")
    
    print(f"\nKey Scenarios:")
    
    # 最短时间 (加一点余量确保可以完成)
    min_scenario = calculate_combined_scenario(combined_min_years * 1.01)
    print(f"\n  1. Minimum Time (~{combined_min_years:.0f} years):")
    print(f"     - Elevator payload: {min_scenario['elevator_payload']/1e6:.1f}M tons ({min_scenario['elevator_fraction']*100:.1f}%)")
    print(f"     - Rocket payload: {min_scenario['rocket_payload']/1e6:.1f}M tons")
    print(f"     - Rocket launches: {min_scenario['rocket_launches']:,}")
    print(f"     - Sites used: {min_scenario['sites_used']}")
    print(f"     - Total energy: {min_scenario['total_energy_PJ']:.0f} PJ")
    
    # 中等时间
    mid_years = (combined_min_years + elevator_only_years) / 2
    mid_scenario = calculate_combined_scenario(mid_years)
    print(f"\n  2. Medium Timeline ({mid_years:.0f} years):")
    print(f"     - Elevator payload: {mid_scenario['elevator_payload']/1e6:.1f}M tons ({mid_scenario['elevator_fraction']*100:.1f}%)")
    print(f"     - Rocket payload: {mid_scenario['rocket_payload']/1e6:.1f}M tons")
    print(f"     - Rocket launches: {mid_scenario['rocket_launches']:,}")
    print(f"     - Sites used: {mid_scenario['sites_used']}")
    print(f"     - Total energy: {mid_scenario['total_energy_PJ']:.0f} PJ")
    
    # 接近仅电梯
    extended_years = elevator_only_years * 0.95
    extended_scenario = calculate_combined_scenario(extended_years)
    print(f"\n  3. Extended Timeline ({extended_years:.0f} years):")
    print(f"     - Elevator payload: {extended_scenario['elevator_payload']/1e6:.1f}M tons ({extended_scenario['elevator_fraction']*100:.1f}%)")
    print(f"     - Rocket payload: {extended_scenario['rocket_payload']/1e6:.1f}M tons")
    print(f"     - Rocket launches: {extended_scenario['rocket_launches']:,}")
    print(f"     - Sites used: {extended_scenario['sites_used']}")
    print(f"     - Total energy: {extended_scenario['total_energy_PJ']:.0f} PJ")
    
    # 仅电梯
    print(f"\n  4. Elevator Only ({elevator_only_years:.0f} years):")
    elevator_energy = TOTAL_PAYLOAD * ELEVATOR_SPECIFIC_ENERGY / 1e15
    print(f"     - Elevator payload: 100M tons (100%)")
    print(f"     - Rocket launches: 0")
    print(f"     - Total energy: {elevator_energy:.0f} PJ")
    
    print("\n" + "=" * 90)
    print("ENERGY SAVINGS ANALYSIS")
    print("=" * 90)
    
    rocket_only_energy = TOTAL_PAYLOAD * sum(rocket_energy_per_ton(s) for s in LAUNCH_SITES) / len(LAUNCH_SITES) / 1e15
    
    print(f"\n  Baseline (Rocket Only): {rocket_only_energy:.0f} PJ")
    print(f"  Combined (Min Time): {min_scenario['total_energy_PJ']:.0f} PJ (saves {(1-min_scenario['total_energy_PJ']/rocket_only_energy)*100:.0f}%)")
    print(f"  Combined (Extended): {extended_scenario['total_energy_PJ']:.0f} PJ (saves {(1-extended_scenario['total_energy_PJ']/rocket_only_energy)*100:.0f}%)")
    print(f"  Elevator Only: {elevator_energy:.0f} PJ (saves {(1-elevator_energy/rocket_only_energy)*100:.0f}%)")
    
    print("=" * 90)


def plot_three_scenarios_comparison(
    year_min: float = 100,
    year_max: float = 250,
    save_path: str = '/Volumes/Files/code/mm/20260130_b/three_scenarios_comparison.png'
):
    """
    绘制三种方案在指定年份范围内的能量消耗对比折线图
    
    方案:
    1. 纯火箭方案 (Rocket Only)
    2. 纯电梯方案 (Elevator Only)  
    3. 混合方案 (Combined)
    """
    years = np.linspace(year_min, year_max, 200)
    
    # 计算各方案的能量
    rocket_energy = []
    elevator_energy = []
    combined_energy = []
    
    # 关键时间点
    elevator_only_min_years = TOTAL_PAYLOAD / TOTAL_ELEVATOR_CAPACITY  # ~186 years
    rocket_only_min_years = TOTAL_PAYLOAD / (len(LAUNCH_SITES) * 365 * PAYLOAD_PER_LAUNCH)  # ~219 years
    combined_min_years = TOTAL_PAYLOAD / (TOTAL_ELEVATOR_CAPACITY + len(LAUNCH_SITES) * 365 * PAYLOAD_PER_LAUNCH)  # ~101 years
    
    for y in years:
        # ===== 纯火箭方案 =====
        if y >= rocket_only_min_years:
            # 可以完成，优先使用低纬度站点
            total_launches = int(TOTAL_PAYLOAD / PAYLOAD_PER_LAUNCH)
            days_available = y * 365
            max_per_site = int(days_available)
            
            energy = 0
            remaining = total_launches
            for site in LAUNCH_SITES:
                if remaining <= 0:
                    break
                allocated = min(remaining, max_per_site)
                energy += rocket_energy_per_launch(site) * allocated
                remaining -= allocated
            
            rocket_energy.append(energy / 1e15)
        else:
            # 无法完成
            rocket_energy.append(np.nan)
        
        # ===== 纯电梯方案 =====
        if y >= elevator_only_min_years:
            # 可以完成
            energy = TOTAL_PAYLOAD * ELEVATOR_SPECIFIC_ENERGY
            elevator_energy.append(energy / 1e15)
        else:
            # 无法完成
            elevator_energy.append(np.nan)
        
        # ===== 混合方案 =====
        if y >= combined_min_years:
            scenario = calculate_combined_scenario(y)
            if scenario:
                combined_energy.append(scenario['total_energy_PJ'])
            else:
                combined_energy.append(np.nan)
        else:
            combined_energy.append(np.nan)
    
    # 创建图表 - 使用0.618黄金比例，缩小尺寸以放大字体
    fig_width = 8
    fig_height = fig_width * 0.618
    fig, ax = plt.subplots(figsize=(fig_width, fig_height))
    
    # 中偏低饱和度配色
    color_rocket = '#B85450'    # 暗砖红色
    color_elevator = '#5B8266'  # 灰绿色
    color_combined = '#4A6FA5'  # 灰蓝色
    
    # 绘制三条折线
    ax.plot(years, rocket_energy, color=color_rocket, linewidth=2.5, label='Rocket Only', marker='', markersize=4)
    ax.plot(years, elevator_energy, color=color_elevator, linewidth=2.5, label='Elevator Only', marker='', markersize=4)
    ax.plot(years, combined_energy, color=color_combined, linewidth=2.5, label='Combined (Elevator + Rocket)', marker='', markersize=4)
    
    # 标记关键时间点
    ax.axvline(x=combined_min_years, color=color_combined, linestyle=':', alpha=0.7, linewidth=1.5)
    ax.axvline(x=elevator_only_min_years, color=color_elevator, linestyle=':', alpha=0.7, linewidth=1.5)
    ax.axvline(x=rocket_only_min_years, color=color_rocket, linestyle=':', alpha=0.7, linewidth=1.5)
    
    # 添加标注 - 使用低饱和度背景色，避免遮挡曲线
    y_max = ax.get_ylim()[1]
    ax.annotate(f'Combined\nMin: {combined_min_years:.0f}y', 
                xy=(combined_min_years, y_max * 0.55), fontsize=9, ha='center',
                bbox=dict(boxstyle='round', facecolor='#C5D5E8', edgecolor=color_combined, alpha=0.9))
    ax.annotate(f'Elevator\nMin: {elevator_only_min_years:.0f}y', 
                xy=(elevator_only_min_years, y_max * 0.72), fontsize=9, ha='center',
                bbox=dict(boxstyle='round', facecolor='#C8DBC8', edgecolor=color_elevator, alpha=0.9))
    ax.annotate(f'Rocket\nMin: {rocket_only_min_years:.0f}y', 
                xy=(rocket_only_min_years, y_max * 0.55), fontsize=9, ha='center',
                bbox=dict(boxstyle='round', facecolor='#E8C8C5', edgecolor=color_rocket, alpha=0.9))
    
    ax.set_xlabel('Completion Time (years)', fontsize=11)
    ax.set_ylabel('Total Energy Consumption (PJ)', fontsize=11)
    ax.legend(loc='upper left', fontsize=9)
    ax.grid(True, alpha=0.3)
    ax.set_xlim(year_min, year_max)
    
    plt.tight_layout()
    
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"\n三方案对比图已保存至: {save_path}")
    
    # 打印关键数据
    print("\n" + "=" * 70)
    print("THREE SCENARIOS ENERGY COMPARISON (100-300 years)")
    print("=" * 70)
    print(f"\n{'Year':<10} {'Combined':>15} {'Elevator':>15} {'Rocket':>15}")
    print(f"{'':10} {'(PJ)':>15} {'(PJ)':>15} {'(PJ)':>15}")
    print("-" * 55)
    
    for ky in [110, 130, 150, 186, 200, 220, 250, 300]:
        idx = np.argmin(np.abs(years - ky))
        c = f'{combined_energy[idx]:.0f}' if not np.isnan(combined_energy[idx]) else 'N/A'
        e = f'{elevator_energy[idx]:.0f}' if not np.isnan(elevator_energy[idx]) else 'N/A'
        r = f'{rocket_energy[idx]:.0f}' if not np.isnan(rocket_energy[idx]) else 'N/A'
        print(f"{ky:<10} {c:>15} {e:>15} {r:>15}")
    
    print("=" * 70)
    
    return fig


# ============== 主程序 ==============

if __name__ == "__main__":
    print("=" * 90)
    print("MOON COLONY: COMBINED SCENARIO ANALYSIS")
    print("(Space Elevator System + Traditional Rockets)")
    print("=" * 90)
    
    # 主分析
    df = plot_combined_analysis()
    
    # 打印摘要
    print_summary(df)
    
    # 方案对比
    print("\n\n正在生成方案对比图...")
    comparison = plot_scenario_comparison()
    
    print("\n" + "=" * 90)
    print("FINAL COMPARISON")
    print("=" * 90)
    print(f"\n{'Scenario':<25} {'Years':>12} {'Energy (PJ)':>15} {'vs Rocket Only':>18}")
    print("-" * 70)
    for name, data in comparison.items():
        rocket_energy_val = comparison['rocket_only']['energy_PJ']
        savings = (1 - data['energy_PJ'] / rocket_energy_val) * 100 if rocket_energy_val > 0 else 0
        print(f"{name:<25} {data['years']:>12.0f} {data['energy_PJ']:>15.0f} {savings:>17.0f}%")
    print("=" * 90)
    
    # 生成三方案对比折线图 (100-250年)
    print("\n\n正在生成三方案对比折线图 (100-250年)...")
    plot_three_scenarios_comparison(year_min=100, year_max=250)
