"""
Lambda Time Cost Analysis for Moon Colony Logistics

This module introduces the time opportunity cost (λ) to find the optimal
operating point on the Energy-Time trade-off curve.

Objective Function: J = E_total + λ × T
where:
    - E_total: Total energy consumption (PJ)
    - T: Construction timeline (years)
    - λ: Time opportunity cost (PJ/year)

The optimal point satisfies: dE/dT = -λ (marginal energy saving = time cost)
"""

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

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

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

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

# Space Elevator parameters
NUM_ELEVATORS = 3
ELEVATOR_CAPACITY_PER_YEAR = 179000  # metric tons per elevator per year
TOTAL_ELEVATOR_CAPACITY = NUM_ELEVATORS * ELEVATOR_CAPACITY_PER_YEAR  # 537,000 tons/year
ELEVATOR_SPECIFIC_ENERGY = 157.2e9  # J per metric ton (157.2 MJ/kg × 1000)

# Rocket parameters - LOX/CH4 (Raptor-class)
PAYLOAD_PER_LAUNCH = 125  # metric tons per launch
ISP = 363  # Specific impulse (s)
SPECIFIC_FUEL_ENERGY = 11.9e6  # J/kg
ALPHA = 0.10  # Structural coefficient
NUM_STAGES = 3
DELTA_V_BASE = 13300  # m/s (LEO + TLI + LOI)


# ============== Launch Site Definition ==============
@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 delta_v_loss(self) -> float:
        v_equator = OMEGA_EARTH * R_EARTH
        v_site = OMEGA_EARTH * R_EARTH * np.cos(np.radians(self.abs_latitude))
        return v_equator - v_site
    
    @property
    def total_delta_v(self) -> float:
        return DELTA_V_BASE + self.delta_v_loss


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)


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

def fuel_ratio_multistage(delta_v: float) -> float:
    """Multi-stage rocket fuel/payload ratio"""
    ve = ISP * G0
    delta_v_per_stage = delta_v / NUM_STAGES
    R_stage = np.exp(delta_v_per_stage / ve)
    
    denominator = 1 - ALPHA * (R_stage - 1)
    if denominator <= 0:
        return np.inf
    
    k_stage = (R_stage - 1) / denominator
    
    total_fuel_ratio = 0
    remaining_ratio = 1.0
    
    for _ in range(NUM_STAGES):
        fuel_this_stage = remaining_ratio * k_stage
        total_fuel_ratio += fuel_this_stage
        remaining_ratio *= (1 + k_stage * (1 + ALPHA))
    
    return total_fuel_ratio


def rocket_energy_per_ton(site: LaunchSite) -> float:
    """Energy consumption per ton of payload for rocket launch (J/ton)"""
    k = fuel_ratio_multistage(site.total_delta_v)
    fuel_per_ton = k * 1000  # kg fuel per metric ton payload
    return fuel_per_ton * SPECIFIC_FUEL_ENERGY


def calculate_scenario(completion_years: float) -> Optional[Dict]:
    """
    Calculate optimal scenario for given completion timeline
    (elevator priority + low-latitude rockets)
    """
    # Space elevator transport
    elevator_payload = min(TOTAL_ELEVATOR_CAPACITY * completion_years, TOTAL_PAYLOAD)
    elevator_energy = elevator_payload * ELEVATOR_SPECIFIC_ENERGY
    
    # Remaining payload for rockets
    remaining_payload = TOTAL_PAYLOAD - elevator_payload
    
    if remaining_payload <= 0:
        return {
            'years': completion_years,
            'elevator_payload': elevator_payload,
            'rocket_payload': 0,
            'elevator_energy_PJ': elevator_energy / 1e15,
            'rocket_energy_PJ': 0,
            'total_energy_PJ': elevator_energy / 1e15,
            'rocket_launches': 0,
            'sites_used': 0,
            'elevator_fraction': 1.0,
        }
    
    # Rocket launches needed
    rocket_launches_needed = int(np.ceil(remaining_payload / PAYLOAD_PER_LAUNCH))
    
    # Allocate by latitude priority
    days_available = completion_years * 365
    max_launches_per_site = int(days_available)
    
    # Check feasibility
    total_rocket_capacity = len(LAUNCH_SITES) * max_launches_per_site * PAYLOAD_PER_LAUNCH
    if remaining_payload > total_rocket_capacity:
        return None
    
    rocket_energy = 0
    sites_used = 0
    remaining_launches = rocket_launches_needed
    
    for site in LAUNCH_SITES:
        if remaining_launches <= 0:
            break
        allocated = min(remaining_launches, max_launches_per_site)
        rocket_energy += rocket_energy_per_ton(site) * PAYLOAD_PER_LAUNCH * allocated
        remaining_launches -= allocated
        if allocated > 0:
            sites_used += 1
    
    rocket_payload = rocket_launches_needed * PAYLOAD_PER_LAUNCH
    total_energy = elevator_energy + rocket_energy
    
    return {
        'years': completion_years,
        'elevator_payload': elevator_payload,
        'rocket_payload': rocket_payload,
        'elevator_energy_PJ': elevator_energy / 1e15,
        'rocket_energy_PJ': rocket_energy / 1e15,
        'total_energy_PJ': total_energy / 1e15,
        'rocket_launches': rocket_launches_needed,
        'sites_used': sites_used,
        'elevator_fraction': elevator_payload / TOTAL_PAYLOAD,
    }


