"""
Moon Colony Construction - Rocket Launch Scenario Analysis

Problem: Deliver 100 million metric tons to Moon using traditional rockets
- 10 launch sites, each max 1 launch per day
- Rocket payload: 100-150 metric tons (use 125 tons average)
- Optimize launch distribution to minimize energy consumption
- Plot completion year vs total energy consumption

Key insight: As completion timeline extends, more launches can be allocated 
to low-latitude sites, reducing the latitude-induced energy penalty.
"""

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
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 = 100e9 kg 的物资，但单位是metric tons
PAYLOAD_PER_LAUNCH = 125  # metric tons per launch (average of 100-150)
TOTAL_LAUNCHES_NEEDED = TOTAL_PAYLOAD / PAYLOAD_PER_LAUNCH  # 800,000 launches

# 火箭参数
ISP = 363  # 比冲 (秒) - 液氧甲烷 (LOX/CH4, Raptor-class)
SPECIFIC_FUEL_ENERGY = 11.9e6  # J/kg 燃料比能量
ALPHA = 0.10  # 结构系数
NUM_STAGES = 3  # 3级火箭

# 地面发射基准 ΔV (赤道发射到月球轨道)
DELTA_V_BASE = 13300  # m/s (LEO + TLI + LOI)


# ============== 发射场定义 ==============
@dataclass
class LaunchSite:
    name: str
    short_name: str
    latitude: float  # degrees
    max_launches_per_day: int = 1
    
    @property
    def abs_latitude(self) -> float:
        return abs(self.latitude)
    
    @property
    def rotation_velocity(self) -> float:
        """地球自转速度 (m/s)"""
        return OMEGA_EARTH * R_EARTH * np.cos(np.radians(self.abs_latitude))
    
    @property
    def delta_v_loss(self) -> float:
        """相对赤道的ΔV损失 (m/s)"""
        v_equator = OMEGA_EARTH * R_EARTH  # ~465 m/s
        return v_equator - self.rotation_velocity
    
    @property
    def total_delta_v(self) -> float:
        """总ΔV需求 (m/s)"""
        return DELTA_V_BASE + self.delta_v_loss


# 10个发射场 (按纬度排序)
LAUNCH_SITES = [
    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),
]

# 按纬度排序
LAUNCH_SITES = sorted(LAUNCH_SITES, key=lambda x: x.abs_latitude)


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

def fuel_ratio_multistage(delta_v: float, isp: float = ISP, alpha: float = ALPHA, 
                          num_stages: int = NUM_STAGES) -> 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 energy_per_launch(site: LaunchSite, payload_tons: float = PAYLOAD_PER_LAUNCH) -> float:
    """
    计算单次发射的能量消耗 (Joules)
    
    参数:
        site: 发射场
        payload_tons: 载荷 (metric tons)
    返回:
        能量消耗 (J)
    """
    delta_v = site.total_delta_v
    k = fuel_ratio_multistage(delta_v)
    
    payload_kg = payload_tons * 1000  # convert to kg
    fuel_kg = k * payload_kg
    
    energy = fuel_kg * SPECIFIC_FUEL_ENERGY
    return energy


def fuel_per_launch(site: LaunchSite, payload_tons: float = PAYLOAD_PER_LAUNCH) -> float:
    """
    计算单次发射的燃料消耗 (metric tons)
    """
    delta_v = site.total_delta_v
    k = fuel_ratio_multistage(delta_v)
    return k * payload_tons


# ============== 发射排布优化 ==============

def calculate_launch_distribution(
    total_launches: int,
    completion_years: float,
    sites: List[LaunchSite] = LAUNCH_SITES
) -> List[Tuple[LaunchSite, int]]:
    """
    计算给定完成年限下的发射分配
    
    策略：优先使用低纬度发射场，直到达到其最大容量
    
    参数:
        total_launches: 总发射次数
        completion_years: 完成年限
        sites: 发射场列表 (已按纬度排序)
    
    返回:
        [(发射场, 分配的发射次数), ...]
    """
    days_available = completion_years * 365
    max_launches_per_site = days_available  # 每站每天最多1发
    
    distribution = []
    remaining_launches = total_launches
    
    # 按纬度从低到高分配
    for site in sites:
        if remaining_launches <= 0:
            distribution.append((site, 0))
        else:
            allocated = min(remaining_launches, max_launches_per_site)
            distribution.append((site, int(allocated)))
            remaining_launches -= allocated
    
    return distribution


