import matplotlib.pyplot as plt
import numpy as np

# Define the data
model_names = ['VGG13', 'VGG16', 'VGG19', 'ResNet18', 'ResNet34', 'ResNet50', 'DenseNet121', 'DenseNet161', 'DenseNet169', 'MobileNet_v2', 'GoogLeNet', 'Inception_v3']
TED_rand = [9.21, 9.23, 9.15, 9.21, 9.21, 9.21, 9.19, 9.20, 9.21, 9.20, 9.23, 9.20]
TED_const = [32.97, 32.89, 32.85, 32.90, 32.92, 32.89, 32.89, 32.88, 32.88, 32.89, 32.96, 32.89]
TED_LVX = [8.21, 8.14, 8.15, 8.52, 8.16, 8.44, 8.20, 8.19, 8.18, 8.41, 8.41, 8.39]
MCS_rand = [28.98, 29.11, 30.67, 28.95, 28.95, 28.96, 29.09, 29.07, 29.25, 29.24, 28.66, 29.19]
MCS_const = [18.77, 19.11, 19.10, 18.92, 18.98, 19.00, 19.11, 19.11, 19.08, 19.09, 18.86, 19.03]
MCS_LVX = [32.31, 32.55, 31.78, 30.24, 32.07, 31.09, 32.13, 32.12, 32.13, 31.61, 30.75, 31.02]
tree_kernel_scores_rand = [59.49, 59.34, 59.44, 58.94, 59.06, 59.16,
                           59.53, 59.35, 59.52, 59.53, 58.62, 59.37]
tree_kernel_scores_const = [58.41, 59.55, 59.39, 58.87, 59.01, 59.21,
                            59.48, 59.48, 59.46, 59.38, 58.71, 59.27]
tree_kernel_scores_lvx = [63.43, 63.79, 63.32, 61.19, 63.27, 62.06,
                          63.44, 63.73, 63.46, 61.87, 61.38, 61.85]

# 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

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, 5)
# 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, 36)
# 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/cifar10_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, 33)
# 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.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/cifar10_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(55, 64)
# 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/cifar10_tk.pdf')