# ============== Generate Trade-off Curve ==============

def generate_tradeoff_curve(
    year_min: float = 100,
    year_max: float = 250,
    num_points: int = 500
) -> pd.DataFrame:
    """Generate Energy-Time trade-off curve data"""
    years_range = np.linspace(year_min, year_max, num_points)
    
    results = []
    for years in years_range:
        scenario = calculate_scenario(years)
        if scenario is not None:
            results.append({
                'years': years,
                'energy_PJ': scenario['total_energy_PJ'],
                'elevator_fraction': scenario['elevator_fraction'],
                'sites_used': scenario['sites_used'],
                'rocket_launches': scenario['rocket_launches'],
            })
    
    return pd.DataFrame(results)


# ============== Lambda Cost Analysis ==============

def calculate_total_cost(df: pd.DataFrame, lambda_cost: float) -> np.ndarray:
    """
    Calculate total cost J = E + λ × T
    
    Args:
        df: Trade-off curve data
        lambda_cost: Time opportunity cost (PJ/year)
    
    Returns:
        Total cost array
    """
    return df['energy_PJ'].values + lambda_cost * df['years'].values


def find_optimal_point(df: pd.DataFrame, lambda_cost: float) -> Dict:
    """
    Find optimal point that minimizes J = E + λ × T
    
    Args:
        df: Trade-off curve data
        lambda_cost: Time opportunity cost (PJ/year)
    
    Returns:
        Optimal point information
    """
    total_cost = calculate_total_cost(df, lambda_cost)
    opt_idx = np.argmin(total_cost)
    
    return {
        'index': opt_idx,
        'years': df['years'].iloc[opt_idx],
        'energy_PJ': df['energy_PJ'].iloc[opt_idx],
        'total_cost': total_cost[opt_idx],
        'elevator_fraction': df['elevator_fraction'].iloc[opt_idx],
        'lambda': lambda_cost,
    }


def calculate_marginal_energy_saving(df: pd.DataFrame) -> np.ndarray:
    """
    Calculate marginal energy saving rate: -dE/dT (PJ/year)
    
    This represents how much energy is saved per additional year of timeline.
    """
    years = df['years'].values
    energy = df['energy_PJ'].values
    
    # Use central difference for interior points
    marginal = -np.gradient(energy, years)
    
    return marginal


def sensitivity_analysis(
    df: pd.DataFrame,
    lambda_range: np.ndarray
) -> pd.DataFrame:
    """
    Perform sensitivity analysis on λ parameter
    
    Args:
        df: Trade-off curve data
        lambda_range: Array of λ values to test
    
    Returns:
        DataFrame with optimal points for each λ
    """
    results = []
    for lam in lambda_range:
        opt = find_optimal_point(df, lam)
        results.append({
            'lambda_PJ_per_year': lam,
            'optimal_years': opt['years'],
            'optimal_energy_PJ': opt['energy_PJ'],
            'total_cost_PJ': opt['total_cost'],
            'elevator_fraction': opt['elevator_fraction'],
        })
    
    return pd.DataFrame(results)


# ============== Visualization Functions ==============

