import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import matplotlib.patches as patches
import numpy as np

# --- 准备工作：设置全局字体和样式以匹配图片 ---
# Matplotlib 默认可能没有 Helvetica，'Arial' 是一个很好的替代品
plt.rcParams['font.family'] = 'Arial'
plt.rcParams['font.weight'] = 'bold'
plt.rcParams['axes.labelweight'] = 'bold'
plt.rcParams['axes.titleweight'] = 'bold'
plt.rcParams['axes.linewidth'] = 1.5 # 坐标轴线宽

# --- 创建 Figure 和 Subplots ---
# 创建一个一行两列的图，设置合适的尺寸
fig, axs = plt.subplots(1, 1, figsize=(7, 6))

# --- 数据估算 ---
# 根据图片估算两条曲线在两个图中的数据点
# 左图数据
data_left = {
    'flow_matching_x': np.array([50e3, 100e3, 200e3, 400e3, 800e3, 2e6, 4e6]),
    'flow_matching_y': np.array([80, 37, 25, 17.2, 10.8, 10.5, 7.5]),
    'contrastive_x': np.array([50e3, 100e3, 200e3, 400e3, 800e3, 1e6]),
    'contrastive_y': np.array([51, 18, 9, 6, 4.5, 3.5,])
}
# 右图数据
data_right = {
    'flow_matching_x': np.array([10, 25, 50, 100, 150, 200, 250]),
    'flow_matching_y': np.array([250, 36, 18, 15, 13.8, 12.9, 12.8]),
    'contrastive_x': np.array([10, 25, 50, 100, 150, 200, 250]),
    'contrastive_y': np.array([70, 12, 9.8, 9.4, 9.1, 8.8, 8.7])
}

# --- 绘制右侧子图 ---
ax_right = axs
ax_right.plot(data_right['flow_matching_x'], data_right['flow_matching_y'], 
              color='#d62728', marker='o', markersize=6, linewidth=2.5, label='DDT-XL/2')
ax_right.plot(data_right['contrastive_x'], data_right['contrastive_y'], 
              color='#1f77b4', marker='o', markersize=6, linewidth=2.5, label='DUPA-XL/2')

# 设置Y轴为对数坐标，X轴为线性坐标
ax_right.set_yscale('log')

# 设置坐标轴标签
ax_right.set_xlabel('Total Denoising Steps', fontsize=20)
ax_right.set_ylabel('FID-50K', fontsize=20)

# 设置Y轴刻度为普通数字
ax_right.yaxis.set_major_formatter(mticker.ScalarFormatter())
ax_right.set_yticks([10, 100])


# 添加水平虚线
ax_right.axhline(y=10, color='gray', linestyle='--', linewidth=1, alpha=0.7)
ax_right.axhline(y=100, color='gray', linestyle='--', linewidth=1, alpha=0.7)
# # 绘制带竖线的箭头
arrow_y_pos = 20
arrow_x_start = 25-1.5
arrow_x_end = 250+1.5
arrow = patches.FancyArrowPatch(
    (arrow_x_start, arrow_y_pos), (arrow_x_end, arrow_y_pos),
    arrowstyle='|-|,widthA=0.4,widthB=0.4', # 带竖线的样式
    mutation_scale=10, # 控制竖线的高度
    color='black',
    linewidth=1.5
)
ax_right.add_patch(arrow)
arrow = patches.FancyArrowPatch(
    (25, 12), (25, 20),
    arrowstyle='-', # 带竖线的样式
    linestyle='--',
    mutation_scale=10, # 控制竖线的高度
    color='#1f77b4',
    linewidth=2
)
ax_right.add_patch(arrow)
arrow = patches.FancyArrowPatch(
    (250, 12.8), (250, 20),
    arrowstyle='-', # 带竖线的样式
    linestyle='--',
    mutation_scale=10, # 控制竖线的高度
    color='#d62728',
    linewidth=2
)
ax_right.add_patch(arrow)

# 添加文字
ax_right.annotate('10x Faster', 
                 xy=((arrow_x_start + arrow_x_end) / 2, arrow_y_pos), 
                 xytext=(0, 5), textcoords='offset points',
                 ha='center', va='bottom', fontsize=14, fontweight='bold')

# --- 对两个子图进行统一的样式微调 ---
# for ax in axs:
axs.tick_params(axis='both', which='major', labelsize=14, direction='in', width=1.5, length=6)
axs.tick_params(axis='both', which='minor', direction='in', width=1, length=3)
axs.legend(fontsize=14, edgecolor='black')
        # 手动设置图例边框线宽
axs.get_legend().get_frame().set_linewidth(1.5)


# --- 调整整体布局并显示 ---
plt.tight_layout()
plt.show()
plt.savefig('zx.svg', format='svg', bbox_inches='tight')