import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
import pandas as pd
import numpy as np
from pandas import DataFrame
from sklearn.decomposition import PCA

df = pd.read_csv('DowJones.csv')

df = df.drop('stock', axis=1)
df = df.drop('date', axis=1)
df = df.dropna()

pca = PCA(n_components=2).fit(df)
data_2d = pca.transform(df)

kmeans = KMeans(n_clusters=3).fit(data_2d)
centroids = kmeans.cluster_centers_

for i in range(0, data_2d.shape[0]):
  if kmeans.labels_[i] == 0:
    plt.scatter(data_2d[i,0], data_2d[i,1], c='green') 
  elif kmeans.labels_[i] == 1:
    plt.scatter(data_2d[i,0], data_2d[i,1], c='yellow') 
  elif kmeans.labels_[i] == 2:
    plt.scatter(data_2d[i,0], data_2d[i,1], c='blue') 

plt.scatter(centroids[:, 0], centroids[:, 1], c='black', marker='X')
plt.title("Clusters")
plt.show()

