import pandas as pd
import numpy as np
from scipy.stats import mode
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from random import sample
import matplotlib.pyplot as plt

X, y = load_iris(return_X_y=True)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3)

class KMeans:
    def __init__(self, n_clusters, num_inits=10, max_iters=10):
        self.n_clusters = n_clusters
        self.model_fitted = False
        self.centroids = None
        self.num_inits = num_inits
        self.max_iters = max_iters
        
    # Priprema matrice podataka
    def _prepare_data_matrix(self, data):
        if isinstance(data, pd.DataFrame):
            return data.values
        return np.array(data)
        
    '''
    Inicijalizacija centroida
    na tačke iz trenung skupa odabrae slučajnim izborom
    '''
    def _initialize_centroids(self):
        n_clusters = self.n_clusters
        X = self.X.copy()
        np.random.shuffle(X)
        return X[:n_clusters]
    
    '''
    Mera udaljenosti
    '''
    def distance(self, u, v, metric='euclidean'):
        if metric == 'euclidean':
            return np.linalg.norm(u - v)
        elif metric == 'manhattan':
            return np.sum(u - v)
        else:
            raise Exception('Unknown metric')
    
    '''
    Izracunavanje matrice rastojanja
    '''
    def _compute_distance_matrix(self, data_matrix, metric='euclidean'):
        X = self.X
        return np.array([
            [self.distance(training_vector, input_vector, metric) for training_vector in X]
            for input_vector in data_matrix
        ])
        
    '''
    Pronalaženje klastera algoritmom K-sredina
    '''
    def fit(self, X, y):
        # Provera ulaznih podataka
        self.X = self._prepare_data_matrix(X)
        self.y = self._prepare_data_matrix(y).ravel()
        
        if X.shape[0] != y.shape[0]:
            raise Exception('X and y have different number of records')
        
        # Najkvalitetnije klasterovanje (označavanje sa najmanjom MSE)
        best_result = None
        
        # Centroidi najboljeg klasterovanja
        best_centroids = None
        
        # Srednje kvadratna greška najboljeg rezultata
        best_result_error = float('inf')
        
        num_inits = self.num_inits
        max_iters = self.max_iters
        
        # Za svaku inicijalizaciju klasterovanja:
        for _ in range(num_inits):

            # Inicijalizacija centroida
            centroids = self._initialize_centroids()
            
            # Za svaku iteraciju jedne inicijalizacije klasterovanja:
            for _ in range(max_iters):
                # Odrediti najbliže tačke svakog centroida
                '''
                [
                    [dist(x0, d0), dist(x0, d1) ...]
                    [dist(x1, d0), dist(x1, d1) ...]
                ]
                '''
                
                distance_matrix = np.array([
                    [self.distance(x, c) for c in centroids]
                    for x in self.X
                ])
                
                # Dodeljivanje oznaka najbližeg centroida za svaku tačku iz skupa
                cluster_labels = distance_matrix.argmin(axis=1)
        
        
                # Inicijalizacije srednje kvadratne greške
                mean_square_error = 0
        
                # Azurirati pozicije centroida
                for ci in range(centroids.shape[0]):
                    centroid_index_filter = cluster_labels == ci
                    closest_points = X[centroid_index_filter,:]
    
                    # Računanje srednje kvadratne greške - rastojanja tačaka od najbližeg centroida
                    mean_square_error += np.array([
                        self.distance(centroids[ci], point) ** 2 for point in closest_points
                    ]).sum().ravel() / closest_points.shape
        
                    # Ažuriranje koordinata centroida na srednju vrednost koordinata najbližih tačaka
                    centroids[ci,:] = closest_points.mean(axis=0)

                # Evaluacija klasterovanja (odabir najboljeg rezultata)
                if mean_square_error[0] < best_result_error:
                    best_result_error = mean_square_error[0]
                    best_result = cluster_labels.ravel()
                    best_centroids = centroids
        
        self.centroids = best_centroids
        self.model_fitted = True
        return best_result

# Očekivani broj klastera
n_clusters = 3

# Broj ponovnih inicijalizacija
n_init = 15

# Maksimalni broj iteracija jedne inicijalizacije
n_iter = 20

cls = KMeans(n_clusters, n_init, n_iter)
labels = cls.fit(X, y)

# Iscrtavanje pozicija tačaka u 2D prostoru
# na osnovu 2. i 3. koordinate
plt.scatter(X[:,2], X[:,3], c=labels)

# Iscrtavanje pozicija centroida
for centroid in cls.centroids:
    plt.scatter(centroid[2], centroid[3], c='red')