def calculate_total_energy(distribution: List[Tuple[LaunchSite, int]]) -> dict:
    """
    计算发射分配的总能量消耗
    
    返回:
        包含各项统计的字典
    """
    total_energy = 0
    total_fuel = 0
    total_launches = 0
    site_details = []
    
    for site, num_launches in distribution:
        if num_launches > 0:
            energy = energy_per_launch(site) * num_launches
            fuel = fuel_per_launch(site) * num_launches
            total_energy += energy
            total_fuel += fuel
            total_launches += num_launches
            
            site_details.append({
                'site': site.name,
                'short_name': site.short_name,
                'latitude': site.abs_latitude,
                'launches': num_launches,
                'energy_PJ': energy / 1e15,  # PetaJoules
                'fuel_Mt': fuel / 1e6,  # Million metric tons
            })
    
    return {
        'total_energy_PJ': total_energy / 1e15,
        'total_fuel_Mt': total_fuel / 1e6,
        'total_launches': total_launches,
        'sites': site_details
    }


def analyze_completion_timeline(
    years_range: np.ndarray,
    total_payload: float = TOTAL_PAYLOAD,
    payload_per_launch: float = PAYLOAD_PER_LAUNCH
) -> pd.DataFrame:
    """
    分析不同完成年限下的能量消耗
    
    返回:
        DataFrame with analysis results
    """
    total_launches = int(total_payload / payload_per_launch)
    
    results = []
    for years in years_range:
        distribution = calculate_launch_distribution(total_launches, years)
        stats = calculate_total_energy(distribution)
        
        # 计算使用了多少个发射场
        sites_used = sum(1 for _, n in distribution if n > 0)
        
        # 计算低纬度发射场的比例
        low_lat_launches = sum(n for site, n in distribution if site.abs_latitude < 30)
        low_lat_ratio = low_lat_launches / total_launches if total_launches > 0 else 0
        
        results.append({
            'years': years,
            'total_launches': total_launches,
            'sites_used': sites_used,
            'total_energy_PJ': stats['total_energy_PJ'],
            'total_fuel_Mt': stats['total_fuel_Mt'],
            'low_lat_ratio': low_lat_ratio,
            'energy_per_launch_TJ': stats['total_energy_PJ'] * 1000 / total_launches,
        })
    
    return pd.DataFrame(results)


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

