"""
地面发射 vs 太空电梯发射：比能量计算与对比

可配置参数进行不同场景的能量分析
"""

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 Optional

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


# ============== 物理常数 ==============
@dataclass
class PhysicalConstants:
    """物理常数"""
    G: float = 6.674e-11          # 万有引力常数 (m³/kg/s²)
    M_earth: float = 5.972e24     # 地球质量 (kg)
    M_moon: float = 7.342e22      # 月球质量 (kg)
    R_earth: float = 6.371e6      # 地球半径 (m)
    R_moon: float = 1.737e6       # 月球半径 (m)
    g0: float = 9.81              # 地表重力加速度 (m/s²)
    d_earth_moon: float = 3.844e8 # 地月平均距离 (m)
    omega_earth: float = 7.27e-5  # 地球自转角速度 (rad/s)
    
    @property
    def GM_earth(self) -> float:
        """地球引力参数"""
        return self.G * self.M_earth
    
    @property
    def GM_moon(self) -> float:
        """月球引力参数"""
        return self.G * self.M_moon


# ============== 发动机参数 ==============
@dataclass
class EngineParams:
    """发动机/推进系统参数"""
    name: str = "液氧甲烷"
    isp: float = 363              # 比冲 (秒) - LOX/CH4 (Raptor-class)
    specific_energy: float = 11.9e6  # 燃料比能量 (J/kg)
    
    @property
    def exhaust_velocity(self) -> float:
        """排气速度 (m/s)"""
        return self.isp * 9.81


# ============== 发射场参数 ==============
@dataclass
class LaunchSite:
    """发射场参数"""
    name: str = "赤道发射场"
    latitude: float = 0.0           # 纬度 (度)
    
    # 常见发射场纬度参考:
    # 赤道: 0°
    # 文昌 (中国): 19.6°
    # 卡纳维拉尔角 (美国): 28.5°
    # 拜科努尔 (哈萨克斯坦): 45.6°
    # 酒泉 (中国): 40.9°
    # 普列谢茨克 (俄罗斯): 62.9°
    
    @property
    def latitude_rad(self) -> float:
        """纬度 (弧度)"""
        return np.radians(self.latitude)


# 预定义的发射场
LAUNCH_SITES = {
    # 赤道参考
    '赤道': LaunchSite("Equator", 0.0),
    
    # 法属圭亚那 (接近赤道)
    'Kourou': LaunchSite("French Guiana", 5.2),
    
    # 印度
    'SDSC': LaunchSite("Satish Dhawan", 13.7),
    
    # 美国
    'Texas': LaunchSite("Texas", 26.0),
    'Florida': LaunchSite("Florida", 28.5),
    'California': LaunchSite("California", 34.7),
    'Virginia': LaunchSite("Virginia", 37.8),
    'Alaska': LaunchSite("Alaska", 57.4),
    
    # 中国
    'Taiyuan': LaunchSite("Taiyuan", 38.8),
    
    # 新西兰
    'Mahia': LaunchSite("Mahia", -39.3),  # 南半球
    
    # 哈萨克斯坦
    'Baikonur': LaunchSite("Kazakhstan", 45.6),
}


# ============== 任务参数 ==============
@dataclass
class MissionParams:
    """任务参数"""
    # 结构系数
    alpha: float = 0.10
    
    # 目标轨道 (环月轨道高度)
    lunar_orbit_altitude: float = 100e3  # 100 km
    
    # 地面发射参数 (赤道发射基准值)
    leo_altitude: float = 400e3          # LEO高度 400 km
    delta_v_ground_to_leo_base: float = 9400  # 赤道发射到LEO的基准ΔV (m/s)
    delta_v_tli: float = 3100            # 地月转移注入 (m/s)
    delta_v_loi: float = 800             # 月球轨道捕获 (m/s)
    
    # 目标轨道倾角 (度), 0 = 赤道轨道
    target_inclination: float = 0.0
    
    # 太空电梯参数
    elevator_length: float = 1.0e8       # 电梯长度 (m), 默认10万km
    
    @property
    def total_delta_v_ground(self) -> float:
        """地面发射总ΔV (赤道发射基准)"""
        return self.delta_v_ground_to_leo_base + self.delta_v_tli + self.delta_v_loi


# ============== 纬度相关计算函数 ==============

def earth_rotation_velocity(latitude: float, constants: PhysicalConstants = PhysicalConstants()) -> float:
    """
    计算地球表面某纬度的自转线速度
    
    参数:
        latitude: 纬度 (度)
        constants: 物理常数
    返回:
        自转线速度 (m/s)
    """
    lat_rad = np.radians(latitude)
    return constants.omega_earth * constants.R_earth * np.cos(lat_rad)