def plot_lambda_analysis(
    df: pd.DataFrame,
    save_path: str = '/Volumes/Files/code/mm/20260130_b/p1/lambda_cost_analysis.png'
):
    """
    Comprehensive visualization of λ time cost analysis
    Focus on critical range λ = 400-600 PJ/year
    """
    fig, axes = plt.subplots(2, 2, figsize=(14, 12))
    
    years = df['years'].values
    energy = df['energy_PJ'].values
    
    # Key boundaries
    T_min = years.min()  # ~100.7 years (fastest)
    T_elev = TOTAL_PAYLOAD / TOTAL_ELEVATOR_CAPACITY  # ~186 years (elevator only)
    E_min = energy[years >= T_elev].min() if any(years >= T_elev) else energy.min()
    
    # ========== Plot 1: Trade-off Curve with λ Iso-cost Lines ==========
    ax1 = axes[0, 0]
    
    # Main curve
    ax1.plot(years, energy, 'b-', linewidth=2.5, label='Energy-Time Trade-off')
    
    # Mark key points
    ax1.axvline(x=T_elev, color='green', linestyle='--', alpha=0.7, label=f'Elevator-only: {T_elev:.1f} years')
    ax1.axvline(x=T_min, color='red', linestyle='--', alpha=0.7, label=f'Minimum time: {T_min:.1f} years')
    
    # Draw iso-cost lines for different λ (focus on 400-600 range)
    lambda_values = [420, 480, 500, 550]
    colors = ['#2ecc71', '#e74c3c', '#9b59b6', '#3498db']
    
    for lam, color in zip(lambda_values, colors):
        opt = find_optimal_point(df, lam)
        # Iso-cost line: E + λT = const → E = const - λT
        T_line = np.linspace(80, 220, 100)
        E_line = opt['total_cost'] - lam * T_line
        valid = (E_line > 0) & (E_line < 70000)
        ax1.plot(T_line[valid], E_line[valid], '--', color=color, alpha=0.6, linewidth=1.5)
        ax1.plot(opt['years'], opt['energy_PJ'], 'o', color=color, markersize=12,
                markeredgecolor='black', markeredgewidth=1.5,
                label=f'λ={lam}: T={opt["years"]:.1f}y, E={opt["energy_PJ"]:.0f}PJ')
    
    ax1.set_xlabel('Construction Timeline T (years)', fontsize=12)
    ax1.set_ylabel('Total Energy E (PJ)', fontsize=12)
    ax1.set_title('Energy-Time Trade-off with λ Iso-cost Lines (λ=400-600)\n$J = E + λT$', fontsize=13)
    ax1.legend(loc='upper right', fontsize=9)
    ax1.grid(True, alpha=0.3)
    ax1.set_xlim(95, 200)
    ax1.set_ylim(10000, 65000)
    
    # ========== Plot 2: Marginal Energy Saving Rate ==========
    ax2 = axes[0, 1]
    
    marginal = calculate_marginal_energy_saving(df)
    
    ax2.plot(years, marginal, 'r-', linewidth=2, label='Marginal Energy Saving -dE/dT')
    ax2.axhline(y=0, color='black', linestyle='-', alpha=0.3)
    
    # Mark critical λ values in 400-600 range
    for lam, color in zip([450, 480, 500, 550], ['#2ecc71', '#e74c3c', '#9b59b6', '#3498db']):
        ax2.axhline(y=lam, color=color, linestyle='--', alpha=0.7, label=f'λ = {lam} PJ/year')
        # Find intersection
        intersect_idx = np.argmin(np.abs(marginal - lam))
        ax2.plot(years[intersect_idx], marginal[intersect_idx], 'o', color=color, markersize=8)
    
    ax2.set_xlabel('Construction Timeline T (years)', fontsize=12)
    ax2.set_ylabel('Marginal Energy Saving -dE/dT (PJ/year)', fontsize=12)
    ax2.set_title('Marginal Analysis: Optimal Condition is -dE/dT = λ', fontsize=13)
    ax2.legend(loc='upper right', fontsize=9)
    ax2.grid(True, alpha=0.3)
    ax2.set_xlim(95, 200)
    ax2.set_ylim(350, 620)
    
    # ========== Plot 3: Sensitivity Analysis (400-600 range) ==========
    ax3 = axes[1, 0]
    
    lambda_range = np.linspace(400, 600, 200)
    sensitivity_df = sensitivity_analysis(df, lambda_range)
    
    # Plot optimal years vs λ
    ax3.plot(sensitivity_df['lambda_PJ_per_year'], sensitivity_df['optimal_years'], 
             'b-', linewidth=2.5, label='Optimal Timeline T*')
    
    # Mark key regions
    ax3.axhline(y=T_elev, color='green', linestyle='--', alpha=0.7, label=f'Elevator-only: {T_elev:.1f}y')
    ax3.axhline(y=T_min, color='red', linestyle='--', alpha=0.7, label=f'Min timeline: {T_min:.1f}y')
    ax3.axhline(y=139, color='orange', linestyle=':', linewidth=2, alpha=0.8, label='Original knee: 139y')
    
    # Find λ corresponding to 139 years
    idx_139 = np.argmin(np.abs(sensitivity_df['optimal_years'] - 139))
    if idx_139 < len(sensitivity_df):
        lambda_139 = sensitivity_df['lambda_PJ_per_year'].iloc[idx_139]
        ax3.axvline(x=lambda_139, color='orange', linestyle=':', linewidth=2, alpha=0.6)
        ax3.scatter([lambda_139], [139], s=150, c='orange', marker='*', zorder=5, 
                   edgecolors='black', linewidths=1.5, label=f'λ≈{lambda_139:.0f} for T*=139y')
    
    # Mark critical transition
    ax3.axvline(x=480, color='purple', linestyle='--', alpha=0.5)
    ax3.text(482, 175, 'Critical\nTransition', fontsize=9, color='purple')
    
    ax3.set_xlabel('Time Opportunity Cost λ (PJ/year)', fontsize=12)
    ax3.set_ylabel('Optimal Timeline T* (years)', fontsize=12)
    ax3.set_title('Sensitivity Analysis: Optimal Timeline vs λ (400-600 range)', fontsize=13)
    ax3.legend(loc='upper right', fontsize=9)
    ax3.grid(True, alpha=0.3)
    ax3.set_xlim(400, 600)
    ax3.set_ylim(95, 195)
    
    # ========== Plot 4: Total Cost Curves (400-600 range) ==========
    ax4 = axes[1, 1]
    
    for lam, color in [(450, '#2ecc71'), (480, '#e74c3c'), (500, '#9b59b6'), (550, '#3498db')]:
        total_cost = calculate_total_cost(df, lam)
        ax4.plot(years, total_cost / 1000, color=color, linestyle='-', 
                linewidth=2, label=f'λ={lam} PJ/year')
        
        # Mark minimum
        opt = find_optimal_point(df, lam)
        ax4.plot(opt['years'], opt['total_cost'] / 1000, 'o', color=color, markersize=10,
                markeredgecolor='black', markeredgewidth=1.5)
    
    ax4.set_xlabel('Construction Timeline T (years)', fontsize=12)
    ax4.set_ylabel('Total Cost J = E + λT (×10³ PJ)', fontsize=12)
    ax4.set_title('Total Cost Function for λ = 400-600 PJ/year', fontsize=13)
    ax4.legend(loc='upper right', fontsize=10)
    ax4.grid(True, alpha=0.3)
    ax4.set_xlim(95, 200)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"Lambda analysis plot saved to: {save_path}")
    
    return fig