def plot_timeline_analysis(
    save_path: str = '/Volumes/Files/code/mm/20260130_b/rocket_launch_timeline.png'
):
    """
    绘制完成年份-能量消耗关系图
    """
    # 计算最短完成时间
    total_launches = int(TOTAL_PAYLOAD / PAYLOAD_PER_LAUNCH)
    num_sites = len(LAUNCH_SITES)
    min_years = total_launches / (num_sites * 365)
    
    print(f"Total payload: {TOTAL_PAYLOAD/1e6:.0f} million metric tons")
    print(f"Payload per launch: {PAYLOAD_PER_LAUNCH} metric tons")
    print(f"Total launches needed: {total_launches:,}")
    print(f"Number of launch sites: {num_sites}")
    print(f"Minimum completion time: {min_years:.1f} years")
    
    # 分析范围：从最短时间到5倍最短时间
    years_range = np.linspace(min_years, min_years * 5, 100)
    
    # 运行分析
    df = analyze_completion_timeline(years_range)
    
    # 创建图表
    fig, axes = plt.subplots(2, 2, figsize=(16, 14))
    
    # ========== 图1: 完成年份 vs 总能量消耗 ==========
    ax1 = axes[0, 0]
    ax1.plot(df['years'], df['total_energy_PJ'], 'b-', linewidth=2)
    ax1.fill_between(df['years'], df['total_energy_PJ'], alpha=0.3)
    
    # 标记关键点
    min_idx = df['total_energy_PJ'].idxmin()
    ax1.axvline(x=df.loc[min_idx, 'years'], color='r', linestyle='--', 
                label=f'Optimal: {df.loc[min_idx, "years"]:.1f} years')
    ax1.axhline(y=df.loc[min_idx, 'total_energy_PJ'], color='r', linestyle=':', alpha=0.5)
    
    # 标记最短时间点
    ax1.axvline(x=min_years, color='g', linestyle='--', 
                label=f'Minimum time: {min_years:.1f} years')
    
    ax1.set_xlabel('Completion Time (years)', fontsize=12)
    ax1.set_ylabel('Total Energy Consumption (PJ)', fontsize=12)
    ax1.set_title('Completion Timeline vs Total Energy\n(100 Million Metric Tons to Moon)', fontsize=14)
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # ========== 图2: 使用的发射场数量 ==========
    ax2 = axes[0, 1]
    ax2.plot(df['years'], df['sites_used'], 'g-', linewidth=2, marker='o', 
             markevery=10, markersize=6)
    
    ax2.set_xlabel('Completion Time (years)', fontsize=12)
    ax2.set_ylabel('Number of Launch Sites Used', fontsize=12)
    ax2.set_title('Launch Sites Required vs Timeline', fontsize=14)
    ax2.set_ylim(0, 11)
    ax2.grid(True, alpha=0.3)
    
    # 添加发射场名称标注
    for i, site in enumerate(LAUNCH_SITES):
        ax2.axhline(y=i+1, color='gray', linestyle=':', alpha=0.3)
        ax2.text(df['years'].max() * 0.98, i+1 + 0.1, site.short_name, 
                fontsize=8, ha='right', va='bottom')
    
    # ========== 图3: 低纬度发射比例 ==========
    ax3 = axes[1, 0]
    ax3.plot(df['years'], df['low_lat_ratio'] * 100, 'm-', linewidth=2)
    ax3.fill_between(df['years'], df['low_lat_ratio'] * 100, alpha=0.3, color='magenta')
    
    ax3.set_xlabel('Completion Time (years)', fontsize=12)
    ax3.set_ylabel('Low Latitude (<30°) Launch Ratio (%)', fontsize=12)
    ax3.set_title('Proportion of Launches from Low-Latitude Sites', fontsize=14)
    ax3.set_ylim(0, 105)
    ax3.grid(True, alpha=0.3)
    
    # ========== 图4: 平均每发能量 ==========
    ax4 = axes[1, 1]
    ax4.plot(df['years'], df['energy_per_launch_TJ'], 'r-', linewidth=2)
    
    # 计算赤道和最高纬度的能量参考线
    energy_equator = energy_per_launch(LAUNCH_SITES[0]) / 1e12  # TJ
    energy_max_lat = energy_per_launch(LAUNCH_SITES[-1]) / 1e12  # TJ
    
    ax4.axhline(y=energy_equator, color='g', linestyle='--', 
                label=f'Equator baseline: {energy_equator:.1f} TJ')
    ax4.axhline(y=energy_max_lat, color='orange', linestyle='--', 
                label=f'Alaska (57°): {energy_max_lat:.1f} TJ')
    
    ax4.set_xlabel('Completion Time (years)', fontsize=12)
    ax4.set_ylabel('Average Energy per Launch (TJ)', fontsize=12)
    ax4.set_title('Energy Efficiency vs Timeline\n(Longer timeline → more low-lat launches)', fontsize=14)
    ax4.legend()
    ax4.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"\n图表已保存至: {save_path}")
    
    return df, fig