def delta_v_rotation_loss(latitude: float, constants: PhysicalConstants = PhysicalConstants()) -> float:
    """
    计算相对赤道发射损失的自转速度贡献
    
    参数:
        latitude: 发射场纬度 (度)
        constants: 物理常数
    返回:
        损失的ΔV (m/s)
    """
    v_equator = earth_rotation_velocity(0, constants)  # 赤道速度 ~465 m/s
    v_latitude = earth_rotation_velocity(latitude, constants)
    return v_equator - v_latitude


def delta_v_inclination_change(
    launch_latitude: float, 
    target_inclination: float,
    orbital_velocity: float = 7800  # LEO轨道速度 (m/s)
) -> float:
    """
    计算轨道倾角改变所需的ΔV
    
    从发射场纬度的轨道转到目标倾角轨道
    注意: 最小轨道倾角 = 发射场纬度
    
    参数:
        launch_latitude: 发射场纬度 (度)
        target_inclination: 目标轨道倾角 (度)
        orbital_velocity: 轨道速度 (m/s)
    返回:
        倾角改变ΔV (m/s)
    """
    # 发射场纬度决定最小可达倾角
    min_inclination = abs(launch_latitude)
    
    if target_inclination < min_inclination:
        # 需要进行纯倾角改变机动 (非常昂贵!)
        # ΔV = 2 * v * sin(Δi/2)
        delta_i = np.radians(min_inclination - target_inclination)
        return 2 * orbital_velocity * np.sin(delta_i / 2)
    else:
        # 可以通过发射方位角直接进入目标倾角，无额外ΔV
        return 0


def total_delta_v_with_latitude(
    launch_site: LaunchSite,
    mission: MissionParams,
    constants: PhysicalConstants = PhysicalConstants()
) -> dict:
    """
    计算考虑纬度因素的总ΔV
    
    参数:
        launch_site: 发射场
        mission: 任务参数
        constants: 物理常数
    返回:
        包含各项ΔV的字典
    """
    # 1. 自转速度损失
    dv_rotation_loss = delta_v_rotation_loss(launch_site.latitude, constants)
    
    # 2. 倾角改变损失 (如果目标是低倾角轨道)
    dv_inclination = delta_v_inclination_change(
        launch_site.latitude, 
        mission.target_inclination,
        orbital_velocity=7800
    )
    
    # 3. 总地面到LEO的ΔV
    dv_to_leo = mission.delta_v_ground_to_leo_base + dv_rotation_loss
    
    # 4. 总ΔV
    dv_total = dv_to_leo + dv_inclination + mission.delta_v_tli + mission.delta_v_loi
    
    return {
        'launch_site': launch_site.name,
        'latitude': launch_site.latitude,
        'rotation_velocity': earth_rotation_velocity(launch_site.latitude, constants),
        'dv_rotation_loss': dv_rotation_loss,
        'dv_inclination_change': dv_inclination,
        'dv_to_leo': dv_to_leo,
        'dv_tli': mission.delta_v_tli,
        'dv_loi': mission.delta_v_loi,
        'dv_total': dv_total
    }


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

def mass_ratio(delta_v: float, engine: EngineParams) -> float:
    """
    计算质量比 R = m0/mf
    
    参数:
        delta_v: 速度增量 (m/s)
        engine: 发动机参数
    返回:
        质量比 R
    """
    return np.exp(delta_v / engine.exhaust_velocity)


def fuel_to_payload_ratio(delta_v: float, engine: EngineParams, alpha: float = 0.1) -> float:
    """
    计算燃料与载荷的质量比（单级火箭）
    
    参数:
        delta_v: 速度增量 (m/s)
        engine: 发动机参数
        alpha: 结构系数
    返回:
        m_fuel / m_payload
    """
    R = mass_ratio(delta_v, engine)
    denominator = 1 - alpha * (R - 1)
    
    if denominator <= 0:
        return np.inf
    
    return (R - 1) / denominator


