"""
任务三：月球殖民地水需求与运输分析

Water Supply Analysis for Moon Colony

分析10万人月球殖民地一年期水需求，以及四种运输方案的成本和时间线。

水需求模型：
- 生存用水：2.5 L/人/天
- 卫生用水：0.4 × α (舒适度系数) L/人/天
- 医疗应急：5 L/次，每百人每日患病率 2%
- 水循环回收率：90%

运输方案：
1. 仅使用空间电梯
2. 仅使用火箭（10个发射场）
3. 混合方案（电梯 + 10个发射场）
4. 混合方案（电梯 + 低纬度发射场）
"""

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, Dict, 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

# ============== 殖民地参数 ==============
POPULATION = 100_000  # 人口

# ============== 水需求参数 ==============
@dataclass
class WaterParameters:
    """水需求参数"""
    # 生存用水 (L/人/天)
    survival_water: float = 2.5
    
    # 卫生用水基础量 (L/人/天)
    hygiene_water_base: float = 0.4
    
    # 舒适度系数 (1=最低, 5=奢侈)
    comfort_factor: float = 1.0
    
    # 医疗应急用水 (L/次)
    medical_water_per_event: float = 5.0
    
    # 每百人每日患病率
    sickness_rate_per_100: float = 0.02  # 2%
    
    # 水循环回收率
    recycle_rate: float = 0.90  # 90%
    
    # 安全缓冲天数
    safety_buffer_days: int = 30
    
    @property
    def hygiene_water(self) -> float:
        """卫生用水 (L/人/天)"""
        return self.hygiene_water_base * self.comfort_factor
    
    @property
    def daily_water_per_person(self) -> float:
        """每人每日循环用水 (L)"""
        return self.survival_water + self.hygiene_water
    
    @property
    def daily_medical_events(self) -> float:
        """每日医疗事件数"""
        return POPULATION * self.sickness_rate_per_100 / 100
    
    @property
    def daily_medical_water(self) -> float:
        """每日医疗用水 (L)"""
        return self.daily_medical_events * self.medical_water_per_event


# ============== 运输参数 ==============
# 太空电梯参数
NUM_ELEVATORS = 3
ELEVATOR_CAPACITY_PER_YEAR = 179_000  # 吨/年/部
TOTAL_ELEVATOR_CAPACITY = NUM_ELEVATORS * ELEVATOR_CAPACITY_PER_YEAR  # 537,000 吨/年
ELEVATOR_SPECIFIC_ENERGY = 157.2e9  # J/吨 (157.2 GJ/吨)

# 火箭参数
PAYLOAD_PER_LAUNCH = 125  # 吨/次
ISP = 450  # 比冲 (秒)
SPECIFIC_FUEL_ENERGY = 15.5e6  # J/kg
ALPHA = 0.10  # 结构系数
NUM_STAGES = 3
DELTA_V_BASE = 13300  # m/s


# ============== 发射场定义 ==============
@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个发射场 (按纬度排序)
ALL_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)