def plot_decision_recommendation(
    df: pd.DataFrame,
    save_path: str = '/Volumes/Files/code/mm/20260130_b/p1/lambda_decision_map.png'
):
    """
    Decision map showing optimal choice based on λ preference
    """
    fig, ax = plt.subplots(figsize=(12, 8))
    
    years = df['years'].values
    energy = df['energy_PJ'].values
    
    # Main trade-off curve
    ax.plot(years, energy, 'b-', linewidth=3, label='Feasible Trade-off Curve')
    
    # Key boundaries
    T_elev = TOTAL_PAYLOAD / TOTAL_ELEVATOR_CAPACITY
    T_min = years.min()
    
    # Shade decision regions
    # Region 1: Low λ (cost priority) → longer timeline
    ax.axvspan(160, 190, alpha=0.2, color='green', label='Low λ (<150): Cost Priority')
    # Region 2: Medium λ (balanced) → middle ground
    ax.axvspan(130, 160, alpha=0.2, color='yellow', label='Medium λ (150-250): Balanced')
    # Region 3: High λ (time priority) → shorter timeline
    ax.axvspan(100, 130, alpha=0.2, color='red', label='High λ (>250): Time Priority')
    
    # Mark specific strategy points
    strategies = [
        {'name': 'A: Cost-Prioritized', 'years': 186, 'lambda': 0, 'color': 'green'},
        {'name': 'C: Balanced (λ≈200)', 'years': 139, 'lambda': 200, 'color': 'orange'},
        {'name': 'B: Time-Prioritized', 'years': 101, 'lambda': 500, 'color': 'red'},
    ]
    
    for s in strategies:
        idx = np.argmin(np.abs(years - s['years']))
        ax.plot(years[idx], energy[idx], 'o', color=s['color'], markersize=15, 
                markeredgecolor='black', markeredgewidth=2)
        ax.annotate(s['name'], (years[idx], energy[idx]), 
                   textcoords="offset points", xytext=(10, 10), fontsize=11,
                   fontweight='bold', color=s['color'])
    
    # Add decision guidance text
    textstr = '\n'.join([
        'Decision Guidance:',
        '─────────────────',
        'λ < 150 PJ/year → Strategy A',
        '  (Long-term cost efficiency)',
        '',
        '150 ≤ λ ≤ 250 → Strategy C',
        '  (Balanced trade-off)',
        '',
        'λ > 250 PJ/year → Strategy B',
        '  (Time-critical mission)',
    ])
    props = dict(boxstyle='round', facecolor='wheat', alpha=0.8)
    ax.text(0.02, 0.98, textstr, transform=ax.transAxes, fontsize=10,
            verticalalignment='top', bbox=props, family='monospace')
    
    ax.set_xlabel('Construction Timeline T (years)', fontsize=13)
    ax.set_ylabel('Total Energy E (PJ)', fontsize=13)
    ax.set_title('Decision Map: Optimal Strategy Selection Based on Time Opportunity Cost λ', fontsize=14)
    ax.legend(loc='upper right', fontsize=10)
    ax.grid(True, alpha=0.3)
    ax.set_xlim(95, 200)
    ax.set_ylim(10000, 65000)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"Decision map saved to: {save_path}")
    
    return fig