def fuel_to_payload_ratio_multistage(
    delta_v: float, 
    engine: EngineParams, 
    alpha: float = 0.1,
    num_stages: int = 3
) -> float:
    """
    计算燃料与载荷的质量比（多级火箭）
    
    假设各级均分ΔV，使用相同发动机和结构系数
    
    参数:
        delta_v: 总速度增量 (m/s)
        engine: 发动机参数
        alpha: 结构系数
        num_stages: 级数
    返回:
        m_fuel_total / m_payload
    """
    # 每级分配的ΔV
    delta_v_per_stage = delta_v / num_stages
    
    # 每级的质量比
    R_stage = mass_ratio(delta_v_per_stage, engine)
    
    # 每级的燃料比 (相对于该级有效载荷)
    denominator = 1 - alpha * (R_stage - 1)
    if denominator <= 0:
        return np.inf
    
    k_stage = (R_stage - 1) / denominator
    
    # 结构质量比 (相对于燃料)
    # 每级末质量 = 载荷 + 结构 = 载荷 + alpha * 燃料
    # 每级初质量 = 载荷 + 结构 + 燃料 = 载荷 * (1 + k_stage * (1 + alpha))
    
    # 对于多级火箭，总质量比是各级质量比的乘积
    # 第i级的"载荷"是第i+1级及以上所有质量
    
    # 简化计算：假设各级结构系数相同
    # 级比 = (1 + k_stage * (1 + alpha))
    stage_ratio = 1 + k_stage * (1 + alpha)
    
    # 总初始质量/最终载荷 = stage_ratio ^ num_stages
    total_ratio = stage_ratio ** num_stages
    
    # 总燃料 = 总初始质量 - 载荷 - 总结构
    # 近似：总燃料/载荷 ≈ total_ratio - 1 (忽略结构细节)
    
    # 更精确的计算
    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 ground_launch_specific_energy(
    engine: EngineParams,
    mission: MissionParams,
    constants: PhysicalConstants = PhysicalConstants(),
    num_stages: int = 3
) -> dict:
    """
    计算地面发射的比能量（使用多级火箭）
    
    参数:
        engine: 发动机参数
        mission: 任务参数
        constants: 物理常数
        num_stages: 火箭级数 (默认3级)
    
    返回:
        包含各项比能量的字典 (单位: J/kg 载荷)
    """
    # 总ΔV
    delta_v_total = mission.total_delta_v_ground
    
    # 燃料/载荷比 (使用多级火箭模型)
    k = fuel_to_payload_ratio_multistage(delta_v_total, engine, mission.alpha, num_stages)
    
    # 推进剂化学能 (J/kg 载荷)
    propellant_energy = k * engine.specific_energy
    
    # 载荷轨道能量增益
    r_moon_orbit = constants.R_moon + mission.lunar_orbit_altitude
    
    # 从地面静止到月球轨道的比能量变化
    E_initial = -constants.GM_earth / constants.R_earth
    E_final = -constants.GM_moon / (2 * r_moon_orbit) - constants.GM_earth / constants.d_earth_moon
    payload_orbital_energy = E_final - E_initial
    
    # 效率
    efficiency = payload_orbital_energy / propellant_energy if propellant_energy > 0 and not np.isinf(propellant_energy) else 0
    
    return {
        'delta_v': delta_v_total,
        'fuel_ratio': k,
        'num_stages': num_stages,
        'propellant_energy': propellant_energy,
        'payload_orbital_energy': payload_orbital_energy,
        'total_energy': propellant_energy,  # 地面发射全靠推进剂
        'efficiency': efficiency,
        'method': f'地面发射({num_stages}级)'
    }


def ground_launch_specific_energy_with_latitude(
    engine: EngineParams,
    mission: MissionParams,
    launch_site: LaunchSite,
    constants: PhysicalConstants = PhysicalConstants(),
    num_stages: int = 3
) -> dict:
    """
    计算考虑发射场纬度的地面发射比能量
    
    参数:
        engine: 发动机参数
        mission: 任务参数
        launch_site: 发射场
        constants: 物理常数
        num_stages: 火箭级数 (默认3级)
    
    返回:
        包含各项比能量的字典 (单位: J/kg 载荷)
    """
    # 计算纬度相关的ΔV
    dv_info = total_delta_v_with_latitude(launch_site, mission, constants)
    delta_v_total = dv_info['dv_total']
    
    # 燃料/载荷比 (使用多级火箭模型)
    k = fuel_to_payload_ratio_multistage(delta_v_total, engine, mission.alpha, num_stages)
    
    # 推进剂化学能 (J/kg 载荷)
    propellant_energy = k * engine.specific_energy
    
    # 载荷轨道能量增益
    r_moon_orbit = constants.R_moon + mission.lunar_orbit_altitude
    E_initial = -constants.GM_earth / constants.R_earth
    E_final = -constants.GM_moon / (2 * r_moon_orbit) - constants.GM_earth / constants.d_earth_moon
    payload_orbital_energy = E_final - E_initial
    
    # 效率
    efficiency = payload_orbital_energy / propellant_energy if propellant_energy > 0 and not np.isinf(propellant_energy) else 0
    
    return {
        'launch_site': launch_site.name,
        'latitude': launch_site.latitude,
        'rotation_velocity': dv_info['rotation_velocity'],
        'dv_rotation_loss': dv_info['dv_rotation_loss'],
        'dv_inclination_change': dv_info['dv_inclination_change'],
        'delta_v': delta_v_total,
        'fuel_ratio': k,
        'num_stages': num_stages,
        'propellant_energy': propellant_energy,
        'payload_orbital_energy': payload_orbital_energy,
        'total_energy': propellant_energy,
        'efficiency': efficiency,
        'method': f'地面发射({launch_site.name}, {num_stages}级)'
    }


