import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

# 设置美赛O奖风格
plt.rcParams['font.family'] = 'Times New Roman'
plt.rcParams['font.size'] = 10
plt.rcParams['axes.labelsize'] = 10
plt.rcParams['axes.titlesize'] = 11
plt.rcParams['xtick.labelsize'] = 10
plt.rcParams['ytick.labelsize'] = 10
plt.rcParams['legend.fontsize'] = 9
plt.rcParams['figure.titlesize'] = 14
plt.rcParams['text.usetex'] = False
plt.rcParams['axes.unicode_minus'] = False

# 使用seaborn美化
sns.set_style("whitegrid")
sns.set_context("paper")

# 数据准备
scenarios = ['Gaming', 'Navigation', 'Movie', 'Chatting', 'Screen Off']
start_caps = [100, 75, 50, 25]
titles_cn = ['大型游戏', '地图导航', '在线观影', '社交聊天', '熄屏待机']

# 时间数据 (h) - 5行4列
data_matrix = [
    [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
]

# 模拟参数
n = 1.7 # 曲线弯曲程度

# 颜色设置 (为每个场景分配不同的专业配色)
# Gaming: 红色系 (高能耗)
# Navigation: 橙色系
# Movie: 紫色系
# Chatting: 蓝色系
# Screen Off: 绿色系 (低能耗)
colors = ['#D32F2F', '#E65100', '#512DA8', '#1976D2', '#388E3C']

# 2. 创建图表 (5行4列)
fig, axes = plt.subplots(5, 4, figsize=(14, 15), dpi=300, constrained_layout=True)

# 遍历生成子图
for i, scenario in enumerate(scenarios):
    scenario_color = colors[i]  # 获取当前场景的颜色
    for j, start_cap in enumerate(start_caps):
        ax = axes[i, j]
        t_verify = data_matrix[i][j]
        
        # 准备曲线数据
        # 采用混合幂律模型模拟电池放电特性：前期平缓，后期陡峭
        # y = Start - k * (w1 * x^n1 + w2 * x^n2)
        if start_cap <= 5:
            # 理论上不应该发生，但作为保护
            x = np.linspace(0, 1, 100)
            y = np.full_like(x, start_cap)
        else:
            x = np.linspace(0, t_verify, 200)
            
            # 模型参数：n1控制前期线性度，n2控制后期陡峭度
            n1 = 1.2  # 接近线性，保持前期一致
            n2 = 5.0  # 高阶项，制造后期"跳水"效果
            w1 = 0.6  # 线性项权重
            w2 = 0.4  # 陡峭项权重
            
            # 计算归一化的形状函数 (x/t)^n 比直接用 x^n 更稳健
            # shape = w1 * (x/t)^n1 + w2 * (x/t)^n2
            # 终点处 shape = w1 + w2 = 1 (如果w1+w2=1)
            # 所以 Delta = Start - 5
            # y = Start - Delta * shape
            
            # 使用归一化时间 tau = x / t_verify，避免数值量级问题
            tau = x / t_verify
            shape_func = w1 * (tau ** n1) + w2 * (tau ** n2)
            
            # 缩放系数确保终点为5%
            # y_end = Start - k * shape_func(1) = 5
            # k = (Start - 5) / (w1 + w2)
            # 如果 w1+w2 = 1, 则 k = Start - 5
            
            delta = start_cap - 5
            y = start_cap - delta * shape_func

        # 3. 绘制主曲线
        ax.plot(x, y, color=scenario_color, linewidth=2.5, label=f'$T_{{verify}}={t_verify}$ h', zorder=3)

        # 4. 绘制阴影区域
        ax.fill_between(x, y, 0, color=scenario_color, alpha=0.15, zorder=1)

        # 5. 绘制红色虚线 (阈值线) - 使用深灰色作为通用阈值线，避免颜色冲突
        ax.axhline(y=5, color='#424242', linestyle='--', linewidth=1.5, zorder=2) 

        # 6. 设置标题 (仅第一行显示列标题)
        if i == 0:
            ax.set_title(f"Initial Charge: {start_cap}%", fontsize=12, fontweight='bold', pad=10)

        # 7. 设置坐标轴标签
        # 最后一行显示X轴标签
        if i == len(scenarios) - 1:
            ax.set_xlabel('Time (hours)', fontsize=10, fontweight='bold')
        
        # 第一列显示Y轴标签（包含场景名称，仿照参考图风格）
        if j == 0:
            # 使用多行文本，第一行是场景名，第二行是单位
            label_text = f"{titles_cn[i]}\nBattery (%)" if 'titles_cn' in globals() else f"{scenario}\nBattery (%)"
            # 如果 titles_cn 未定义，回退到 scenario 英文名. 
            # 检查上下文，发现 titles_cn = ['大型游戏', ...] 已经定义在前面了 (Line 23)
            # 但用户给的数据是中文名+英文名, 之前的代码 titles_cn = ... 定义了.
            # 让我检查下 2.py 的内容
            ax.set_ylabel(f"{scenario}\nSOC (%)", fontsize=11, fontweight='bold')

        # 8. 设置坐标轴范围
        # x轴动态范围，留出20%空间给图例
        ax.set_xlim(0, t_verify * 1.3) 
        ax.set_ylim(0, 105)

        # 9. 设置网格
        ax.grid(True, linestyle=':', alpha=0.3, color='gray', linewidth=0.8, zorder=0)
        ax.set_axisbelow(True)

        # 10. 设置图例
        # 把T_verify放在图例里，Threshold线就不放了，或者都放但字体缩小
        ax.legend(loc='upper right', fontsize=9, frameon=True, 
                 edgecolor='black', fancybox=False, shadow=False, framealpha=0.9)

        # 11. 美化边框
        for spine in ax.spines.values():
            spine.set_linewidth(1.0)
            spine.set_color('black')
        ax.tick_params(labelsize=9, width=1.0, length=4, direction='out')

# 保存
plt.savefig('2.png', dpi=300, bbox_inches='tight', facecolor='white')
plt.savefig('2.pdf', bbox_inches='tight', facecolor='white')
print("图表已生成：2.png")