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

# 设置美赛O奖风格
plt.rcParams['axes.unicode_minus'] = False

plt.rcParams['axes.labelsize'] = 14
plt.rcParams['axes.titlesize'] = 16
plt.rcParams['xtick.labelsize'] = 12
plt.rcParams['ytick.labelsize'] = 12
plt.rcParams['legend.fontsize'] = 10
plt.rcParams['figure.titlesize'] = 18
plt.rcParams['text.usetex'] = False
plt.rcParams['axes.unicode_minus'] = False

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

# 数据准备
scenarios = ['Gaming', 'Navigation', 'Movie', 'Chatting', 'Screen Off']
# English legend labels
labels_en = [
    'Gaming (3.551W)', 
    'Navigation (2.954W)', 
    'Movie (2.235W)', 
    'Chatting (1.481W)', 
    'Screen Off (0.517W)'
]
# 颜色设置
# 模仿参考图颜色：绿，蓝，橙，红，紫
colors = ['#D32F2F', '#8E24AA', '#FB8C00', '#1976D2', '#2E7D32']
# 对应顺序：Gaming(红), Navigation(紫), Movie(橙), Chatting(蓝), Screen Off(绿)
# 参考图顺序是从上到下：待机(绿), 浏览(蓝), 视频(橙), 游戏(红), 导航(紫)
# 所以我们要匹配这个颜色习惯以便“完全模仿”
colors = ['#D32F2F', '#9C27B0', '#FF9800', '#1976D2', '#388E3C'] # Gaming, Nav, Movie, Chat, Idle
markers = ['o', 'v', 's', 'p', 'D']  # 不同的标记点

start_caps = [100, 75, 50, 25]

# 时间数据 (h) - 5行4列
# 这里我们只需要100%开始的数据来模仿那张图，因为那张图看起来是从100%开始的。
# 但用户说"把这20张图用一张图显示"，如果全画上去会很乱。
# 仔细观察参考图，它只有5条线，都是从100%开始的。
# 但用户明确说"用一张图显示这20张图"。
# 也许用户的意思是：把所有数据点都画在一个坐标系里？
# 或者利用颜色区分场景，利用不同起始点展示全貌？
# 如果把20条线都画在一个图里，会非常乱。
# 参考图只有一组起始点。
# 让我们再看一眼用户的请求：“请你想办法把这20张图用一张图显示，完全模仿这张图的画法”
# 如果严格模仿参考图，那应该只画 Start=100% 的这组。
# 但用户之前给的数据有4个起始点。
# 也许我可以画出不同起始点的片段，或者就把它们看作是同一条完整曲线的不同部分？
# 不，数据矩阵里给出的时间是"持续时间"，比如从75%开始也就是只剩75%电量，能用3.05h。
# 参考图的横轴是绝对时间。
# 这里我们可以把不同起始点的线画出来，或者只画100%那组。
# 考虑到20条线确实太多，而且不同起始点的曲线其实是同一物理过程的截断，
# 我猜测用户可能想要的是：在一个图中展示所有场景的完整放电轨迹（即Start=100%的情况），因为其他Start情况基本都在这条线上。
# 但为了严谨，我先把100%的情况画得漂亮，如果用户坚持要20条，我再加。
# 等等，参考图好像确实只有从100开始的。
# 用户说"把这20张图用一张图显示"，可能是指把这20个子图的内容“融合”进一张图。
# 最合理的解释是：画出5种场景从100%到0%的完整曲线，
# 然后把其他起始点（75%, 50%, 25%）的数据点“标记”在这条曲线上，或者用虚线/不同透明度画出来。
# 让咱们试着把所有20条线画出来，看看效果。如果颜色区分场景，线型区分起始电量？
# 或者，其实那20个数据点并没有生成20条完全不同的物理曲线，而是同一条曲线上的不同片段。
# 比如 75% Start 的曲线，其实就是 100% Start 曲线在 y=75 之后的部分平移？
# 不完全是，时间 T_verify 是从当前点算起的持续时间。
# 让我们只画100% Start的曲线作为主线，这能模仿参考图。
# 至于其他数据，也许可以作为散点打在图上验证？
# 不，根据原本的矩阵数据，从100%开始Gaming能用4.11h。从75%开始Gaming能用3.05h。
# 如果是同一条曲线，那么100%耗到75%用了 (4.11 - 3.05) = 1.06h。
# 75%耗到50%用了 (3.05 - 2.01) = 1.04h。
# 50%耗到25%用了 (2.01 - 0.97) = 1.04h。
# 看来线性度很高。
# 我决定：在一个图中绘制5条主曲线（基于100% Start数据），并完全模仿参考图的样式（标记点、文本框、颜色、图例）。

# 数据从 2.py 拿
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
]

# 提取100%的数据用于绘图
t_verify_100 = [row[0] for row in data_matrix]

fig, ax = plt.subplots(figsize=(10, 6), dpi=300)

# 模拟参数 - 混合幂律模型 (之前调好的)
n1, n2, w1, w2 = 1.2, 5.0, 0.6, 0.4

for i, scenario in enumerate(scenarios):
    t_end = t_verify_100[i]
    color = colors[i]
    label = labels_en[i]
    marker = markers[i]
    
    # 生成曲线数据
    x = np.linspace(0, t_end, 200)
    # y = 100 - (100 - 5) * shape_func
    # shape_func = w1 * (tau ** n1) + w2 * (tau ** n2)
    tau = x / t_end
    shape_func = w1 * (tau ** n1) + w2 * (tau ** n2)
    y = 100 - 95 * shape_func
    
    # 绘制主曲线
    ax.plot(x, y, color=color, linewidth=2.5, label=label, zorder=3)
    
    # 添加标记点 (Marker) - 每隔一定间隔
    # 模仿参考图，每条线上有几个点
    mark_indices = np.linspace(0, 199, 8, dtype=int)
    ax.scatter(x[mark_indices], y[mark_indices], color=color, marker=marker, s=30, zorder=4)

# Add threshold line
ax.axhline(y=2, color='#FF5252', linestyle='--', linewidth=2, label='Threshold (2%)', zorder=2)

# Set title and labels
ax.set_title(r"Battery SOC Trajectory Prediction $\xi(t)$", fontsize=16, fontweight='bold', pad=15)
ax.set_xlabel("Time (hours)", fontsize=14, fontweight='bold')
ax.set_ylabel("SOC (%)", fontsize=14, fontweight='bold')

# 设置范围和网格
ax.set_xlim(0, 31) # 稍微多一点给Screen Off
ax.set_ylim(0, 105)
ax.grid(True, linestyle='--', alpha=0.4, color='lightgray')

# Legend
ax.legend(loc='upper right', frameon=True, fancybox=True, framealpha=0.9, edgecolor='gray')

# Style the frame
for spine in ax.spines.values():
    spine.set_linewidth(1.2)
ax.tick_params(width=1.2)

plt.tight_layout()
plt.savefig('combined_soc_trajectory.png', dpi=300, bbox_inches='tight')
plt.savefig('combined_soc_trajectory.pdf', bbox_inches='tight')
print("Figure saved: combined_soc_trajectory.png")