def compare_launch_sites(
    engine: EngineParams = EngineParams(),
    mission: MissionParams = MissionParams(),
    constants: PhysicalConstants = PhysicalConstants(),
    num_stages: int = 3
) -> list:
    """
    比较不同发射场的燃料和能量需求
    
    返回:
        各发射场的分析结果列表 (按纬度绝对值排序)
    """
    results = []
    for name, site in LAUNCH_SITES.items():
        # 使用纬度绝对值进行计算（南北半球对称）
        abs_lat = abs(site.latitude)
        site_normalized = LaunchSite(site.name, abs_lat)
        result = ground_launch_specific_energy_with_latitude(
            engine, mission, site_normalized, constants, num_stages
        )
        result['original_latitude'] = site.latitude
        result['abs_latitude'] = abs_lat
        results.append(result)
    
    return sorted(results, key=lambda x: x['abs_latitude'])


def print_latitude_comparison(results: list):
    """打印不同纬度发射场的对比"""
    print("=" * 110)
    print("Launch Site Latitude Effects on Fuel and Energy")
    print("=" * 110)
    
    print(f"\n{'Launch Site':<30} {'Lat':>6} {'V_rot':>8} {'ΔV_loss':>8} {'ΔV_total':>9} {'Fuel/PL':>9} {'Energy':>10}")
    print(f"{'':30} {'(°)':>6} {'(m/s)':>8} {'(m/s)':>8} {'(km/s)':>9} {'Ratio':>9} {'(MJ/kg)':>10}")
    print("-" * 110)
    
    for r in results:
        lat = r.get('abs_latitude', abs(r['latitude']))
        print(f"{r['launch_site']:<30} {lat:>6.1f} {r['rotation_velocity']:>8.0f} {r['dv_rotation_loss']:>8.0f} {r['delta_v']/1000:>9.2f} {r['fuel_ratio']:>9.1f} {r['total_energy']/1e6:>10.1f}")
    
    print("-" * 110)
    
    # 计算相对赤道的损失
    equator_result = next((r for r in results if abs(r['latitude']) < 0.1), results[0])
    print(f"\nExtra consumption relative to Equator launch:")
    for r in results:
        lat = r.get('abs_latitude', abs(r['latitude']))
        if lat > 0.1:
            fuel_increase = (r['fuel_ratio'] / equator_result['fuel_ratio'] - 1) * 100
            energy_increase = (r['total_energy'] / equator_result['total_energy'] - 1) * 100
            print(f"  {r['launch_site']:<30}: Fuel +{fuel_increase:>5.1f}%, Energy +{energy_increase:>5.1f}%")
    
    print("=" * 110)


