import pandas as pd
from scipy.cluster.hierarchy import linkage, fcluster
import matplotlib.pyplot as plt
import seaborn as sns

# Кластеризация: Группировка стран на основе их характеристик:
# количество выбросов от добычи нефтепродуктов, газа, угля в 2020 году.

# Загружаем данные, отбираем данные за 2020 год
# и удаляем строки, в которых данные отсутствуют
data_first = pd.read_csv('CO2.csv')
data_first = data_first.dropna()
data = data_first[data_first['Year'] == 2020]
data = data[data.Country != 'Global']
names = data['Country']

# Выделение признаков
features = data[['Coal', 'Oil', 'Gas']]

# Применение алгоритма linkage
link_cl = linkage(features, method='ward', metric='euclidean')

# Кластеризация на 3 кластера
clusters = fcluster(link_cl, 4, criterion='maxclust')

# Добавление информации о кластерах в исходные данные
features['cluster'] = clusters

# Диаграмма результатов кластеризации
plt.figure(figsize=(12, 6))
sns.scatterplot(x='Coal', y='Gas', size='Oil',
                sizes=(10, 200), hue='cluster', palette='viridis', data=features)
plt.title('Clustering Results')
plt.xlabel('Coal')
plt.ylabel('Gas')

# Вывод таблицы с результатами кластеризации
# Добавление наименований стран
features.insert(loc=0, column='country', value=names)
# Вывод результатов кластеризации
result_clust = features[['cluster', 'country', 'Coal', 'Oil', 'Gas']]
print(result_clust)

plt.show()