def plot_launch_distribution(
    years: float,
    save_path: str = '/Volumes/Files/code/mm/20260130_b/launch_distribution.png'
):
    """
    绘制特定年限下的发射分配详情
    """
    total_launches = int(TOTAL_PAYLOAD / PAYLOAD_PER_LAUNCH)
    distribution = calculate_launch_distribution(total_launches, years)
    stats = calculate_total_energy(distribution)
    
    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
    
    # ========== 图1: 发射次数分配 ==========
    ax1 = axes[0]
    sites = [s.short_name for s, _ in distribution]
    launches = [n for _, n in distribution]
    latitudes = [s.abs_latitude for s, _ in distribution]
    
    colors = plt.cm.RdYlGn_r(np.array(latitudes) / 60)
    bars = ax1.barh(sites, launches, color=colors)
    
    ax1.set_xlabel('Number of Launches', fontsize=12)
    ax1.set_ylabel('Launch Site', fontsize=12)
    ax1.set_title(f'Launch Distribution ({years:.1f} years completion)\n'
                  f'Total: {total_launches:,} launches', fontsize=13)
    ax1.grid(True, alpha=0.3, axis='x')
    
    # 添加数值标签
    for bar, n in zip(bars, launches):
        if n > 0:
            ax1.text(bar.get_width() + 1000, bar.get_y() + bar.get_height()/2,
                    f'{n:,}', va='center', fontsize=9)
    
    # ========== 图2: 能量消耗分配 ==========
    ax2 = axes[1]
    
    site_names = [d['short_name'] for d in stats['sites']]
    energies = [d['energy_PJ'] for d in stats['sites']]
    
    colors = plt.cm.RdYlGn_r(np.array([d['latitude'] for d in stats['sites']]) / 60)
    bars = ax2.barh(site_names, energies, color=colors)
    
    ax2.set_xlabel('Energy Consumption (PJ)', fontsize=12)
    ax2.set_ylabel('Launch Site', fontsize=12)
    ax2.set_title(f'Energy Distribution\n'
                  f'Total: {stats["total_energy_PJ"]:.1f} PJ', fontsize=13)
    ax2.grid(True, alpha=0.3, axis='x')
    
    # 添加数值标签
    for bar, e in zip(bars, energies):
        if e > 0:
            ax2.text(bar.get_width() + 0.5, bar.get_y() + bar.get_height()/2,
                    f'{e:.1f}', va='center', fontsize=9)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"分配详情图已保存至: {save_path}")
    
    return distribution, stats


def print_analysis_summary(df: pd.DataFrame):
    """打印分析摘要"""
    print("\n" + "=" * 80)
    print("ROCKET LAUNCH SCENARIO ANALYSIS SUMMARY")
    print("=" * 80)
    
    total_launches = int(TOTAL_PAYLOAD / PAYLOAD_PER_LAUNCH)
    
    print(f"\nMission Parameters:")
    print(f"  - Total payload: {TOTAL_PAYLOAD/1e6:.0f} million metric tons")
    print(f"  - Payload per launch: {PAYLOAD_PER_LAUNCH} metric tons")
    print(f"  - Total launches required: {total_launches:,}")
    print(f"  - Number of launch sites: {len(LAUNCH_SITES)}")
    
    # 最短时间场景
    min_years = df['years'].min()
    min_time_row = df[df['years'] == min_years].iloc[0]
    print(f"\nMinimum Time Scenario ({min_years:.1f} years):")
    print(f"  - All {int(min_time_row['sites_used'])} sites at full capacity")
    print(f"  - Total energy: {min_time_row['total_energy_PJ']:.1f} PJ")
    print(f"  - Total fuel: {min_time_row['total_fuel_Mt']:.1f} million metric tons")
    print(f"  - Low-latitude ratio: {min_time_row['low_lat_ratio']*100:.1f}%")
    
    # 最优能量场景 (只使用最低纬度站点)
    max_years = df['years'].max()
    max_time_row = df[df['years'] == max_years].iloc[0]
    print(f"\nExtended Timeline Scenario ({max_years:.1f} years):")
    print(f"  - Using {int(max_time_row['sites_used'])} lowest-latitude sites")
    print(f"  - Total energy: {max_time_row['total_energy_PJ']:.1f} PJ")
    print(f"  - Total fuel: {max_time_row['total_fuel_Mt']:.1f} million metric tons")
    print(f"  - Low-latitude ratio: {max_time_row['low_lat_ratio']*100:.1f}%")
    
    # 能量节省
    energy_saving = (1 - max_time_row['total_energy_PJ'] / min_time_row['total_energy_PJ']) * 100
    print(f"\nEnergy Savings (Extended vs Minimum time):")
    print(f"  - Energy reduction: {energy_saving:.1f}%")
    print(f"  - Time trade-off: {max_years - min_years:.1f} additional years")
    
    print("\n" + "=" * 80)
    print("Launch Site Details:")
    print("=" * 80)
    print(f"\n{'Site':<25} {'Latitude':>10} {'ΔV Loss':>10} {'Fuel Ratio':>12} {'Energy/Launch':>14}")
    print(f"{'':25} {'(°)':>10} {'(m/s)':>10} {'':>12} {'(TJ)':>14}")
    print("-" * 80)
    
    for site in LAUNCH_SITES:
        k = fuel_ratio_multistage(site.total_delta_v)
        e = energy_per_launch(site) / 1e12
        print(f"{site.name:<25} {site.abs_latitude:>10.1f} {site.delta_v_loss:>10.0f} {k:>12.1f} {e:>14.1f}")
    
    print("=" * 80)