# ============== Critical Point Analysis ==============

def analyze_curve_structure(df: pd.DataFrame) -> Dict:
    """
    Analyze the structure of the trade-off curve to understand
    why certain λ values lead to specific optimal points.
    """
    years = df['years'].values
    energy = df['energy_PJ'].values
    
    # Calculate marginal rates
    marginal = -np.gradient(energy, years)
    
    # Find the elevator-only boundary
    T_elev = TOTAL_PAYLOAD / TOTAL_ELEVATOR_CAPACITY
    idx_elev = np.argmin(np.abs(years - T_elev))
    
    # Analyze marginal rate distribution
    # Region 1: T > T_elev (pure elevator, marginal ≈ 0)
    # Region 2: T < T_elev (need rockets, marginal > 0)
    
    # Find critical points where marginal rate changes significantly
    marginal_near_elev = marginal[idx_elev - 5:idx_elev + 5]
    marginal_at_139 = marginal[np.argmin(np.abs(years - 139))]
    marginal_at_101 = marginal[np.argmin(np.abs(years - 101))]
    
    return {
        'T_elev': T_elev,
        'idx_elev': idx_elev,
        'marginal_at_elev': marginal[idx_elev],
        'marginal_at_139': marginal_at_139,
        'marginal_at_101': marginal_at_101,
        'energy_at_elev': energy[idx_elev],
        'energy_at_139': energy[np.argmin(np.abs(years - 139))],
        'energy_at_101': energy[np.argmin(np.abs(years - 101))],
    }


def find_critical_lambda(df: pd.DataFrame) -> Dict:
    """
    Find the critical λ values where optimal point transitions occur.
    """
    years = df['years'].values
    energy = df['energy_PJ'].values
    
    # Dense lambda scan
    lambda_range = np.linspace(1, 1000, 2000)
    
    transitions = []
    prev_years = None
    
    for lam in lambda_range:
        opt = find_optimal_point(df, lam)
        if prev_years is not None and abs(opt['years'] - prev_years) > 5:
            transitions.append({
                'lambda': lam,
                'from_years': prev_years,
                'to_years': opt['years'],
            })
        prev_years = opt['years']
    
    return transitions