def plot_latitude_effects(
    engine: EngineParams = EngineParams(),
    mission: MissionParams = MissionParams(),
    constants: PhysicalConstants = PhysicalConstants(),
    save_path: str = '/Volumes/Files/code/mm/20260130_b/p1/latitude_effects.png',
    lunar_mission: bool = True
):
    """
    Plot fuel ratio vs launch latitude.
    
    Args:
        engine: Engine parameters
        mission: Mission parameters
        constants: Physical constants
        save_path: Output file path
        lunar_mission: If True, use lunar mission scenario (no orbital plane change penalty)
    """
    
    # Lunar mission scenario: set high target inclination to avoid plane change penalty
    if lunar_mission:
        mission_plot = MissionParams(
            alpha=mission.alpha,
            lunar_orbit_altitude=mission.lunar_orbit_altitude,
            delta_v_ground_to_leo_base=mission.delta_v_ground_to_leo_base,
            delta_v_tli=mission.delta_v_tli,
            delta_v_loi=mission.delta_v_loi,
            target_inclination=90.0,  # Allow any inclination
            elevator_length=mission.elevator_length
        )
        title_suffix = " (Lunar Mission)"
    else:
        mission_plot = mission
        title_suffix = " (Target: Equatorial Orbit)"
    
    # Golden ratio aspect: 缩小尺寸使字体相对更大
    fig, ax = plt.subplots(figsize=(8, 8 * 0.618))
    
    # 中低饱和度配色
    color_curve = '#52796F'      # 暗绿色 - 曲线
    color_point = '#8E7B9A'      # 灰紫色 - 数据点
    
    # Continuous latitude range
    latitudes = np.linspace(0, 65, 100)
    
    # Calculate fuel ratio for each latitude
    fuel_ratios = []
    for lat in latitudes:
        site = LaunchSite("temp", lat)
        result = ground_launch_specific_energy_with_latitude(engine, mission_plot, site, constants)
        fuel_ratios.append(result['fuel_ratio'])
    
    # Plot continuous curve
    ax.plot(latitudes, fuel_ratios, color=color_curve, linewidth=2, label='Fuel Ratio vs Latitude')
    
    # Mark launch sites with custom offsets to avoid overlap with curve
    # Offsets: (x, y) in points - positive y moves text up, positive x moves right
    # 纬度参考: Equator(0), French Guiana(5.2), Satish Dhawan(13.7), Texas(26), 
    #          Florida(28.5), California(34.7), Virginia(37.8), Taiyuan(38.8), 
    #          Mahia(39.3), Kazakhstan(45.6), Alaska(57.4)
    label_offsets = {
        '赤道': (5, -15),          # Below curve
        'Kourou': (-75, 15),       # Left above (French Guiana)
        'SDSC': (8, -15),          # Below right (Satish Dhawan)
        'Texas': (-38, -15),       # Left below
        'Florida': (8, -15),       # Right below
        'California': (8, -15),    # Right below
        'Virginia': (8, 12),       # Right above
        'Taiyuan': (-52, 12),      # Left above (避开California)
        'Mahia': (8, -18),         # Right below
        'Alaska': (8, 8),          # Right above
        'Baikonur': (8, 8),        # Right above (Kazakhstan)
    }
    
    for name, site in LAUNCH_SITES.items():
        abs_lat = abs(site.latitude)
        site_temp = LaunchSite(site.name, abs_lat)
        result = ground_launch_specific_energy_with_latitude(engine, mission_plot, site_temp, constants)
        ax.plot(abs_lat, result['fuel_ratio'], 'o', color=color_point, markersize=8)
        # Simplified label
        label = site.name.split('(')[0].strip()
        # Get custom offset for this site, default to (5, 8)
        offset = label_offsets.get(name, (5, 8))
        ax.annotate(label, (abs_lat, result['fuel_ratio']), 
                    textcoords="offset points", xytext=offset, fontsize=11)
    
    ax.set_xlabel('Latitude |φ| (°)', fontsize=12)
    ax.set_ylabel('Fuel / Payload Mass Ratio', fontsize=12)
    # ax.set_title(f'Fuel Requirement vs Launch Latitude{title_suffix}', fontsize=13)
    ax.grid(True, alpha=0.3)
    ax.set_xlim(-2, 65)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"Latitude effects plot saved to: {save_path}")
    
    return fig


# ============== 电梯发射比能量计算 ==============

