import re
import matplotlib.pyplot as plt

file = './logs/cifar10-c/conceptEM/numcluster-conceptem-adam-adapt.txt'

with open(file, 'r') as f:
    lines = f.readlines()
    lines = [x.strip() for x in lines]
    lines = [x.split(', ') for x in lines]
    lines = [[re.sub(r'[^\d\.]', '', y) for y in x] for x in lines]
    lines = [[float(y) for y in x] for x in lines]

weights = []
num_clusters = 6
current_cluster_num = 6
current_tobe_removed = 0
removed_clusters = []

for x in lines:
    if len(x) < current_cluster_num:
        removed_clusters.append(current_tobe_removed)
        current_cluster_num = len(x)
    min_value = min(x)
    new_weights_vector = []
    K = 0
    for i in range(num_clusters):
        if i in removed_clusters:
            new_weights_vector.append(0)
        else:
            if x[K] == min_value:
                current_tobe_removed = K
            new_weights_vector.append(x[K])
            K += 1
    weights.append(new_weights_vector)

cluster_weights_0 = [weight[0] for weight in weights]
cluster_weights_1 = [weight[1] for weight in weights]
cluster_weights_2 = [weight[2] for weight in weights]
cluster_weights_3 = [weight[3] for weight in weights]
cluster_weights_4 = [weight[4] for weight in weights]
cluster_weights_5 = [weight[5] for weight in weights]
x = [i for i in range(len(cluster_weights_0))]


plt.plot(x, cluster_weights_0, label='Model 1')
plt.plot(x, cluster_weights_1, label='Model 2')
plt.plot(x, cluster_weights_2, label='Model 3')
plt.plot(x, cluster_weights_3, label='Model 4')
plt.plot(x, cluster_weights_4, label='Model 5')
plt.plot(x, cluster_weights_5, label='Model 6')

plt.legend()
plt.xlabel('Communication Rounds')
plt.ylabel('Weights')
plt.title('Adaptive process of ConceptEM_Adam')
plt.savefig('adaptive-conceptem-adam.pdf')