#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Richards S-Curve Fit for 2010-2025 Data with K=4298

Fits Richards model to 2010-2025 launch data with carrying capacity
constrained to K=4298 (close to physical limit of 3650).
"""

import pandas as pd
import numpy as np
from scipy.optimize import curve_fit
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')

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


def richards(t, K, r, t0, v):
    """Richards curve (generalized logistic model)"""
    exp_term = np.exp(-r * (t - t0))
    exp_term = np.clip(exp_term, 1e-10, 1e10)
    return K / np.power(1 + exp_term, 1/v)


def richards_fixed_K(t, r, t0, v):
    """Richards curve with fixed K=4298"""
    K = 4298
    exp_term = np.exp(-r * (t - t0))
    exp_term = np.clip(exp_term, 1e-10, 1e10)
    return K / np.power(1 + exp_term, 1/v)


def load_data(filepath="rocket_launch_counts.csv"):
    """Load rocket launch data"""
    df = pd.read_csv(filepath)
    df = df.rename(columns={"YDate": "year", "Total": "launches"})
    df["year"] = pd.to_numeric(df["year"], errors="coerce")
    df["launches"] = pd.to_numeric(df["launches"], errors="coerce")
    df = df.dropna(subset=["year", "launches"])
    df = df[(df["year"] >= 1957) & (df["year"] <= 2025)]
    df = df.astype({"year": int, "launches": int})
    df = df.sort_values("year").reset_index(drop=True)
    return df


def main():
    print("=" * 60)
    print("Richards S-Curve Fit (2010-2025 Data, K=4298)")
    print("=" * 60)
    
    # Load data
    df = load_data()
    
    # Filter 2010-2025
    start_year = 2010
    end_year = 2025
    data = df[(df["year"] >= start_year) & (df["year"] <= end_year)].copy()
    years = data["year"].values
    launches = data["launches"].values
    
    print(f"Data range: {start_year} - {end_year}")
    print(f"Data points: {len(data)}")
    print(f"\nHistorical data:")
    for y, l in zip(years, launches):
        print(f"  {y}: {l} launches")
    
    # Fit with fixed K=4298
    K_fixed = 4298
    t = (years - start_year).astype(float)
    
    p0 = [0.2, 20.0, 2.0]  # r, t0, v
    bounds = ([0.01, 5, 0.5], [1.0, 100, 10.0])
    
    try:
        popt, pcov = curve_fit(richards_fixed_K, t, launches, p0=p0, bounds=bounds, maxfev=100000)
        r, t0, v = popt
        
        # Calculate R²
        y_pred = richards_fixed_K(t, *popt)
        ss_res = np.sum((launches - y_pred) ** 2)
        ss_tot = np.sum((launches - np.mean(launches)) ** 2)
        r_squared = 1 - (ss_res / ss_tot)
        
        print(f"\nFitted Parameters (K fixed at {K_fixed}):")
        print(f"  K (carrying capacity) = {K_fixed} launches/year (FIXED)")
        print(f"  r (growth rate)       = {r:.4f}")
        print(f"  t0 (inflection point) = {start_year + t0:.1f}")
        print(f"  v (shape parameter)   = {v:.3f}")
        print(f"  R²                    = {r_squared:.4f}")
        
    except Exception as e:
        print(f"Fitting error: {e}")
        return
    
    # Physical limit
    physical_max = 3650
    print(f"\nPhysical limit: {physical_max} (10 sites × 365 days)")
    print(f"K / Physical limit = {K_fixed/physical_max:.2f}x")
    
    # ========== Create Visualization ==========
    # 缩小图片尺寸（比例不变），使字体相对更大
    fig, ax = plt.subplots(figsize=(8, 4.67))
    
    # 中低饱和度配色
    color_data = '#5D6D7E'       # 灰蓝色 - 数据点
    color_curve = '#52796F'      # 暗绿色 - S曲线
    color_target = '#7B9EA8'     # 灰蓝绿色 - 2050标记
    
    # Historical data points
    ax.scatter(years, launches, color=color_data, s=80, alpha=0.85, 
               label='Historical Data (2010-2025)', zorder=4, edgecolor='white', linewidth=1)
    
    # Generate smooth S-curve
    years_smooth = np.linspace(start_year, 2055, 500)
    t_smooth = years_smooth - start_year
    pred_smooth = richards(t_smooth, K_fixed, r, t0, v)
    
    # S-curve prediction
    ax.plot(years_smooth, pred_smooth, color=color_curve, lw=2.5, 
            label=f'Richards Model (K={K_fixed}, R²={r_squared:.3f})', zorder=2)
    
    # K=4298 saturation line
    ax.axhline(K_fixed, color=color_curve, ls=':', lw=1.5, alpha=0.6,
               label=f'K = {K_fixed}')
    
    # Mark 2050 line only
    ax.axvline(2050, color=color_target, ls=':', lw=1.5, alpha=0.7)
    ax.text(2050.5, K_fixed*0.83, '2050\n(Target)', fontsize=9, color=color_target)
    
    # Only show 2050 prediction point
    t_2050 = 2050 - start_year
    pred_2050 = richards(t_2050, K_fixed, r, t0, v)
    ax.scatter([2050], [pred_2050], color=color_target, s=60, marker='D', zorder=4)
    ax.annotate(f'{pred_2050:.0f}', xy=(2050, pred_2050), 
                xytext=(2050.5, pred_2050+180),
                fontsize=9, color=color_target, fontweight='bold')
    
    # Formatting
    ax.set_xlabel('Year', fontsize=11)
    ax.set_ylabel('Annual Launches', fontsize=11)
    ax.legend(loc='upper left', fontsize=9)
    ax.grid(True, alpha=0.25)
    ax.set_xlim(2010, 2055)
    ax.set_ylim(0, K_fixed * 1.15)
    
    plt.tight_layout()
    plt.savefig('richards_curve_2010_K4298.png', dpi=150, bbox_inches='tight')
    plt.close()
    
    print("\nPlot saved: richards_curve_2010_K4298.png")
    print("=" * 60)


if __name__ == "__main__":
    main()