import matplotlib.pyplot as plt
import numpy as np

# Define the data
model_names = ['ResNet20', 'ResNet32', 'ResNet44', 'ResNet56', 'MBNv2-x0.5', 'MBNv2-x0.75', 'MBNv2-x1.0', 'MBNv2-x1.4', 'RepVGG A0', 'RepVGG A1', 'RepVGG A2']
TED_rand = [9.61, 9.57, 9.54, 9.52, 9.55, 9.43, 9.48, 9.43, 9.44, 9.42, 9.40]
TED_const = [28.77, 28.67, 28.60, 28.54, 28.58, 28.43, 28.35, 28.16, 28.28, 28.21, 28.08]
TED_LVX = [8.96, 8.86, 8.81, 8.83, 8.87, 8.76, 8.73, 8.65, 8.74, 8.72, 8.70]
MCS_rand = [22.65, 22.84, 23.48, 23.87, 23.43, 24.47, 24.35, 24.87, 24.65, 25.27, 25.46]
MCS_const = [17.60, 17.92, 18.34, 18.60, 18.19, 18.87, 19.02, 19.47, 19.21, 19.59, 19.82]
MCS_LVX = [24.89, 25.39, 25.97, 26.47, 25.50, 26.71, 27.09, 27.41, 26.84, 27.43, 27.99]
tree_kernel_scores_rand = [45.52, 46.45, 47.42, 48.04, 47.13, 49.19, 49.27, 50.52, 49.78, 50.63, 51.26]
tree_kernel_scores_const = [46.96, 47.87, 48.87, 49.58, 48.58, 50.55, 50.68, 52.08, 51.37, 52.17, 52.87]
tree_kernel_scores_lvx = [47.70, 48.58, 49.79, 50.17, 49.08, 51.21, 51.47, 52.91, 52.01, 52.81, 53.04]


# 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, 33)
# 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/cifar100_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(15, 29)
# 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/cifar100_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(43, 54)
# 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/cifar100_tk.pdf')