def analyze_launch_frequency_scenarios(
    save_path: str = '/Volumes/Files/code/mm/20260130_b/launch_frequency_analysis.png'
):
    """
    分析不同发射频率场景下的完成时间和能量消耗
    """
    total_launches = int(TOTAL_PAYLOAD / PAYLOAD_PER_LAUNCH)
    num_sites = len(LAUNCH_SITES)
    
    # 不同发射频率场景
    frequencies = [1, 2, 4, 8, 12, 24]  # launches per day per site
    
    results = []
    for freq in frequencies:
        # 最短完成时间
        min_years = total_launches / (num_sites * freq * 365)
        
        # 分析该时间下的能量（所有站点满负荷）
        distribution = calculate_launch_distribution(total_launches, min_years)
        stats = calculate_total_energy(distribution)
        
        results.append({
            'frequency': freq,
            'min_years': min_years,
            'total_energy_PJ': stats['total_energy_PJ'],
            'total_fuel_Mt': stats['total_fuel_Mt'],
        })
    
    df = pd.DataFrame(results)
    
    # 绘图
    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
    
    # 图1: 发射频率 vs 完成时间
    ax1 = axes[0]
    ax1.bar(range(len(frequencies)), df['min_years'], color='steelblue', alpha=0.8)
    ax1.set_xticks(range(len(frequencies)))
    ax1.set_xticklabels([f'{f}/day' for f in frequencies])
    ax1.set_xlabel('Launch Frequency (per site)', fontsize=12)
    ax1.set_ylabel('Minimum Completion Time (years)', fontsize=12)
    ax1.set_title('Completion Time vs Launch Frequency\n(10 sites, 800,000 launches)', fontsize=13)
    ax1.grid(True, alpha=0.3, axis='y')
    
    # 添加数值标签
    for i, (freq, years) in enumerate(zip(frequencies, df['min_years'])):
        ax1.text(i, years + 2, f'{years:.1f}y', ha='center', fontsize=10, fontweight='bold')
    
    # 图2: 能量消耗比较
    ax2 = axes[1]
    ax2.bar(range(len(frequencies)), df['total_energy_PJ'], color='coral', alpha=0.8)
    ax2.set_xticks(range(len(frequencies)))
    ax2.set_xticklabels([f'{f}/day' for f in frequencies])
    ax2.set_xlabel('Launch Frequency (per site)', fontsize=12)
    ax2.set_ylabel('Total Energy (PJ)', fontsize=12)
    ax2.set_title('Total Energy Consumption\n(Same for all frequencies at max capacity)', fontsize=13)
    ax2.grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"\n发射频率分析图已保存至: {save_path}")
    
    return df


