import matplotlib.pyplot as plt
import numpy as np

# Define the data
model_names = ['DenseNet121']
TED_rand = [9.98]
TED_const = [33.86]
TED_LVX = [9.38]
MCS_rand = [22.93]
MCS_const = [16.82]
MCS_LVX = [24.17]
tree_kernel_scores_rand = [45.78]
tree_kernel_scores_const = [45.06]
tree_kernel_scores_lvx = [47.23]


# Set the figure size
# plt.figure(figsize=(10, 4))

# Define the colors for the three methods
color_rand = '#2ca02c'  # Green
color_const = '#1f77b4'  # Blue
color_LVX = '#ff7f0e'  # Orange

# color_rand = '#4e79a7'  # Dark blue
# color_const = '#f28e2b'  # Orange
# color_LVX = '#e15759'  # Red

highlight_color_LVX = '#ff5500'  # Stronger orange
highlight_width_LVX = 3

# Set the bar positions
bar_width = 0.6
index = np.arange(len(model_names))*2 
bias = 0.7
figsize = (3, 6)
# Plot the TED indicators
fig, ax = plt.subplots(figsize=figsize)
ax.bar(index, TED_const, width=bar_width, label='Constant', alpha=0.7, color=color_const)
ax.bar(index+bar_width, TED_rand, width=bar_width, label='Random', alpha=0.6,color=color_rand)
ax.bar(index+bar_width*2, TED_LVX, width=bar_width, label='LVX', alpha=0.5, edgecolor=highlight_color_LVX, linewidth=highlight_width_LVX, color=color_LVX)

ax.set_ylim(5, 35)
# Customize the plot
ax.set_xlabel('Model', fontsize=40, fontname='Times New Roman')
ax.set_ylabel('TED', fontsize=40, fontname='Times New Roman')
# ax.set_title('Baseline Comparison on CIFAR-10 Dataset - TED')
ax.set_xticks(index+ bar_width)
ax.set_xticklabels(model_names, rotation=0, fontsize=24, fontname='Times New Roman')
# Set the font properties for the legend
font_properties = {'family': 'Times New Roman', 'size': 26}
# ax.legend(prop=font_properties,loc='upper center', bbox_to_anchor=(0.5, 1.0), ncol=3)

# Increase font size for tick labels
ax.tick_params(axis='both', labelsize=18)
# Show the TED plot
plt.tight_layout()
# plt.show()
plt.savefig('tools/plot/xray_ted.pdf')



fig, ax = plt.subplots(figsize=figsize)
ax.bar(index, MCS_const, width=bar_width, label='Constant', alpha=0.7, color=color_const)
ax.bar(index+bar_width, MCS_rand, width=bar_width, label='Random', alpha=0.6,color=color_rand)
ax.bar(index+bar_width*2, MCS_LVX, width=bar_width, label='LVX', alpha=0.5, edgecolor=highlight_color_LVX, linewidth=highlight_width_LVX,  color=color_LVX)

ax.set_ylim(10, 26)
# Customize the plot
ax.set_xlabel('Model', fontsize=40, fontname='Times New Roman')
ax.set_ylabel('Normalized MCS', fontsize=40, fontname='Times New Roman')
# ax.set_title('Baseline Comparison on CIFAR-10 Dataset - TED')
ax.set_xticks(index+ bar_width)

ax.set_xticklabels(model_names, rotation=0, fontsize=24, fontname='Times New Roman')
# Set the font properties for the legend

# ax.legend(prop=font_properties,loc='upper center', bbox_to_anchor=(0.5, 1.3), ncol=3)

# Increase font size for tick labels
ax.tick_params(axis='both', labelsize=18)
# Show the TED plot
plt.tight_layout()
# plt.show()
plt.savefig('tools/plot/xray_mcs.pdf')

fig, ax = plt.subplots(figsize=figsize)
ax.bar(index, tree_kernel_scores_const, width=bar_width, label='Constant', alpha=0.7, color=color_const)
ax.bar(index+bar_width, tree_kernel_scores_rand, width=bar_width, label='Random', alpha=0.6,color=color_rand)
ax.bar(index+bar_width*2, tree_kernel_scores_lvx, width=bar_width, label='LVX', alpha=0.5, edgecolor=highlight_color_LVX, linewidth=highlight_width_LVX,  color=color_LVX)

# Customize the plot
ax.set_xlabel('Model', fontsize=40, fontname='Times New Roman')
ax.set_ylabel('Normalized TK', fontsize=40, fontname='Times New Roman')
# ax.set_title('Baseline Comparison on CIFAR-10 Dataset - TED')
ax.set_xticks(index+ bar_width)
ax.set_xticklabels(model_names, rotation=0, fontsize=24, fontname='Times New Roman')
# Set the font properties for the legend
# ax.legend(prop=font_properties,loc='upper center', bbox_to_anchor=(0.5, 1.0), ncol=3)
ax.set_ylim(43, 49)
# Increase font size for tick labels
ax.tick_params(axis='both', labelsize=18)
# Show the TED plot
plt.tight_layout()
# plt.show()
plt.savefig('tools/plot/xray_tk.pdf')