"""
Figure: Battery SOC Trajectory under Different Usage Scenarios
MCM/ICM 2026 - Problem A
"""
import matplotlib.pyplot as plt
import numpy as np

# === MCM O-Award Style Configuration ===
plt.rcParams.update({
    'font.family': 'serif',
    'font.serif': ['Times New Roman'],
    'mathtext.fontset': 'stix',
    'axes.labelsize': 12,
    'axes.titlesize': 14,
    'xtick.labelsize': 11,
    'ytick.labelsize': 11,
    'legend.fontsize': 9,
    'axes.linewidth': 1.0,
    'axes.unicode_minus': False,
    'figure.dpi': 300,
})

# === Data Configuration ===
scenarios = ['Gaming', 'Navigation', 'Movie', 'Chatting', 'Screen Off']
power_values = [3.551, 2.954, 2.235, 1.481, 0.517]  # W

# Time-to-empty from different starting SOC [100%, 75%, 50%, 25%]
data_matrix = np.array([
    [4.11, 3.05, 2.01, 0.97],   # Gaming
    [5.01, 3.72, 2.45, 1.18],   # Navigation
    [6.63, 4.92, 3.24, 1.56],   # Movie
    [10.02, 7.43, 4.89, 2.36],  # Chatting
    [29.45, 21.85, 14.39, 6.95] # Screen Off
])
start_soc = [100, 75, 50, 25]
z_min = 0.02  # SOC threshold (2%)

# Professional color palette (colorblind-friendly, Nature-style)
colors = ['#E64B35', '#4DBBD5', '#00A087', '#3C5488', '#F39B7F']
markers = ['o', 's', '^', 'D', 'v']

# === Figure Setup ===
fig, ax = plt.subplots(figsize=(8, 5))

# Mixed power-law model parameters
n1, n2, w1, w2 = 1.2, 5.0, 0.6, 0.4

for i, scenario in enumerate(scenarios):
    t_end = data_matrix[i, 0]  # Time from 100% to z_min
    color = colors[i]
    marker = markers[i]
    
    # Mixed power-law SOC decay model
    t = np.linspace(0, t_end, 200)
    tau = t / t_end
    shape_func = w1 * (tau ** n1) + w2 * (tau ** n2)
    z = 100 - 98 * shape_func  # 100% -> 2%
    
    # Plot trajectory
    label = f'{scenario} ({power_values[i]:.3f} W)'
    ax.plot(t, z, color=color, linewidth=1.8, label=label, zorder=3)
    
    # Add markers along curve
    mark_indices = np.linspace(0, 199, 8, dtype=int)
    ax.scatter(t[mark_indices], z[mark_indices], color=color, marker=marker, 
               s=35, edgecolors='white', linewidths=0.5, zorder=4)

# Threshold line
ax.axhline(y=z_min*100, color='#B71C1C', linestyle='--', linewidth=1.5, 
           label=f'Cutoff Threshold ({z_min*100:.0f}%)', zorder=2)

# === Axis Configuration ===
ax.set_xlabel('Time $t$ (hours)', fontweight='bold')
ax.set_ylabel('State of Charge $z(t)$ (%)', fontweight='bold')
ax.set_xlim(0, 32)
ax.set_ylim(0, 105)
ax.set_xticks(np.arange(0, 35, 5))
ax.set_yticks(np.arange(0, 120, 20))

# Grid styling
ax.grid(True, linestyle='-', alpha=0.3, linewidth=0.5, color='gray')
ax.set_axisbelow(True)

# Legend (outside plot area for clarity)
ax.legend(loc='upper right', frameon=True, fancybox=False, 
          edgecolor='black', framealpha=0.95, ncol=1)

# Minor ticks
ax.minorticks_on()
ax.tick_params(which='minor', length=2, width=0.5)
ax.tick_params(which='major', length=4, width=1.0)

# === Output ===
plt.tight_layout()
plt.savefig('combined_soc_trajectory.png', dpi=300, bbox_inches='tight', 
            facecolor='white', edgecolor='none')
plt.savefig('combined_soc_trajectory.pdf', bbox_inches='tight',
            facecolor='white', edgecolor='none')
print("Figure saved: combined_soc_trajectory.png / .pdf")