import matplotlib.pyplot as plt
import numpy as np

# Define the data
model_names = ['ResNet18', 'ResNet34', 'ResNet50', 'ResNet101', 'ViT-T16', 'ViT-S16', 'ViT-B16', 'ViT-L16']
TED_rand = [9.83, 9.75, 9.68, 9.64, 10.42, 9.69, 9.62, 9.45]
TED_const = [34.15, 33.78, 33.58, 33.48, 35.44, 33.61, 33.37, 32.84]
TED_LVX = [9.30, 9.17, 9.09, 9.04, 9.99, 9.10, 8.99, 8.79]
MCS_rand = [22.52, 23.74, 24.59, 24.94, 15.07, 24.16, 25.20, 27.27]
MCS_const = [16.82, 17.71, 18.35, 18.66, 11.25, 18.01, 18.79, 20.35]
MCS_LVX = [23.87, 25.09, 25.99, 26.51, 15.91, 25.73, 27.01, 29.29]
tree_kernel_scores_rand = [45.32, 47.66, 49.38, 50.25, 30.30, 48.53, 50.64, 54.83]
tree_kernel_scores_const = [44.85, 47.16, 48.97, 49.77, 29.92, 48.05, 50.22, 54.36]
tree_kernel_scores_lvx = [46.85, 49.24, 51.19, 51.99, 31.24, 50.34, 52.76, 57.14]


# 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 = (10, 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, 37)
# 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=30, 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/imagenet_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, 31)
# 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=30, 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/imagenet_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=30, 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(25, 58)
# 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/imagenet_tk.pdf')