def elevator_launch_specific_energy(
    engine: EngineParams,
    mission: MissionParams,
    constants: PhysicalConstants = PhysicalConstants(),
    release_height: Optional[float] = None
) -> dict:
    """
    计算太空电梯发射的比能量
    
    参数:
        engine: 发动机参数
        mission: 任务参数
        constants: 物理常数
        release_height: 释放高度 (m)，None则使用电梯顶端
    
    返回:
        包含各项比能量的字典 (单位: J/kg 载荷)
    """
    # 释放点参数
    if release_height is None:
        release_height = mission.elevator_length
    
    r_release = constants.R_earth + release_height
    
    # 释放点速度 (随地球自转)
    v_release = constants.omega_earth * r_release
    
    # 释放点的逃逸速度
    v_escape = np.sqrt(2 * constants.GM_earth / r_release)
    
    # 释放点的圆轨道速度
    v_circular = np.sqrt(constants.GM_earth / r_release)
    
    # 计算C3 (特征能量)
    C3 = v_release**2 - v_escape**2
    
    # 电梯提升能量 (从地面到释放点)
    # 包括势能变化和动能获得
    lift_energy = (
        constants.GM_earth / constants.R_earth 
        - constants.GM_earth / r_release 
        + 0.5 * v_release**2
    )
    
    # 计算所需ΔV
    if C3 > 0:
        # 已超过逃逸速度，需要减速
        # 目标：进入地月转移轨道 (C3 ≈ -2 km²/s²)
        C3_target = -2e6  # -2 km²/s² 转换为 m²/s²
        v_needed = np.sqrt(max(0, v_escape**2 + C3_target))
        delta_v_adjust = abs(v_release - v_needed)
        maneuver_type = "减速"
    else:
        # 未达逃逸速度，需要加速
        # 计算到达月球轨道所需的额外速度
        v_needed = np.sqrt(v_escape**2 - 2e6)  # 地月转移
        delta_v_adjust = v_needed - v_release
        maneuver_type = "加速"
    
    # 加上月球轨道捕获
    delta_v_total = abs(delta_v_adjust) + mission.delta_v_loi
    
    # 燃料/载荷比
    k = fuel_to_payload_ratio(delta_v_total, engine, mission.alpha)
    
    # 推进剂化学能
    propellant_energy = k * engine.specific_energy
    
    # 载荷轨道能量增益 (与地面发射相同)
    r_moon_orbit = constants.R_moon + mission.lunar_orbit_altitude
    E_initial = -constants.GM_earth / constants.R_earth
    E_final = -constants.GM_moon / (2 * r_moon_orbit) - constants.GM_earth / constants.d_earth_moon
    payload_orbital_energy = E_final - E_initial
    
    # 总能量 = 电梯提升 + 推进剂
    total_energy = lift_energy + propellant_energy
    
    # 效率
    efficiency = payload_orbital_energy / total_energy if total_energy > 0 else 0
    
    return {
        'release_height': release_height,
        'release_velocity': v_release,
        'escape_velocity': v_escape,
        'C3': C3,
        'maneuver_type': maneuver_type,
        'delta_v': delta_v_total,
        'fuel_ratio': k,
        'lift_energy': lift_energy,
        'propellant_energy': propellant_energy,
        'payload_orbital_energy': payload_orbital_energy,
        'total_energy': total_energy,
        'efficiency': efficiency,
        'method': '电梯发射'
    }


# ============== 对比分析函数 ==============

def compare_launch_methods(
    engine: EngineParams = EngineParams(),
    mission: MissionParams = MissionParams(),
    constants: PhysicalConstants = PhysicalConstants()
) -> dict:
    """
    对比两种发射方式
    
    返回:
        对比结果字典
    """
    ground = ground_launch_specific_energy(engine, mission, constants)
    elevator = elevator_launch_specific_energy(engine, mission, constants)
    
    return {
        'ground': ground,
        'elevator': elevator,
        'fuel_saving': 1 - elevator['fuel_ratio'] / ground['fuel_ratio'],
        'energy_saving': 1 - elevator['total_energy'] / ground['total_energy'],
        'delta_v_saving': 1 - elevator['delta_v'] / ground['delta_v']
    }


def print_comparison(comparison: dict):
    """打印对比结果"""
    ground = comparison['ground']
    elevator = comparison['elevator']
    
    ground_label = ground['method']
    elevator_label = elevator['method']
    
    print("=" * 70)
    print("地面发射 vs 太空电梯发射：比能量对比")
    print("=" * 70)
    
    print(f"\n{'参数':<25} {ground_label:>18} {elevator_label:>18}")
    print("-" * 70)
    
    print(f"{'ΔV (km/s)':<25} {ground['delta_v']/1000:>18.2f} {elevator['delta_v']/1000:>18.2f}")
    print(f"{'燃料/载荷比':<25} {ground['fuel_ratio']:>18.2f} {elevator['fuel_ratio']:>18.2f}")
    print(f"{'电梯提升能量 (MJ/kg)':<25} {'0':>18} {elevator['lift_energy']/1e6:>18.1f}")
    print(f"{'推进剂能量 (MJ/kg)':<25} {ground['propellant_energy']/1e6:>18.1f} {elevator['propellant_energy']/1e6:>18.1f}")
    print(f"{'总能量 (MJ/kg)':<25} {ground['total_energy']/1e6:>18.1f} {elevator['total_energy']/1e6:>18.1f}")
    print(f"{'载荷能量增益 (MJ/kg)':<25} {ground['payload_orbital_energy']/1e6:>18.1f} {elevator['payload_orbital_energy']/1e6:>18.1f}")
    print(f"{'系统效率':<25} {ground['efficiency']*100:>17.1f}% {elevator['efficiency']*100:>17.1f}%")
    
    print("\n" + "-" * 70)
    print(f"{'ΔV 节省':<25} {comparison['delta_v_saving']*100:>18.1f}%")
    print(f"{'燃料节省':<25} {comparison['fuel_saving']*100:>18.1f}%")
    print(f"{'总能量节省':<25} {comparison['energy_saving']*100:>18.1f}%")
    print("=" * 70)