# 低纬度发射场 (纬度 < 30°)
LOW_LAT_SITES = [s for s in ALL_LAUNCH_SITES if s.abs_latitude < 30]


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

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/吨)"""
    k = fuel_ratio_multistage(site.total_delta_v)
    fuel_per_ton = k * 1000  # kg燃料/吨载荷
    return fuel_per_ton * SPECIFIC_FUEL_ENERGY


def average_rocket_energy(sites: List[LaunchSite]) -> float:
    """计算发射场列表的平均比能量 (J/吨)"""
    if not sites:
        return 0
    return sum(rocket_energy_per_ton(s) for s in sites) / len(sites)


def rocket_daily_capacity(sites: List[LaunchSite]) -> float:
    """计算发射场列表的日运力 (吨/天)"""
    return len(sites) * PAYLOAD_PER_LAUNCH  # 每站每天1发


def rocket_annual_capacity(sites: List[LaunchSite]) -> float:
    """计算发射场列表的年运力 (吨/年)"""
    return rocket_daily_capacity(sites) * 365


# ============== 水需求计算 ==============

def calculate_water_demand(params: WaterParameters) -> Dict:
    """
    计算水需求
    
    返回:
        包含各项水需求的字典
    """
    # 每日循环用水量 (参与水循环的水)
    daily_circulation_liters = POPULATION * params.daily_water_per_person
    daily_circulation_tons = daily_circulation_liters / 1000
    
    # 每日医疗用水 (不参与循环，直接消耗)
    daily_medical_liters = params.daily_medical_water
    daily_medical_tons = daily_medical_liters / 1000
    
    # 每日循环损耗 (需补充)
    daily_circulation_loss = daily_circulation_tons * (1 - params.recycle_rate)
    
    # 每日总补充量
    daily_supply_tons = daily_circulation_loss + daily_medical_tons
    
    # 每月补充量
    monthly_supply_tons = daily_supply_tons * 30
    
    # 年补充量
    annual_supply_tons = daily_supply_tons * 365
    
    # 水库初始存量 (支撑单日循环)
    reservoir_tons = daily_circulation_tons
    
    # 安全缓冲 (N天的补充量)
    safety_buffer_tons = daily_supply_tons * params.safety_buffer_days
    
    # 首批运输量 = 水库存量 + 安全缓冲
    initial_transport_tons = reservoir_tons + safety_buffer_tons
    
    return {
        'population': POPULATION,
        'comfort_factor': params.comfort_factor,
        'recycle_rate': params.recycle_rate,
        'safety_buffer_days': params.safety_buffer_days,
        
        # 每人每日
        'daily_per_person_liters': params.daily_water_per_person,
        'survival_water': params.survival_water,
        'hygiene_water': params.hygiene_water,
        
        # 医疗
        'daily_medical_events': params.daily_medical_events,
        'daily_medical_tons': daily_medical_tons,
        
        # 每日
        'daily_circulation_tons': daily_circulation_tons,
        'daily_circulation_loss_tons': daily_circulation_loss,
        'daily_supply_tons': daily_supply_tons,
        
        # 每月
        'monthly_supply_tons': monthly_supply_tons,
        
        # 年度
        'annual_supply_tons': annual_supply_tons,
        
        # 初始
        'reservoir_tons': reservoir_tons,
        'safety_buffer_tons': safety_buffer_tons,
        'initial_transport_tons': initial_transport_tons,
    }


# ============== 运输方案分析 ==============

@dataclass
class TransportScenario:
    """运输方案"""
    name: str
    name_cn: str
    use_elevator: bool
    rocket_sites: List[LaunchSite]
    description: str


def analyze_transport_scenario(
    scenario: TransportScenario,
    water_demand: Dict
) -> Dict:
    """
    分析单个运输方案
    
    返回:
        包含能量消耗和时间需求的字典
    """
    initial_tons = water_demand['initial_transport_tons']
    monthly_tons = water_demand['monthly_supply_tons']
    
    # 计算运力和比能量
    if scenario.use_elevator and scenario.rocket_sites:
        # 混合方案：电梯优先
        elevator_daily = TOTAL_ELEVATOR_CAPACITY / 365
        rocket_daily = rocket_daily_capacity(scenario.rocket_sites)
        total_daily = elevator_daily + rocket_daily
        
        # 比能量：按运力加权平均
        elevator_ratio = elevator_daily / total_daily
        rocket_ratio = rocket_daily / total_daily
        avg_energy = (elevator_ratio * ELEVATOR_SPECIFIC_ENERGY + 
                      rocket_ratio * average_rocket_energy(scenario.rocket_sites))
        
        annual_capacity = TOTAL_ELEVATOR_CAPACITY + rocket_annual_capacity(scenario.rocket_sites)
        
    elif scenario.use_elevator:
        # 纯电梯
        total_daily = TOTAL_ELEVATOR_CAPACITY / 365
        avg_energy = ELEVATOR_SPECIFIC_ENERGY
        annual_capacity = TOTAL_ELEVATOR_CAPACITY
        
    else:
        # 纯火箭
        total_daily = rocket_daily_capacity(scenario.rocket_sites)
        avg_energy = average_rocket_energy(scenario.rocket_sites)
        annual_capacity = rocket_annual_capacity(scenario.rocket_sites)
    
    # 首批运输
    initial_days = initial_tons / total_daily
    initial_energy_j = initial_tons * avg_energy
    initial_energy_pj = initial_energy_j / 1e15
    initial_energy_gj = initial_energy_j / 1e9
    
    # 每月补充
    monthly_days = monthly_tons / total_daily
    monthly_energy_j = monthly_tons * avg_energy
    monthly_energy_pj = monthly_energy_j / 1e15
    monthly_energy_gj = monthly_energy_j / 1e9
    
    # 年度补充
    annual_tons = water_demand['annual_supply_tons']
    annual_days = annual_tons / total_daily
    annual_energy_j = annual_tons * avg_energy
    annual_energy_pj = annual_energy_j / 1e15
    
    return {
        'scenario_name': scenario.name,
        'scenario_name_cn': scenario.name_cn,
        'description': scenario.description,
        
        # 运力参数
        'daily_capacity_tons': total_daily,
        'annual_capacity_tons': annual_capacity,
        'specific_energy_gj_per_ton': avg_energy / 1e9,
        
        # 首批运输 (水库 + 安全缓冲)
        'initial_tons': initial_tons,
        'initial_days': initial_days,
        'initial_energy_gj': initial_energy_gj,
        'initial_energy_pj': initial_energy_pj,
        
        # 每月补充
        'monthly_tons': monthly_tons,
        'monthly_days': monthly_days,
        'monthly_energy_gj': monthly_energy_gj,
        'monthly_energy_pj': monthly_energy_pj,
        
        # 年度补充
        'annual_tons': annual_tons,
        'annual_days': annual_days,
        'annual_energy_pj': annual_energy_pj,
    }


def define_scenarios() -> List[TransportScenario]:
    """定义四种运输方案"""
    return [
        TransportScenario(
            name="Elevator Only",
            name_cn="方案1: 仅空间电梯",
            use_elevator=True,
            rocket_sites=[],
            description="仅使用3个太空电梯运输水资源"
        ),
        TransportScenario(
            name="Rocket Only (All Sites)",
            name_cn="方案2: 仅火箭（10发射场）",
            use_elevator=False,
            rocket_sites=ALL_LAUNCH_SITES,
            description="使用全部10个发射场进行火箭运输"
        ),
        TransportScenario(
            name="Combined (All Sites)",
            name_cn="方案3: 混合（电梯+10发射场）",
            use_elevator=True,
            rocket_sites=ALL_LAUNCH_SITES,
            description="电梯与全部10个发射场火箭混合运输"
        ),
        TransportScenario(
            name="Combined (Low-Lat Sites)",
            name_cn="方案4: 混合（电梯+低纬发射场）",
            use_elevator=True,
            rocket_sites=LOW_LAT_SITES,
            description="电梯与低纬度(<30°)发射场混合运输"
        ),
    ]


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

def plot_scenario_comparison(
    results: List[Dict],
    water_demand: Dict,
    save_dir: str = '/Volumes/Files/code/mm/20260130_b/p3'
):
    """绘制方案对比图"""
    
    fig, axes = plt.subplots(2, 2, figsize=(16, 14))
    
    scenarios = [r['scenario_name_cn'] for r in results]
    colors = ['green', 'red', 'blue', 'purple']
    
    # ========== 图1: 首批运输能量消耗 ==========
    ax1 = axes[0, 0]
    initial_energy = [r['initial_energy_gj'] for r in results]
    bars1 = ax1.bar(range(len(scenarios)), initial_energy, color=colors, alpha=0.7)
    ax1.set_ylabel('Energy (GJ)', fontsize=12)
    ax1.set_title('Initial Transport Energy\n(Reservoir + 30-day Safety Buffer)', fontsize=13)
    ax1.set_xticks(range(len(scenarios)))
    ax1.set_xticklabels([s.replace('方案', 'S').replace(':', '\n').replace('（', '\n(').replace('）', ')') 
                         for s in scenarios], fontsize=9)
    ax1.grid(True, alpha=0.3, axis='y')
    
    for bar, val in zip(bars1, initial_energy):
        ax1.text(bar.get_x() + bar.get_width()/2, bar.get_height() + max(initial_energy)*0.02,
                f'{val:.1f}', ha='center', fontsize=10, fontweight='bold')
    
    # ========== 图2: 首批运输天数 ==========
    ax2 = axes[0, 1]
    initial_days = [r['initial_days'] for r in results]
    bars2 = ax2.bar(range(len(scenarios)), initial_days, color=colors, alpha=0.7)
    ax2.set_ylabel('Days', fontsize=12)
    ax2.set_title('Initial Transport Time\n(Reservoir + 30-day Safety Buffer)', fontsize=13)
    ax2.set_xticks(range(len(scenarios)))
    ax2.set_xticklabels([s.replace('方案', 'S').replace(':', '\n').replace('（', '\n(').replace('）', ')') 
                         for s in scenarios], fontsize=9)
    ax2.grid(True, alpha=0.3, axis='y')
    
    for bar, val in zip(bars2, initial_days):
        ax2.text(bar.get_x() + bar.get_width()/2, bar.get_height() + max(initial_days)*0.02,
                f'{val:.2f}', ha='center', fontsize=10, fontweight='bold')
    
    # ========== 图3: 每月补充能量消耗 ==========
    ax3 = axes[1, 0]
    monthly_energy = [r['monthly_energy_gj'] for r in results]
    bars3 = ax3.bar(range(len(scenarios)), monthly_energy, color=colors, alpha=0.7)
    ax3.set_ylabel('Energy (GJ)', fontsize=12)
    ax3.set_title('Monthly Supply Energy', fontsize=13)
    ax3.set_xticks(range(len(scenarios)))
    ax3.set_xticklabels([s.replace('方案', 'S').replace(':', '\n').replace('（', '\n(').replace('）', ')') 
                         for s in scenarios], fontsize=9)
    ax3.grid(True, alpha=0.3, axis='y')
    
    for bar, val in zip(bars3, monthly_energy):
        ax3.text(bar.get_x() + bar.get_width()/2, bar.get_height() + max(monthly_energy)*0.02,
                f'{val:.1f}', ha='center', fontsize=10, fontweight='bold')
    
    # ========== 图4: 每月补充天数 ==========
    ax4 = axes[1, 1]
    monthly_days = [r['monthly_days'] for r in results]
    bars4 = ax4.bar(range(len(scenarios)), monthly_days, color=colors, alpha=0.7)
    ax4.set_ylabel('Days', fontsize=12)
    ax4.set_title('Monthly Supply Time', fontsize=13)
    ax4.set_xticks(range(len(scenarios)))
    ax4.set_xticklabels([s.replace('方案', 'S').replace(':', '\n').replace('（', '\n(').replace('）', ')') 
                         for s in scenarios], fontsize=9)
    ax4.grid(True, alpha=0.3, axis='y')
    
    for bar, val in zip(bars4, monthly_days):
        ax4.text(bar.get_x() + bar.get_width()/2, bar.get_height() + max(monthly_days)*0.02,
                f'{val:.3f}', ha='center', fontsize=10, fontweight='bold')
    
    plt.suptitle(f'Water Transport Analysis for {POPULATION:,} Person Moon Colony\n'
                 f'(Recycle Rate: {water_demand["recycle_rate"]*100:.0f}%, '
                 f'Comfort Factor: {water_demand["comfort_factor"]:.1f})',
                 fontsize=15, y=1.02)
    
    plt.tight_layout()
    plt.savefig(f'{save_dir}/water_transport_comparison.png', dpi=150, bbox_inches='tight')
    print(f"方案对比图已保存至: {save_dir}/water_transport_comparison.png")
    
    return fig


def plot_water_demand_breakdown(
    water_demand: Dict,
    save_dir: str = '/Volumes/Files/code/mm/20260130_b/p3'
):
    """绘制水需求分解图"""
    
    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
    
    # ========== 图1: 每日用水构成 ==========
    ax1 = axes[0]
    
    daily_survival = POPULATION * water_demand['survival_water'] / 1000  # 吨
    daily_hygiene = POPULATION * water_demand['hygiene_water'] / 1000  # 吨
    daily_medical = water_demand['daily_medical_tons']
    
    categories = ['Survival\n(2.5 L/p/d)', 'Hygiene\n(0.4×α L/p/d)', 'Medical\n(Emergency)']
    values = [daily_survival, daily_hygiene, daily_medical]
    colors = ['#2ecc71', '#3498db', '#e74c3c']
    
    bars = ax1.bar(categories, values, color=colors, alpha=0.8)
    ax1.set_ylabel('Daily Water (tons)', fontsize=12)
    ax1.set_title('Daily Water Usage Breakdown', fontsize=13)
    ax1.grid(True, alpha=0.3, axis='y')
    
    for bar, val in zip(bars, values):
        ax1.text(bar.get_x() + bar.get_width()/2, bar.get_height() + max(values)*0.02,
                f'{val:.1f} t', ha='center', fontsize=11, fontweight='bold')
    
    # ========== 图2: 水流分析（桑基图简化版） ==========
    ax2 = axes[1]
    
    circulation = water_demand['daily_circulation_tons']
    recycled = circulation * water_demand['recycle_rate']
    loss = water_demand['daily_circulation_loss_tons']
    medical = water_demand['daily_medical_tons']
    supply = water_demand['daily_supply_tons']
    
    # 使用堆叠条形图展示水流
    bar_width = 0.6
    
    # 循环水
    ax2.barh(0, circulation, height=bar_width, color='#3498db', alpha=0.8, label='Circulation Water')
    ax2.barh(1, recycled, height=bar_width, color='#2ecc71', alpha=0.8, label='Recycled (90%)')
    ax2.barh(1, loss, height=bar_width, left=recycled, color='#e74c3c', alpha=0.8, label='Loss (10%)')
    ax2.barh(2, supply, height=bar_width, color='#9b59b6', alpha=0.8, label='Daily Supply Needed')
    
    ax2.set_yticks([0, 1, 2])
    ax2.set_yticklabels(['Daily\nCirculation', 'After\nRecycle', 'Daily\nSupply'])
    ax2.set_xlabel('Water (tons)', fontsize=12)
    ax2.set_title('Daily Water Flow Analysis', fontsize=13)
    ax2.legend(loc='lower right')
    ax2.grid(True, alpha=0.3, axis='x')
    
    # 添加数值标注
    ax2.text(circulation + 5, 0, f'{circulation:.1f} t', va='center', fontsize=10)
    ax2.text(recycled/2, 1, f'{recycled:.1f} t', va='center', ha='center', fontsize=9, color='white')
    ax2.text(recycled + loss/2, 1, f'{loss:.1f} t', va='center', ha='center', fontsize=9, color='white')
    ax2.text(supply + 5, 2, f'{supply:.1f} t\n(Loss + Medical)', va='center', fontsize=10)
    
    plt.suptitle(f'Water Demand Model for {POPULATION:,} Person Colony', fontsize=14, y=1.02)
    
    plt.tight_layout()
    plt.savefig(f'{save_dir}/water_demand_breakdown.png', dpi=150, bbox_inches='tight')
    print(f"水需求分解图已保存至: {save_dir}/water_demand_breakdown.png")
    
    return fig


def plot_annual_comparison(
    results: List[Dict],
    save_dir: str = '/Volumes/Files/code/mm/20260130_b/p3'
):
    """绘制年度能耗对比图"""
    
    fig, ax = plt.subplots(figsize=(12, 7))
    
    scenarios = [r['scenario_name_cn'].replace('方案', 'Scenario ').replace(':', ':\n') for r in results]
    
    # 数据
    initial_energy = [r['initial_energy_pj'] for r in results]
    annual_energy = [r['annual_energy_pj'] for r in results]
    total_first_year = [i + a for i, a in zip(initial_energy, annual_energy)]
    
    x = np.arange(len(scenarios))
    width = 0.35
    
    bars1 = ax.bar(x - width/2, initial_energy, width, label='Initial (Reservoir + Buffer)', 
                   color='#3498db', alpha=0.8)
    bars2 = ax.bar(x + width/2, annual_energy, width, label='Annual Supply', 
                   color='#e74c3c', alpha=0.8)
    
    ax.set_ylabel('Energy (PJ)', fontsize=12)
    ax.set_title('First Year Water Transport Energy Consumption', fontsize=14)
    ax.set_xticks(x)
    ax.set_xticklabels(scenarios, fontsize=9)
    ax.legend()
    ax.grid(True, alpha=0.3, axis='y')
    
    # 添加总计标签
    for i, (b1, b2, total) in enumerate(zip(bars1, bars2, total_first_year)):
        ax.text(i, max(initial_energy + annual_energy) * 1.05,
               f'Total: {total:.4f} PJ', ha='center', fontsize=10, fontweight='bold')
    
    plt.tight_layout()
    plt.savefig(f'{save_dir}/annual_energy_comparison.png', dpi=150, bbox_inches='tight')
    print(f"年度能耗对比图已保存至: {save_dir}/annual_energy_comparison.png")
    
    return fig


# ============== 报告生成 ==============

def generate_report(
    water_demand: Dict,
    results: List[Dict],
    params: WaterParameters
) -> str:
    """生成分析报告"""
    
    report = []
    report.append("=" * 80)
    report.append("TASK 3: WATER SUPPLY ANALYSIS FOR MOON COLONY")
    report.append("=" * 80)
    
    report.append("\n## 一、水需求模型参数")
    report.append(f"- 殖民地人口: {POPULATION:,} 人")
    report.append(f"- 舒适度系数 (α): {params.comfort_factor}")
    report.append(f"- 水循环回收率: {params.recycle_rate * 100:.0f}%")
    report.append(f"- 安全缓冲天数: {params.safety_buffer_days} 天")
    report.append(f"- 每百人每日患病率: {params.sickness_rate_per_100 * 100:.1f}%")
    
    report.append("\n## 二、每日用水标准")
    report.append(f"- 生存用水: {params.survival_water} L/人/天")
    report.append(f"- 卫生用水: {params.hygiene_water_base} × {params.comfort_factor} = {params.hygiene_water:.2f} L/人/天")
    report.append(f"- 医疗应急: {params.medical_water_per_event} L/次")
    report.append(f"- 每人每日循环用水: {params.daily_water_per_person:.2f} L")
    
    report.append("\n## 三、水需求计算结果")
    report.append(f"\n### 每日需求")
    report.append(f"- 每日循环用水量: {water_demand['daily_circulation_tons']:.1f} 吨")
    report.append(f"- 每日医疗事件数: {water_demand['daily_medical_events']:.0f} 次")
    report.append(f"- 每日医疗用水: {water_demand['daily_medical_tons']:.1f} 吨")
    report.append(f"- 每日循环损耗: {water_demand['daily_circulation_loss_tons']:.1f} 吨 (10%)")
    report.append(f"- 每日总补充量: {water_demand['daily_supply_tons']:.1f} 吨")
    
    report.append(f"\n### 每月需求")
    report.append(f"- 每月补充量: {water_demand['monthly_supply_tons']:.1f} 吨")
    
    report.append(f"\n### 年度需求")
    report.append(f"- 年补充量: {water_demand['annual_supply_tons']:.1f} 吨")
    
    report.append(f"\n### 初始存量")
    report.append(f"- 水库存量: {water_demand['reservoir_tons']:.1f} 吨")
    report.append(f"- 安全缓冲 ({params.safety_buffer_days}天): {water_demand['safety_buffer_tons']:.1f} 吨")
    report.append(f"- 首批运输总量: {water_demand['initial_transport_tons']:.1f} 吨")
    
    report.append("\n" + "=" * 80)
    report.append("## 四、四种运输方案对比")
    report.append("=" * 80)
    
    for r in results:
        report.append(f"\n### {r['scenario_name_cn']}")
        report.append(f"    {r['description']}")
        report.append(f"\n    运力参数:")
        report.append(f"    - 日运力: {r['daily_capacity_tons']:.1f} 吨/天")
        report.append(f"    - 年运力: {r['annual_capacity_tons']:,.0f} 吨/年")
        report.append(f"    - 比能量: {r['specific_energy_gj_per_ton']:.1f} GJ/吨")
        
        report.append(f"\n    首批运输 (水库 + {params.safety_buffer_days}天缓冲):")
        report.append(f"    - 运输量: {r['initial_tons']:.1f} 吨")
        report.append(f"    - 所需天数: {r['initial_days']:.2f} 天")
        report.append(f"    - 能量消耗: {r['initial_energy_gj']:.1f} GJ ({r['initial_energy_pj']:.6f} PJ)")
        
        report.append(f"\n    每月补充:")
        report.append(f"    - 运输量: {r['monthly_tons']:.1f} 吨")
        report.append(f"    - 所需天数: {r['monthly_days']:.3f} 天")
        report.append(f"    - 能量消耗: {r['monthly_energy_gj']:.1f} GJ ({r['monthly_energy_pj']:.6f} PJ)")
        
        report.append(f"\n    年度总计:")
        report.append(f"    - 年补充量: {r['annual_tons']:.1f} 吨")
        report.append(f"    - 年运输天数: {r['annual_days']:.2f} 天")
        report.append(f"    - 年能量消耗: {r['annual_energy_pj']:.4f} PJ")
    
    report.append("\n" + "=" * 80)
    report.append("## 五、方案对比汇总表")
    report.append("=" * 80)
    
    report.append(f"\n{'方案':<35} {'首批天数':>10} {'首批能耗(GJ)':>15} {'月补天数':>10} {'月补能耗(GJ)':>15}")
    report.append("-" * 90)
    for r in results:
        name = r['scenario_name_cn'][:30]
        report.append(f"{name:<35} {r['initial_days']:>10.2f} {r['initial_energy_gj']:>15.1f} {r['monthly_days']:>10.3f} {r['monthly_energy_gj']:>15.1f}")
    
    report.append("\n" + "=" * 80)
    report.append("## 六、结论与建议")
    report.append("=" * 80)
    
    # 找出最优方案
    min_energy_scenario = min(results, key=lambda x: x['initial_energy_gj'])
    min_time_scenario = min(results, key=lambda x: x['initial_days'])
    
    report.append(f"\n- 能耗最优方案: {min_energy_scenario['scenario_name_cn']}")
    report.append(f"  (首批能耗 {min_energy_scenario['initial_energy_gj']:.1f} GJ)")
    report.append(f"\n- 时间最优方案: {min_time_scenario['scenario_name_cn']}")
    report.append(f"  (首批用时 {min_time_scenario['initial_days']:.2f} 天)")
    
    report.append(f"\n- 水运输仅占电梯年运力的 {water_demand['annual_supply_tons']/TOTAL_ELEVATOR_CAPACITY*100:.2f}%")
    report.append(f"  因此水供应不会对建设任务造成显著影响。")
    
    return "\n".join(report)


# ============== 多场景分析 ==============

def analyze_multiple_scenarios(comfort_factors: List[float] = [1, 50, 250]):
    """
    分析多个舒适度系数场景
    
    参数:
        comfort_factors: 舒适度系数列表
    """
    all_results = {}
    all_demands = {}
    
    for alpha in comfort_factors:
        params = WaterParameters(
            survival_water=2.5,
            hygiene_water_base=0.4,
            comfort_factor=alpha,
            medical_water_per_event=5.0,
            sickness_rate_per_100=0.02,
            recycle_rate=0.90,
            safety_buffer_days=30
        )
        
        water_demand = calculate_water_demand(params)
        scenarios = define_scenarios()
        
        results = []
        for scenario in scenarios:
            result = analyze_transport_scenario(scenario, water_demand)
            result['comfort_factor'] = alpha
            results.append(result)
        
        all_results[alpha] = results
        all_demands[alpha] = water_demand
    
    return all_results, all_demands


def plot_multi_scenario_comparison(
    all_results: Dict,
    all_demands: Dict,
    save_dir: str = '/Volumes/Files/code/mm/20260130_b/p3'
):
    """绘制多场景对比图"""
    
    comfort_factors = list(all_results.keys())
    scenarios = define_scenarios()
    scenario_names = [s.name_cn.replace('方案', 'S').replace(':', '').replace('（', '\n(').replace('）', ')') 
                      for s in scenarios]
    
    colors = ['green', 'red', 'blue', 'purple']
    
    # ========== 图1: 首批运输能耗对比 ==========
    fig1, ax1 = plt.subplots(figsize=(14, 8))
    
    x = np.arange(len(scenarios))
    width = 0.25
    
    for i, alpha in enumerate(comfort_factors):
        results = all_results[alpha]
        energy = [r['initial_energy_gj'] / 1000 for r in results]  # 转换为TJ
        bars = ax1.bar(x + i * width, energy, width, label=f'α = {alpha}', alpha=0.8)
        
        for bar, val in zip(bars, energy):
            ax1.text(bar.get_x() + bar.get_width()/2, bar.get_height() + max(energy)*0.02,
                    f'{val:.1f}', ha='center', fontsize=8, rotation=0)
    
    ax1.set_ylabel('Energy (TJ)', fontsize=12)
    ax1.set_title('Initial Transport Energy by Comfort Factor\n(Reservoir + 30-day Buffer)', fontsize=14)
    ax1.set_xticks(x + width)
    ax1.set_xticklabels(scenario_names, fontsize=9)
    ax1.legend(title='Comfort Factor')
    ax1.grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    plt.savefig(f'{save_dir}/multi_scenario_initial_energy.png', dpi=150, bbox_inches='tight')
    print(f"多场景首批能耗图已保存")
    
    # ========== 图2: 首批运输天数对比 ==========
    fig2, ax2 = plt.subplots(figsize=(14, 8))
    
    for i, alpha in enumerate(comfort_factors):
        results = all_results[alpha]
        days = [r['initial_days'] for r in results]
        bars = ax2.bar(x + i * width, days, width, label=f'α = {alpha}', alpha=0.8)
        
        for bar, val in zip(bars, days):
            ax2.text(bar.get_x() + bar.get_width()/2, bar.get_height() + max(days)*0.02,
                    f'{val:.1f}', ha='center', fontsize=8)
    
    ax2.set_ylabel('Days', fontsize=12)
    ax2.set_title('Initial Transport Time by Comfort Factor\n(Reservoir + 30-day Buffer)', fontsize=14)
    ax2.set_xticks(x + width)
    ax2.set_xticklabels(scenario_names, fontsize=9)
    ax2.legend(title='Comfort Factor')
    ax2.grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    plt.savefig(f'{save_dir}/multi_scenario_initial_days.png', dpi=150, bbox_inches='tight')
    print(f"多场景首批天数图已保存")
    
    # ========== 图3: 每月补充能耗对比 ==========
    fig3, ax3 = plt.subplots(figsize=(14, 8))
    
    for i, alpha in enumerate(comfort_factors):
        results = all_results[alpha]
        energy = [r['monthly_energy_gj'] / 1000 for r in results]  # 转换为TJ
        bars = ax3.bar(x + i * width, energy, width, label=f'α = {alpha}', alpha=0.8)
        
        for bar, val in zip(bars, energy):
            ax3.text(bar.get_x() + bar.get_width()/2, bar.get_height() + max(energy)*0.02,
                    f'{val:.1f}', ha='center', fontsize=8)
    
    ax3.set_ylabel('Energy (TJ)', fontsize=12)
    ax3.set_title('Monthly Supply Energy by Comfort Factor', fontsize=14)
    ax3.set_xticks(x + width)
    ax3.set_xticklabels(scenario_names, fontsize=9)
    ax3.legend(title='Comfort Factor')
    ax3.grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    plt.savefig(f'{save_dir}/multi_scenario_monthly_energy.png', dpi=150, bbox_inches='tight')
    print(f"多场景月补充能耗图已保存")
    
    # ========== 图4: 每月补充天数对比 ==========
    fig4, ax4 = plt.subplots(figsize=(14, 8))
    
    for i, alpha in enumerate(comfort_factors):
        results = all_results[alpha]
        days = [r['monthly_days'] for r in results]
        bars = ax4.bar(x + i * width, days, width, label=f'α = {alpha}', alpha=0.8)
        
        for bar, val in zip(bars, days):
            ax4.text(bar.get_x() + bar.get_width()/2, bar.get_height() + max(days)*0.02,
                    f'{val:.2f}', ha='center', fontsize=8)
    
    ax4.set_ylabel('Days', fontsize=12)
    ax4.set_title('Monthly Supply Time by Comfort Factor', fontsize=14)
    ax4.set_xticks(x + width)
    ax4.set_xticklabels(scenario_names, fontsize=9)
    ax4.legend(title='Comfort Factor')
    ax4.grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    plt.savefig(f'{save_dir}/multi_scenario_monthly_days.png', dpi=150, bbox_inches='tight')
    print(f"多场景月补充天数图已保存")
    
    # ========== 图5: 综合对比热力图 ==========
    fig5, axes5 = plt.subplots(2, 2, figsize=(16, 14))
    
    # 准备数据矩阵
    scenario_labels = ['Elevator\nOnly', 'Rocket\nOnly', 'Combined\n(All)', 'Combined\n(Low-Lat)']
    alpha_labels = [f'α={a}' for a in comfort_factors]
    
    # 首批能耗
    ax = axes5[0, 0]
    data = np.array([[all_results[a][i]['initial_energy_gj']/1000 for i in range(4)] for a in comfort_factors])
    im = ax.imshow(data, cmap='YlOrRd', aspect='auto')
    ax.set_xticks(range(4))
    ax.set_xticklabels(scenario_labels, fontsize=9)
    ax.set_yticks(range(len(comfort_factors)))
    ax.set_yticklabels(alpha_labels)
    ax.set_title('Initial Energy (TJ)', fontsize=12)
    for i in range(len(comfort_factors)):
        for j in range(4):
            ax.text(j, i, f'{data[i,j]:.1f}', ha='center', va='center', fontsize=10, 
                   color='white' if data[i,j] > data.max()*0.5 else 'black')
    plt.colorbar(im, ax=ax)
    
    # 首批天数
    ax = axes5[0, 1]
    data = np.array([[all_results[a][i]['initial_days'] for i in range(4)] for a in comfort_factors])
    im = ax.imshow(data, cmap='YlOrRd', aspect='auto')
    ax.set_xticks(range(4))
    ax.set_xticklabels(scenario_labels, fontsize=9)
    ax.set_yticks(range(len(comfort_factors)))
    ax.set_yticklabels(alpha_labels)
    ax.set_title('Initial Days', fontsize=12)
    for i in range(len(comfort_factors)):
        for j in range(4):
            ax.text(j, i, f'{data[i,j]:.1f}', ha='center', va='center', fontsize=10,
                   color='white' if data[i,j] > data.max()*0.5 else 'black')
    plt.colorbar(im, ax=ax)
    
    # 月补充能耗
    ax = axes5[1, 0]
    data = np.array([[all_results[a][i]['monthly_energy_gj']/1000 for i in range(4)] for a in comfort_factors])
    im = ax.imshow(data, cmap='YlOrRd', aspect='auto')
    ax.set_xticks(range(4))
    ax.set_xticklabels(scenario_labels, fontsize=9)
    ax.set_yticks(range(len(comfort_factors)))
    ax.set_yticklabels(alpha_labels)
    ax.set_title('Monthly Energy (TJ)', fontsize=12)
    for i in range(len(comfort_factors)):
        for j in range(4):
            ax.text(j, i, f'{data[i,j]:.1f}', ha='center', va='center', fontsize=10,
                   color='white' if data[i,j] > data.max()*0.5 else 'black')
    plt.colorbar(im, ax=ax)
    
    # 月补充天数
    ax = axes5[1, 1]
    data = np.array([[all_results[a][i]['monthly_days'] for i in range(4)] for a in comfort_factors])
    im = ax.imshow(data, cmap='YlOrRd', aspect='auto')
    ax.set_xticks(range(4))
    ax.set_xticklabels(scenario_labels, fontsize=9)
    ax.set_yticks(range(len(comfort_factors)))
    ax.set_yticklabels(alpha_labels)
    ax.set_title('Monthly Days', fontsize=12)
    for i in range(len(comfort_factors)):
        for j in range(4):
            ax.text(j, i, f'{data[i,j]:.2f}', ha='center', va='center', fontsize=10,
                   color='white' if data[i,j] > data.max()*0.5 else 'black')
    plt.colorbar(im, ax=ax)
    
    plt.suptitle('Water Transport Analysis: Multi-Scenario Comparison\n(100,000 Person Moon Colony)', 
                 fontsize=15, y=1.02)
    plt.tight_layout()
    plt.savefig(f'{save_dir}/multi_scenario_heatmap.png', dpi=150, bbox_inches='tight')
    print(f"多场景热力图已保存")
    
    return fig1, fig2, fig3, fig4, fig5


def generate_multi_scenario_report(
    all_results: Dict,
    all_demands: Dict,
    comfort_factors: List[float]
) -> str:
    """生成多场景分析报告"""
    
    report = []
    report.append("=" * 100)
    report.append("TASK 3: WATER SUPPLY ANALYSIS - MULTI-SCENARIO COMPARISON")
    report.append("=" * 100)
    
    report.append("\n## 模型参数")
    report.append(f"- 殖民地人口: {POPULATION:,} 人")
    report.append(f"- 生存用水: 2.5 L/人/天")
    report.append(f"- 卫生用水: 0.4 × α L/人/天")
    report.append(f"- 医疗应急: 5.0 L/次 (每百人每日患病率 2%)")
    report.append(f"- 水循环回收率: 90%")
    report.append(f"- 安全缓冲: 30 天")
    report.append(f"- 舒适度系数场景: α = {comfort_factors}")
    
    report.append("\n" + "=" * 100)
    report.append("## 各场景水需求汇总")
    report.append("=" * 100)
    
    report.append(f"\n{'α':<10} {'人均日用水(L)':<15} {'水库(吨)':<12} {'日补充(吨)':<12} {'月补充(吨)':<12} {'首批(吨)':<12} {'年补充(吨)':<12}")
    report.append("-" * 100)
    for alpha in comfort_factors:
        d = all_demands[alpha]
        report.append(f"{alpha:<10} {d['daily_per_person_liters']:<15.1f} {d['reservoir_tons']:<12.1f} {d['daily_supply_tons']:<12.1f} {d['monthly_supply_tons']:<12.1f} {d['initial_transport_tons']:<12.1f} {d['annual_supply_tons']:<12.1f}")
    
    # 各场景详细结果
    for alpha in comfort_factors:
        report.append(f"\n\n{'=' * 100}")
        report.append(f"## 场景: 舒适度系数 α = {alpha}")
        report.append(f"{'=' * 100}")
        
        d = all_demands[alpha]
        report.append(f"\n### 水需求参数")
        report.append(f"- 每人每日用水: {d['daily_per_person_liters']:.1f} L (生存 {d['survival_water']} + 卫生 {d['hygiene_water']:.1f})")
        report.append(f"- 每日循环水量 (水库): {d['reservoir_tons']:.1f} 吨")
        report.append(f"- 每日医疗用水: {d['daily_medical_tons']:.1f} 吨 ({d['daily_medical_events']:.0f} 次)")
        report.append(f"- 每日循环损耗 (10%): {d['daily_circulation_loss_tons']:.1f} 吨")
        report.append(f"- 每日总补充量: {d['daily_supply_tons']:.1f} 吨")
        report.append(f"- 每月补充量: {d['monthly_supply_tons']:.1f} 吨")
        report.append(f"- 首批运输量 (水库+30天缓冲): {d['initial_transport_tons']:.1f} 吨")
        report.append(f"- 年补充量: {d['annual_supply_tons']:.1f} 吨")
        
        report.append(f"\n### 四种运输方案对比")
        report.append(f"\n{'方案':<40} {'首批天数':<12} {'首批能耗(TJ)':<15} {'月补天数':<12} {'月补能耗(TJ)':<15}")
        report.append("-" * 100)
        
        for r in all_results[alpha]:
            name = r['scenario_name_cn'][:35]
            report.append(f"{name:<40} {r['initial_days']:<12.2f} {r['initial_energy_gj']/1000:<15.2f} {r['monthly_days']:<12.3f} {r['monthly_energy_gj']/1000:<15.2f}")
    
    # 横向对比汇总
    report.append(f"\n\n{'=' * 100}")
    report.append("## 横向对比汇总表")
    report.append("=" * 100)
    
    scenarios = define_scenarios()
    
    for i, scenario in enumerate(scenarios):
        report.append(f"\n### {scenario.name_cn}")
        report.append(f"\n{'α':<10} {'首批天数':<12} {'首批能耗(TJ)':<15} {'月补天数':<12} {'月补能耗(TJ)':<15} {'年能耗(PJ)':<12}")
        report.append("-" * 80)
        for alpha in comfort_factors:
            r = all_results[alpha][i]
            report.append(f"{alpha:<10} {r['initial_days']:<12.2f} {r['initial_energy_gj']/1000:<15.2f} {r['monthly_days']:<12.3f} {r['monthly_energy_gj']/1000:<15.2f} {r['annual_energy_pj']:<12.4f}")
    
    report.append(f"\n\n{'=' * 100}")
    report.append("## 结论与分析")
    report.append("=" * 100)
    
    report.append("\n### 1. 舒适度系数影响")
    d1 = all_demands[comfort_factors[0]]
    d3 = all_demands[comfort_factors[-1]]
    increase = (d3['annual_supply_tons'] / d1['annual_supply_tons'] - 1) * 100
    report.append(f"- α从{comfort_factors[0]}增加到{comfort_factors[-1]}时，年补充量从{d1['annual_supply_tons']:.0f}吨增加到{d3['annual_supply_tons']:.0f}吨 (+{increase:.0f}%)")
    
    report.append("\n### 2. 运输方案比较")
    report.append("- 能耗最优: 方案1 (仅电梯) - 比能量 157.2 GJ/吨")
    report.append("- 时间最优: 方案3 (混合全部) - 日运力 2,721 吨/天")
    report.append("- 平衡方案: 方案4 (混合低纬) - 兼顾能耗和时间")
    
    report.append("\n### 3. 占用运力比例")
    for alpha in comfort_factors:
        d = all_demands[alpha]
        ratio = d['annual_supply_tons'] / TOTAL_ELEVATOR_CAPACITY * 100
        report.append(f"- α={alpha}: 年补充量占电梯年运力的 {ratio:.2f}%")
    
    return "\n".join(report)


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

def run_water_analysis():
    """运行完整的水需求分析"""
    
    print("=" * 80)
    print("任务三：月球殖民地水需求与运输分析")
    print("=" * 80)
    
    # 三个舒适度场景
    comfort_factors = [1, 50, 250]
    
    print(f"\n分析三个舒适度场景: α = {comfort_factors}")
    
    # 多场景分析
    print("\n正在计算各场景水需求...")
    all_results, all_demands = analyze_multiple_scenarios(comfort_factors)
    
    # 打印各场景摘要
    for alpha in comfort_factors:
        d = all_demands[alpha]
        print(f"\n场景 α = {alpha}:")
        print(f"  - 每人每日用水: {d['daily_per_person_liters']:.1f} L")
        print(f"  - 水库存量: {d['reservoir_tons']:.1f} 吨")
        print(f"  - 每日补充量: {d['daily_supply_tons']:.1f} 吨")
        print(f"  - 每月补充量: {d['monthly_supply_tons']:.1f} 吨")
        print(f"  - 首批运输量: {d['initial_transport_tons']:.1f} 吨")
        print(f"  - 年补充量: {d['annual_supply_tons']:.1f} 吨")
    
    # 打印各方案结果
    print("\n" + "=" * 80)
    print("各场景各方案详细结果")
    print("=" * 80)
    
    for alpha in comfort_factors:
        print(f"\n--- α = {alpha} ---")
        for r in all_results[alpha]:
            print(f"  {r['scenario_name_cn']}:")
            print(f"    首批: {r['initial_days']:.2f} 天, {r['initial_energy_gj']/1000:.2f} TJ")
            print(f"    每月: {r['monthly_days']:.3f} 天, {r['monthly_energy_gj']/1000:.2f} TJ")
    
    # 生成可视化
    print("\n正在生成可视化图表...")
    
    # 单场景图（使用α=1）
    plot_water_demand_breakdown(all_demands[1])
    plot_scenario_comparison(all_results[1], all_demands[1])
    plot_annual_comparison(all_results[1])
    
    # 多场景对比图
    plot_multi_scenario_comparison(all_results, all_demands)
    
    # 生成报告
    print("\n正在生成分析报告...")
    report = generate_multi_scenario_report(all_results, all_demands, comfort_factors)
    print(report)
    
    # 保存报告
    with open('/Volumes/Files/code/mm/20260130_b/p3/water_analysis_report.txt', 'w', encoding='utf-8') as f:
        f.write(report)
    print(f"\n报告已保存至: p3/water_analysis_report.txt")
    
    # 保存数据
    all_data = []
    for alpha in comfort_factors:
        for r in all_results[alpha]:
            r_copy = r.copy()
            r_copy['comfort_factor'] = alpha
            all_data.append(r_copy)
    
    df = pd.DataFrame(all_data)
    df.to_csv('/Volumes/Files/code/mm/20260130_b/p3/water_transport_results.csv', index=False)
    print(f"数据已保存至: p3/water_transport_results.csv")
    
    return all_results, all_demands, comfort_factors


if __name__ == "__main__":
    all_results, all_demands, comfort_factors = run_water_analysis()