def plot_comprehensive_analysis(
    save_path: str = '/Volumes/Files/code/mm/20260130_b/rocket_comprehensive.png'
):
    """
    综合分析图：不同时间线下的能量消耗和发射分配
    """
    total_launches = int(TOTAL_PAYLOAD / PAYLOAD_PER_LAUNCH)
    num_sites = len(LAUNCH_SITES)
    min_years = total_launches / (num_sites * 365)
    
    # 选择几个关键时间点
    key_years = [min_years, min_years*1.5, min_years*2, min_years*3, min_years*5]
    
    fig, axes = plt.subplots(2, 3, figsize=(18, 12))
    
    # 第一行：能量趋势
    ax_main = axes[0, 0]
    years_range = np.linspace(min_years, min_years * 5, 100)
    df = analyze_completion_timeline(years_range)
    
    ax_main.plot(df['years'], df['total_energy_PJ'], 'b-', linewidth=2)
    ax_main.fill_between(df['years'], df['total_energy_PJ'], alpha=0.3)
    
    # 标记关键点
    for y in key_years[1:]:
        ax_main.axvline(x=y, color='gray', linestyle='--', alpha=0.5)
    
    ax_main.set_xlabel('Completion Time (years)', fontsize=11)
    ax_main.set_ylabel('Total Energy (PJ)', fontsize=11)
    ax_main.set_title('Energy vs Timeline\n(Longer time → Lower energy)', fontsize=12)
    ax_main.grid(True, alpha=0.3)
    
    # 能量效率
    ax_eff = axes[0, 1]
    ax_eff.plot(df['years'], df['energy_per_launch_TJ'], 'r-', linewidth=2)
    ax_eff.set_xlabel('Completion Time (years)', fontsize=11)
    ax_eff.set_ylabel('Energy per Launch (TJ)', fontsize=11)
    ax_eff.set_title('Average Energy Efficiency', fontsize=12)
    ax_eff.grid(True, alpha=0.3)
    
    # 发射场使用数量
    ax_sites = axes[0, 2]
    ax_sites.plot(df['years'], df['sites_used'], 'g-', linewidth=2)
    ax_sites.set_xlabel('Completion Time (years)', fontsize=11)
    ax_sites.set_ylabel('Sites Used', fontsize=11)
    ax_sites.set_title('Number of Active Sites', fontsize=12)
    ax_sites.set_ylim(0, 11)
    ax_sites.grid(True, alpha=0.3)
    
    # 第二行：三个时间点的发射分配饼图
    scenario_years = [min_years, min_years*2, min_years*5]
    scenario_names = ['Minimum Time', '2x Minimum', '5x Minimum']
    
    for idx, (years, name) in enumerate(zip(scenario_years, scenario_names)):
        ax = axes[1, idx]
        distribution = calculate_launch_distribution(total_launches, years)
        
        # 只显示有发射的站点
        labels = [s.short_name for s, n in distribution if n > 0]
        sizes = [n for _, n in distribution if n > 0]
        latitudes = [s.abs_latitude for s, n in distribution if n > 0]
        
        colors = plt.cm.RdYlGn_r(np.array(latitudes) / 60)
        
        wedges, texts, autotexts = ax.pie(sizes, labels=labels, autopct='%1.0f%%',
                                          colors=colors, pctdistance=0.75)
        
        # 计算该场景的能量
        stats = calculate_total_energy(distribution)
        
        ax.set_title(f'{name}\n({years:.0f} years, {stats["total_energy_PJ"]:.0f} PJ)', 
                    fontsize=11)
    
    plt.suptitle(f'Rocket Launch Scenario: {TOTAL_PAYLOAD/1e6:.0f}M tons to Moon\n'
                 f'({total_launches:,} launches needed)', fontsize=14, y=1.02)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"\n综合分析图已保存至: {save_path}")
    
    return fig


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

if __name__ == "__main__":
    print("=" * 80)
    print("MOON COLONY ROCKET LAUNCH SCENARIO ANALYSIS")
    print("Mission: Deliver 100 million metric tons to Moon")
    print("Constraint: 10 launch sites, max 1 launch/day each")
    print("=" * 80)
    
    # 运行时间线分析
    df, fig = plot_timeline_analysis()
    
    # 打印摘要
    print_analysis_summary(df)
    
    # 绘制综合分析图
    plot_comprehensive_analysis()
    
    # 发射频率场景分析
    print("\n" + "=" * 80)
    print("LAUNCH FREQUENCY SENSITIVITY ANALYSIS")
    print("=" * 80)
    freq_df = analyze_launch_frequency_scenarios()
    print("\nIf launch frequency can be increased:")
    print(f"{'Frequency':>15} {'Min Years':>15} {'Energy (PJ)':>15}")
    print("-" * 50)
    for _, row in freq_df.iterrows():
        print(f"{row['frequency']:>12}/day {row['min_years']:>15.1f} {row['total_energy_PJ']:>15.0f}")
    
    # 关键结论
    print("\n" + "=" * 80)
    print("KEY FINDINGS")
    print("=" * 80)
    print(f"""
    1. MINIMUM COMPLETION TIME: {df['years'].min():.0f} years
       (All 10 sites at 1 launch/day capacity)
    
    2. TOTAL ENERGY REQUIRED: ~{df['total_energy_PJ'].mean():.0f} PJ
       (~{df['total_fuel_Mt'].mean():.0f} million metric tons of fuel)
    
    3. LATITUDE EFFECT: 
       - Energy savings from using only low-lat sites: ~2-3%
       - Trade-off: Significantly longer timeline
    
    4. CONCLUSION: 
       Traditional rockets alone are IMPRACTICAL for this mission.
       219+ years is far beyond reasonable project timeline.
       Space Elevator System is essential for feasibility.
    """)