def plot_comparison(
    engine: EngineParams = EngineParams(),
    mission: MissionParams = MissionParams(),
    constants: PhysicalConstants = PhysicalConstants(),
    save_path: str = '/Volumes/Files/code/mm/20260130_b/specific_energy_comparison.png'
):
    """绘制对比图表"""
    
    fig, axes = plt.subplots(2, 2, figsize=(14, 12))
    
    # ========== 图1: 不同电梯高度的ΔV需求 ==========
    ax1 = axes[0, 0]
    heights = np.linspace(3.6e7, 1.5e8, 100)  # 36,000 km 到 150,000 km
    
    delta_vs = []
    for h in heights:
        mission_temp = MissionParams(elevator_length=h)
        result = elevator_launch_specific_energy(engine, mission_temp, constants, h)
        delta_vs.append(result['delta_v'] / 1000)
    
    ax1.plot(heights/1e6, delta_vs, 'b-', linewidth=2, label='电梯发射')
    ax1.axhline(y=mission.total_delta_v_ground/1000, color='r', linestyle='--', 
                linewidth=2, label='地面发射')
    ax1.axvline(x=35.786, color='g', linestyle=':', label='GEO高度')
    
    ax1.set_xlabel('电梯释放高度 (千km)', fontsize=12)
    ax1.set_ylabel('ΔV 需求 (km/s)', fontsize=12)
    ax1.set_title('释放高度 vs ΔV需求', fontsize=14)
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # ========== 图2: 燃料/载荷比对比 ==========
    ax2 = axes[0, 1]
    
    fuel_ratios_elevator = []
    for h in heights:
        mission_temp = MissionParams(elevator_length=h)
        result = elevator_launch_specific_energy(engine, mission_temp, constants, h)
        fuel_ratios_elevator.append(result['fuel_ratio'])
    
    ground_result = ground_launch_specific_energy(engine, mission, constants, num_stages=3)
    
    ax2.plot(heights/1e6, fuel_ratios_elevator, 'b-', linewidth=2, label='电梯发射')
    ax2.axhline(y=ground_result['fuel_ratio'], color='r', linestyle='--', 
                linewidth=2, label='地面发射(3级)')
    
    ax2.set_xlabel('电梯释放高度 (千km)', fontsize=12)
    ax2.set_ylabel('燃料/载荷 质量比', fontsize=12)
    ax2.set_title('释放高度 vs 燃料效率', fontsize=14)
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    ax2.set_ylim(0, 30)
    
    # ========== 图3: 能量构成对比 (柱状图) ==========
    ax3 = axes[1, 0]
    
    comparison = compare_launch_methods(engine, mission, constants)
    
    categories = ['地面发射', '电梯发射']
    lift_energy = [0, comparison['elevator']['lift_energy']/1e6]
    prop_energy = [comparison['ground']['propellant_energy']/1e6, 
                   comparison['elevator']['propellant_energy']/1e6]
    
    x = np.arange(len(categories))
    width = 0.35
    
    bars1 = ax3.bar(x, lift_energy, width, label='电梯提升能量', color='green')
    bars2 = ax3.bar(x, prop_energy, width, bottom=lift_energy, label='推进剂能量', color='orange')
    
    ax3.set_ylabel('比能量 (MJ/kg)', fontsize=12)
    ax3.set_title('能量构成对比', fontsize=14)
    ax3.set_xticks(x)
    ax3.set_xticklabels(categories)
    ax3.legend()
    ax3.grid(True, alpha=0.3, axis='y')
    
    # 添加数值标签
    for i, (l, p) in enumerate(zip(lift_energy, prop_energy)):
        total = l + p
        ax3.text(i, total + 5, f'{total:.0f}', ha='center', fontsize=11, fontweight='bold')
    
    # ========== 图4: 不同发动机对比 ==========
    ax4 = axes[1, 1]
    
    engines = [
        EngineParams(name="固体火箭", isp=280, specific_energy=5e6),
        EngineParams(name="液氧煤油", isp=350, specific_energy=10.3e6),
        EngineParams(name="液氧甲烷", isp=363, specific_energy=11.9e6),
        EngineParams(name="液氧液氢", isp=450, specific_energy=15.5e6),
    ]
    
    x_pos = np.arange(len(engines))
    ground_ratios = []
    elevator_ratios = []
    
    for eng in engines:
        g = ground_launch_specific_energy(eng, mission, constants, num_stages=3)
        e = elevator_launch_specific_energy(eng, mission, constants)
        ground_ratios.append(min(g['fuel_ratio'], 100))  # 限制显示范围
        elevator_ratios.append(e['fuel_ratio'])
    
    width = 0.35
    ax4.bar(x_pos - width/2, ground_ratios, width, label='地面发射(3级)', color='red', alpha=0.7)
    ax4.bar(x_pos + width/2, elevator_ratios, width, label='电梯发射', color='blue', alpha=0.7)
    
    ax4.set_ylabel('燃料/载荷 质量比', fontsize=12)
    ax4.set_title('不同发动机的燃料效率对比', fontsize=14)
    ax4.set_xticks(x_pos)
    ax4.set_xticklabels([e.name for e in engines])
    ax4.legend()
    ax4.grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"图表已保存至: {save_path}")
    
    return fig


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

