import pandas as pd
from sklearn.preprocessing import LabelEncoder
from sklearn.manifold import TSNE
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
import seaborn as sns

# Загрузка данных
data = pd.read_csv("ds_salaries.csv")

# Преобразование категориальных переменных в числовые
le = LabelEncoder()
data['company_location'] = le.fit_transform(data['company_location'])
data['company_size'] = le.fit_transform(data['company_size'])

# Выбор нужных признаков
features = ['company_location', 'company_size']

# Применение t-SNE для уменьшения размерности
tsne = TSNE(n_components=2, random_state=42)
data_tsne = tsne.fit_transform(data[features])

# Кластеризация данных
kmeans = KMeans(n_clusters=3, random_state=42)
labels = kmeans.fit_predict(data_tsne)

# Создание датафрейма с новыми координатами и метками кластеров
data_tsne_df = pd.DataFrame(data=data_tsne, columns=['Dimension 1', 'Dimension 2'])
data_tsne_df['Cluster'] = labels

# Добавление номера кластера в исходные данные
data['Cluster'] = labels

# Визуализация
plt.figure(figsize=(10, 6))
sns.scatterplot(x='Dimension 1', y='Dimension 2', hue='Cluster', data=data_tsne_df, palette='viridis', s=50)
plt.title('t-SNE Clustering of Companies')
plt.show()

# Вывод номера кластера в исходных данных
print(data[['company_location', 'company_size', 'Cluster']])