import numpy as np
import matplotlib.pyplot as plt
from matplotlib import style
style.use('ggplot')

# colors for our groups 
colors = 5*['r', 'b', 'g', 'm', 'y', 'c']

class K_Means:
    def __init__(self, k=2, tol=0.01, max_iter=300):
        self.k = k # num groups to cluster
        self.tol = tol # tolerance: how much centroid can  move before converging
        self.max_iter = max_iter # if tolerance is not reached (did not converge)
    def fit(self, data):
        self.centroids = {}
        # first k centroids are chosen arbitraraly
        # in this case the first k features are the initial centroids
        for i in range(self.k):
            self.centroids[i] = data[i]

        for i in range(self.max_iter):
            # contains (centroid, fetures) pairs
            self.classifications = {}

            for i in range(self.k):
                self.classifications[i] = []

            for featureset in data:
                # list of distances from each centroid from the feature
                distances = [np.linalg.norm(featureset - self.centroids[centroid]) for centroid in self.centroids]
                classification = distances.index(min(distances))
                self.classifications[classification].append(featureset)
            prev_centroids = dict(self.centroids)

            for classification in self.classifications:
                # centroid changes to mean of features
                self.centroids[classification] = np.average(self.classifications[classification], axis=0)
            # see if converged
            converged = True
            for c in self.centroids:
                origanal_centroid = prev_centroids[c]
                current_centroid = self.centroids[c]
                if np.sum((current_centroid-origanal_centroid) / origanal_centroid*100.0) > self.tol:
                    print('centroid moved: ', np.sum((current_centroid-origanal_centroid) / origanal_centroid*100.0))
                    converged = False
    def predict(self, data):
        distances = [np.linalg.norm(data - self.centroids[centroid]) for centroid in self.centroids]
        classification = distances.index(min(distances))
        return classification