#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Sep 13 05:22:49 2020

@author: nishantn
"""

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

from sklearn.cluster import KMeans

def KMeans(X, n=6):
    kmeans = KMeans(n_clusters = n, init = 'k-means++')
    y_kmeans = kmeans.fit_predict(X)
    
    # plt.scatter(X[y_kmeans == 0, 0], X[y_kmeans == 0, 1], s = 100, c = 'red', label = 'Cluster 1')
    # plt.scatter(X[y_kmeans == 1, 0], X[y_kmeans == 1, 1], s = 100, c = 'blue', label = 'Cluster 2')
    # plt.scatter(X[y_kmeans == 2, 0], X[y_kmeans == 2, 1], s = 100, c = 'green', label = 'Cluster 3')
    # plt.scatter(X[y_kmeans == 3, 0], X[y_kmeans == 3, 1], s = 100, c = 'cyan', label = 'Cluster 4')
    # plt.scatter(X[y_kmeans == 4, 0], X[y_kmeans == 4, 1], s = 100, c = 'magenta', label = 'Cluster 5')
    # plt.scatter(kmeans.cluster_centers_[:, 0], kmeans.cluster_centers_[:, 1], s = 300, c = 'yellow', label = 'Centroids')
    # plt.title('Clusters of customers')
    # plt.xlabel('Annual Income (k$)')
    # plt.ylabel('Spending Score (1-100)')
    # plt.legend()
    # plt.show()
    return y_kmeans, kmeans

from sklearn.cluster import AgglomerativeClustering

def HC(X, n = 6):
    hc = AgglomerativeClustering(n_clusters = n, affinity = 'euclidean', linkage = 'ward')
    y_hc = hc.fit_predict(X)
    
    # plt.scatter(X[y_hc == 0, 0], X[y_hc == 0, 1], s = 100, c = 'red', label = 'Cluster 1')
    # plt.scatter(X[y_hc == 1, 0], X[y_hc == 1, 1], s = 100, c = 'blue', label = 'Cluster 2')
    # plt.scatter(X[y_hc == 2, 0], X[y_hc == 2, 1], s = 100, c = 'green', label = 'Cluster 3')
    # plt.scatter(X[y_hc == 3, 0], X[y_hc == 3, 1], s = 100, c = 'cyan', label = 'Cluster 4')
    # plt.scatter(X[y_hc == 4, 0], X[y_hc == 4, 1], s = 100, c = 'magenta', label = 'Cluster 5')
    # plt.title('Clusters of customers')
    # plt.xlabel('Annual Income (k$)')
    # plt.ylabel('Spending Score (1-100)')
    # plt.legend()
    # plt.show()
    
    return y_hc, hc