def plot_comprehensive_analysis(
    df: pd.DataFrame,
    save_path: str = '/Volumes/Files/code/mm/20260130_b/p1/lambda_comprehensive.png'
):
    """
    Comprehensive plot showing:
    1. Trade-off curve with marginal rates
    2. Critical λ transitions
    3. Decision regions
    """
    fig = plt.figure(figsize=(16, 14))
    
    years = df['years'].values
    energy = df['energy_PJ'].values
    marginal = calculate_marginal_energy_saving(df)
    
    T_elev = TOTAL_PAYLOAD / TOTAL_ELEVATOR_CAPACITY
    T_min = years.min()
    
    # ========== Plot 1: Trade-off curve with annotations ==========
    ax1 = fig.add_subplot(2, 2, 1)
    
    ax1.plot(years, energy / 1000, 'b-', linewidth=2.5, label='Trade-off Curve')
    
    # Mark key points
    key_points = [
        (T_min, 'Minimum Time\n(~101 years)', 'red'),
        (139, 'Original Knee\n(139 years)', 'orange'),
        (T_elev, 'Elevator-Only\n(~186 years)', 'green'),
    ]
    
    for t, label, color in key_points:
        idx = np.argmin(np.abs(years - t))
        ax1.plot(years[idx], energy[idx] / 1000, 'o', color=color, markersize=12,
                markeredgecolor='black', markeredgewidth=2, zorder=5)
        ax1.annotate(label, (years[idx], energy[idx] / 1000),
                    textcoords="offset points", xytext=(10, 10), fontsize=10,
                    color=color, fontweight='bold')
    
    ax1.set_xlabel('Construction Timeline T (years)', fontsize=12)
    ax1.set_ylabel('Total Energy E (×10³ PJ)', fontsize=12)
    ax1.set_title('Energy-Time Trade-off Curve', fontsize=14)
    ax1.grid(True, alpha=0.3)
    ax1.set_xlim(95, 200)
    ax1.legend(loc='upper right')
    
    # ========== Plot 2: Marginal Rate Analysis ==========
    ax2 = fig.add_subplot(2, 2, 2)
    
    ax2.plot(years, marginal, 'r-', linewidth=2, label='Marginal Saving Rate')
    ax2.fill_between(years, 0, marginal, alpha=0.3, color='red')
    
    # Mark critical thresholds
    ax2.axhline(y=480, color='purple', linestyle='--', linewidth=2,
               label='Critical λ ≈ 480 PJ/year')
    ax2.axvline(x=T_elev, color='green', linestyle=':', alpha=0.7)
    
    # Annotate the jump at elevator boundary
    ax2.annotate('Marginal rate jumps\nat elevator capacity limit',
                xy=(T_elev, marginal[np.argmin(np.abs(years - T_elev))]),
                xytext=(150, 400), fontsize=10,
                arrowprops=dict(arrowstyle='->', color='black'))
    
    ax2.set_xlabel('Construction Timeline T (years)', fontsize=12)
    ax2.set_ylabel('Marginal Energy Saving -dE/dT (PJ/year)', fontsize=12)
    ax2.set_title('Marginal Analysis: Why 186 years is optimal for most λ', fontsize=14)
    ax2.grid(True, alpha=0.3)
    ax2.set_xlim(95, 200)
    ax2.set_ylim(-50, 600)
    ax2.legend(loc='upper right')
    
    # ========== Plot 3: λ Sensitivity with Phase Transitions (400-600 focus) ==========
    ax3 = fig.add_subplot(2, 2, 3)
    
    lambda_range = np.linspace(400, 600, 200)
    opt_years = []
    for lam in lambda_range:
        opt = find_optimal_point(df, lam)
        opt_years.append(opt['years'])
    
    ax3.plot(lambda_range, opt_years, 'b-', linewidth=2.5)
    
    # Shade phases
    ax3.axhspan(180, 190, alpha=0.3, color='green', label='Phase 1: Elevator-Only (λ < 480)')
    ax3.axhspan(100, 145, alpha=0.3, color='red', label='Phase 2: Hybrid (λ > 480)')
    
    # Mark critical λ
    ax3.axvline(x=480, color='purple', linestyle='--', linewidth=2)
    ax3.annotate('Critical Transition\nλ ≈ 480 PJ/year',
                xy=(480, 160), fontsize=11, color='purple', fontweight='bold',
                ha='center')
    
    # Mark 139 year point
    ax3.axhline(y=139, color='orange', linestyle=':', linewidth=2, alpha=0.8)
    ax3.scatter([500], [139], s=200, c='orange', marker='*', zorder=5, 
               edgecolors='black', linewidths=2, label='T*=139y at λ≈500')
    
    ax3.set_xlabel('Time Opportunity Cost λ (PJ/year)', fontsize=12)
    ax3.set_ylabel('Optimal Timeline T* (years)', fontsize=12)
    ax3.set_title('Phase Transition in Optimal Strategy (λ=400-600)', fontsize=14)
    ax3.grid(True, alpha=0.3)
    ax3.set_xlim(400, 600)
    ax3.set_ylim(95, 195)
    ax3.legend(loc='upper right')
    
    # ========== Plot 4: Decision Summary Table ==========
    ax4 = fig.add_subplot(2, 2, 4)
    ax4.axis('off')
    
    # Create summary table
    summary_text = """
    ┌─────────────────────────────────────────────────────────────────┐
    │            KEY FINDINGS FROM λ COST ANALYSIS                    │
    ├─────────────────────────────────────────────────────────────────┤
    │                                                                 │
    │  1. CURVE STRUCTURE                                             │
    │     • Sharp discontinuity at T = 186 years (elevator capacity)  │
    │     • Marginal rate jumps from ~0 to ~480 PJ/year at boundary   │
    │                                                                 │
    │  2. OPTIMAL POINT SELECTION                                     │
    │     • λ < 480 PJ/year  →  T* = 186 years (elevator-only)        │
    │     • λ ≈ 500 PJ/year  →  T* = 139 years (original knee)        │
    │     • λ > 600 PJ/year  →  T* → 101 years (time-priority)        │
    │                                                                 │
    │  3. IMPLICATIONS FOR THE 139-YEAR KNEE POINT                    │
    │     • Requires implicit assumption: λ ≈ 500 PJ/year             │
    │     • This means: 1 year delay costs ~500 PJ opportunity cost   │
    │     • Equivalent to: ~0.5% of total elevator energy per year    │
    │                                                                 │
    │  4. RECOMMENDATION                                              │
    │     • Either justify λ ≈ 500 PJ/year with physical reasoning    │
    │     • Or acknowledge 186 years as cost-optimal baseline         │
    │     • Present 139 years as a time-constrained scenario          │
    │                                                                 │
    └─────────────────────────────────────────────────────────────────┘
    """
    
    ax4.text(0.05, 0.95, summary_text, transform=ax4.transAxes,
            fontsize=11, verticalalignment='top', family='monospace',
            bbox=dict(boxstyle='round', facecolor='lightyellow', alpha=0.8))
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"Comprehensive analysis saved to: {save_path}")
    
    return fig


