import matplotlib.pyplot as plt
import numpy as np
import os
from matplotlib import font_manager

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 数据
pollution_levels = ['原始', '0.5倍', '1倍', '1.5倍', '2倍']

vector_search = {
    'hit_rate': [91.24, 95.99, 92.80, 86.45, 72.10],
    'mrr': [97.66, 96.26, 95.67, 94.98, 96.23]
}

keyword_search = {
    'hit_rate': [98.77, 98.28, 96.86, 93.93, 84.14],
    'mrr': [96.97, 96.02, 94.03, 92.18, 86.23]
}

hybrid_search = {
    'hit_rate': [99.80, 99.62, 99.37, 94.49, 94.10],
    'mrr': [94.14, 94.55, 92.09, 89.42, 82.12]
}

hybrid_search_after_decoupling = {
    'hit_rate': [99.71, 99.48, 99.04, 97.89, 92.19],
    'mrr': [95.88, 94.12, 91.65, 87.90, 79.86]
}

# 创建图表
fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4, figsize=(24, 6))
# fig.suptitle('搜索方法比较实验', fontsize=20)


# 函数来绘制单个子图
def plot_subplot(ax, data, title):
    ax.plot(pollution_levels, data['hit_rate'], 'ro-', label='命中率')
    ax.plot(pollution_levels, data['mrr'], 'bo-', label='平均倒数排名')
    ax.set_title(title, fontsize=20)
    ax.set_ylim(70, 100)
    ax.set_xlabel('干扰程度', fontsize=18)
    ax.set_ylabel('百分比', fontsize=18)
    ax.legend(fontsize=16)
    ax.grid(True, linestyle='--', alpha=0.7)

    ax.tick_params(axis='both', which='major', labelsize=16)

    # 设置x轴标签旋转
    ax.set_xticklabels(pollution_levels, ha='right')


# 绘制子图
plot_subplot(ax1, vector_search, '向量搜索')
plot_subplot(ax2, keyword_search, '关键字搜索')
plot_subplot(ax3, hybrid_search, '混合搜索')
plot_subplot(ax4, hybrid_search_after_decoupling, '混合搜索_架构拆解')

# 调整布局
plt.tight_layout()

# 保存图表
file_path = os.path.join(os.getcwd(), 'search_methods_comparison.png')
plt.savefig(file_path, dpi=300, bbox_inches='tight')
print(f"图表已保存为: {file_path}")

# 检查文件是否成功创建
if os.path.exists(file_path):
    print(f"文件成功创建在: {file_path}")
    print(f"文件大小: {os.path.getsize(file_path)} 字节")
else:
    print(f"文件创建失败，请检查权限或其他问题")