if __name__ == "__main__":
    # 可以在这里修改参数
    
    # 发动机参数 - 液氧甲烷 (LOX/CH4, Raptor-class)
    engine = EngineParams(
        name="液氧甲烷",
        isp=363,                    # 比冲 (秒)
        specific_energy=11.9e6     # 燃料比能量 (J/kg)
    )
    
    # 任务参数
    mission = MissionParams(
        alpha=0.10,                 # 结构系数
        lunar_orbit_altitude=100e3, # 环月轨道高度 (m)
        delta_v_ground_to_leo_base=9400, # 赤道发射到LEO的基准ΔV (m/s)
        delta_v_tli=3100,           # 地月转移注入 (m/s)
        delta_v_loi=800,            # 月球轨道捕获 (m/s)
        target_inclination=0.0,     # 目标轨道倾角 (度)
        elevator_length=1.0e8       # 电梯长度 10万km (m)
    )
    
    # 物理常数 (通常不需要修改)
    constants = PhysicalConstants()
    
    # ========== 1. 地面发射 vs 电梯发射对比 ==========
    comparison = compare_launch_methods(engine, mission, constants)
    print_comparison(comparison)
    
    # 绘制对比图表
    print("\n正在生成对比图表...")
    plot_comparison(engine, mission, constants)
    
    # 电梯发射详细信息
    print("\n" + "=" * 70)
    print("电梯发射详细参数:")
    print("=" * 70)
    elevator = comparison['elevator']
    print(f"  释放高度: {elevator['release_height']/1e6:.1f} 千km")
    print(f"  释放点速度: {elevator['release_velocity']/1000:.2f} km/s")
    print(f"  当地逃逸速度: {elevator['escape_velocity']/1000:.2f} km/s")
    print(f"  特征能量 C3: {elevator['C3']/1e6:.1f} km²/s²")
    print(f"  机动类型: {elevator['maneuver_type']}")
    
    # ========== 2. 发射场纬度影响分析 (月球任务) ==========
    print("\n")
    print("=" * 110)
    print("LUNAR MISSION: Launch Site Comparison")
    print("(No orbital plane change required - only rotation velocity loss)")
    print("=" * 110)
    
    # 月球任务配置（允许倾角与发射场纬度匹配）
    mission_lunar = MissionParams(
        alpha=mission.alpha,
        lunar_orbit_altitude=mission.lunar_orbit_altitude,
        delta_v_ground_to_leo_base=mission.delta_v_ground_to_leo_base,
        delta_v_tli=mission.delta_v_tli,
        delta_v_loi=mission.delta_v_loi,
        target_inclination=90.0,  # 允许任意倾角
        elevator_length=mission.elevator_length
    )
    
    latitude_results = compare_launch_sites(engine, mission_lunar, constants)
    print_latitude_comparison(latitude_results)
    
    # 绘制纬度影响图表 (月球任务模式)
    print("\n正在生成纬度影响图表 (Lunar Mission)...")
    plot_latitude_effects(engine, mission_lunar, constants, lunar_mission=True)
    
    # ========== 3. 公式总结 ==========
    print("\n" + "=" * 70)
    print("数学关系总结:")
    print("=" * 70)
    print("""
  地球自转速度:
    v(φ) = ω × R_E × cos(φ)
    
  自转速度损失:
    Δv_loss = v(0°) - v(φ) = 465 × (1 - cos(φ)) m/s
    
  纬度对燃料的影响:
    由于 ΔV 增加 → R = exp(ΔV/ve) 指数增长 → 燃料比指数增长
    
  关键结论:
    - 赤道发射: 自转贡献 465 m/s, 燃料最省
    - 高纬度发射: 损失自转贡献, 且可能需要额外倾角机动
    - 每损失 100 m/s 自转速度 ≈ 增加 ~3% 燃料需求
""")