# ============== Main Analysis ==============

def main():
    print("=" * 70)
    print("Lambda Time Cost Analysis for Moon Colony Logistics")
    print("=" * 70)
    
    # Generate trade-off curve
    print("\n[1] Generating Energy-Time trade-off curve...")
    df = generate_tradeoff_curve(year_min=100, year_max=220, num_points=500)
    print(f"    Generated {len(df)} data points")
    
    # Key statistics
    T_min = df['years'].min()
    T_max = df['years'].max()
    E_min = df['energy_PJ'].min()
    E_max = df['energy_PJ'].max()
    T_elev = TOTAL_PAYLOAD / TOTAL_ELEVATOR_CAPACITY
    
    print(f"\n[2] Trade-off Curve Statistics:")
    print(f"    Timeline range: {T_min:.1f} - {T_max:.1f} years")
    print(f"    Energy range: {E_min:.0f} - {E_max:.0f} PJ")
    print(f"    Elevator-only timeline: {T_elev:.1f} years")
    
    # Curve structure analysis
    print("\n[3] Curve Structure Analysis:")
    structure = analyze_curve_structure(df)
    print(f"    At elevator boundary (T={structure['T_elev']:.1f}y):")
    print(f"      - Energy: {structure['energy_at_elev']:.0f} PJ")
    print(f"      - Marginal rate: {structure['marginal_at_elev']:.1f} PJ/year")
    print(f"    At T=139 years:")
    print(f"      - Energy: {structure['energy_at_139']:.0f} PJ")
    print(f"      - Marginal rate: {structure['marginal_at_139']:.1f} PJ/year")
    print(f"    At T=101 years:")
    print(f"      - Energy: {structure['energy_at_101']:.0f} PJ")
    print(f"      - Marginal rate: {structure['marginal_at_101']:.1f} PJ/year")
    
    # Find critical transitions
    print("\n[4] Critical λ Transitions:")
    transitions = find_critical_lambda(df)
    for t in transitions:
        print(f"    λ ≈ {t['lambda']:.0f} PJ/year: T* jumps from {t['from_years']:.0f}y to {t['to_years']:.0f}y")
    
    # Sensitivity analysis
    print("\n[5] Sensitivity Analysis Results:")
    print("    " + "-" * 55)
    print(f"    {'λ (PJ/year)':<15} {'Optimal T (years)':<20} {'Energy (PJ)':<15}")
    print("    " + "-" * 55)
    
    test_lambdas = [100, 200, 300, 400, 450, 480, 500, 550, 600]
    sensitivity_results = []
    
    for lam in test_lambdas:
        opt = find_optimal_point(df, lam)
        print(f"    {lam:<15.0f} {opt['years']:<20.1f} {opt['energy_PJ']:<15.0f}")
        sensitivity_results.append(opt)
    
    # Find λ that gives 139 years
    print("\n[6] Original Knee Point Analysis (139 years):")
    lambda_for_139 = None
    for lam in np.linspace(400, 600, 1000):
        opt = find_optimal_point(df, lam)
        if abs(opt['years'] - 139) < 1:
            lambda_for_139 = lam
            print(f"    λ ≈ {lam:.0f} PJ/year corresponds to T* ≈ 139 years")
            break
    
    if lambda_for_139:
        print(f"\n    INTERPRETATION:")
        print(f"    To justify 139-year knee point, one must argue that:")
        print(f"    • Each year of delay costs ~{lambda_for_139:.0f} PJ in opportunity")
        print(f"    • This is {lambda_for_139/E_min*100:.1f}% of minimum total energy per year")
        print(f"    • Over 47 years (139→186), this amounts to {lambda_for_139*47:.0f} PJ")
    
    # Generate plots
    print("\n[7] Generating visualization plots...")
    plot_lambda_analysis(df)
    plot_decision_recommendation(df)
    plot_comprehensive_analysis(df)
    
    # Save sensitivity data
    sensitivity_df = sensitivity_analysis(df, np.linspace(50, 600, 120))
    sensitivity_df.to_csv('/Volumes/Files/code/mm/20260130_b/p1/lambda_sensitivity.csv', index=False)
    print("    Data saved to: lambda_sensitivity.csv")
    
    # Paper modification recommendations
    print("\n" + "=" * 70)
    print("RECOMMENDATIONS FOR PAPER MODIFICATION")
    print("=" * 70)
    print("""
    1. REFRAME THE OPTIMIZATION PROBLEM
       Current: "Multi-objective Pareto optimization"
       Suggested: "Constrained optimization with time-cost trade-off"
       
    2. INTRODUCE λ AS DECISION PARAMETER
       Add equation: J = E_total + λ × T
       where λ represents "time opportunity cost" (PJ/year)
       
    3. JUSTIFY THE KNEE POINT SELECTION
       Option A: Argue λ ≈ 480-500 PJ/year is reasonable because:
                 - Delayed lunar resource extraction
                 - Extended Earth-side operational costs
                 - Strategic/geopolitical value of early completion
       
       Option B: Present multiple scenarios:
                 - "Cost-optimal" (λ < 480): T* = 186 years
                 - "Balanced" (λ ≈ 500): T* = 139 years  
                 - "Time-critical" (λ > 600): T* → 101 years
    
    4. ADD SENSITIVITY ANALYSIS FIGURE
       Show how optimal T* changes with λ (Figure generated)
       
    5. ACKNOWLEDGE THE DISCONTINUITY
       Note that the trade-off curve has a sharp transition at
       T = 186 years due to elevator capacity constraints
    """)
    
    print("=" * 70)
    print("Analysis Complete!")
    print("=" * 70)
    
    return df, sensitivity_results


def plot_total_cost_standalone(
    df: pd.DataFrame = None,
    save_path: str = '/Volumes/Files/code/mm/20260130_b/p1/total_cost_curves.png'
):
    """
    单独绘制总成本曲线图
    - 长宽比 0.618
    - 删去标题
    - 删去 λ=500，添加目标 λ=504
    - 低偏中饱和度色系
    """
    if df is None:
        df = generate_tradeoff_curve()
    
    # 使用0.618黄金比例
    fig_width = 8
    fig_height = fig_width * 0.618
    fig, ax = plt.subplots(figsize=(fig_width, fig_height))
    
    years = df['years'].values
    
    # 低偏中饱和度配色（删去500，添加504）
    lambda_colors = [
        (450, '#6B9B78'),  # 灰绿色
        (480, '#B87B6B'),  # 灰橙红色
        (504, '#8B7BA8'),  # 灰紫色（目标λ）
        (550, '#5B8FA8'),  # 灰蓝色
    ]
    
    for lam, color in lambda_colors:
        total_cost = calculate_total_cost(df, lam)
        linewidth = 2.5 if lam == 504 else 2
        linestyle = '-'
        label = f'λ={lam} PJ/year' if lam != 504 else f'λ={lam} PJ/year (target)'
        
        ax.plot(years, total_cost / 1000, color=color, linestyle=linestyle, 
                linewidth=linewidth, label=label)
        
        # 标记最小值点
        opt = find_optimal_point(df, lam)
        markersize = 11 if lam == 504 else 9
        ax.plot(opt['years'], opt['total_cost'] / 1000, 'o', color=color, 
                markersize=markersize, markeredgecolor='#333333', markeredgewidth=1.5)
    
    ax.set_xlabel('Construction Timeline T (years)', fontsize=11)
    ax.set_ylabel('Total Cost J = E + λT (×10³ PJ)', fontsize=11)
    ax.legend(loc='upper right', fontsize=9)
    ax.grid(True, alpha=0.3)
    ax.set_xlim(95, 200)
    ax.set_ylim(95, 130)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"Total cost curves saved to: {save_path}")
    
    return fig


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