import numpy as np
from scipy.stats import mode
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

class KNN:
    def __init__(self, n_neighbors):
        self.n_neighbors = n_neighbors
        self.model_trained = False
        
    '''
    Ucitavanje trening podataka
    * X - matrica ulaznih podataka
    * y - kolona oznaka klasa
    '''
    def fit(self, X, y):
        if isinstance(X, pd.DataFrame):
            self.X = data.values
        else:
            self.X = X
            
        if isinstance(y, pd.DataFrame):
            self.y = data.values
        else:
            self.y = y
            
        if self.X.shape[0] != self.y.shape[0]:
            raise Exception('X and y have different shapes')
            self.X = None
            self.y = None
            
        self.model_trained = True
        
    '''
    Priprema matrica podataka
    * data - matrica ulaznih podataka (numpy matrica/pandas dataframe/2D niz/...)
    '''
    def _prepare_data_matrix(self, data):
        data_matrix = None
        # Provera formata ulaznih podataka
        if isinstance(data, pd.DataFrame):
            data_matrix = data.values
        else:
            data_matrix = np.array(data)
            
        if len(data_matrix.shape) < 2:
            raise Exception('Invalid data shape')
            
        # Konacni oblik pripremljenih podataka za obradu je numpy matrica
        return data_matrix
        
    '''
    Izracunavanje udaljenosti izmedju dva vektora
    * u - vektor
    * v - vektor
    * metric - odabrana metrika (podrazumevano je euklidsko rastojanje)
    '''
    def distance(self, u, v,  metric='euclidean'):
        if metric == 'euclidean':
            return np.linalg.norm(u-v)
        else:
            raise Exception('Unknown metric')
        
    '''
    Izracunavanje matrice udaljenosti izmedju ulaznih i trening podataka
    * data_matrix - matrica ulaznih podataka
    '''
    def _compute_distance_matrix(self, data_matrix):
        X = self.X
        return np.array(
        [
            [self.distance(train_vector, prediction_vector) for train_vector in X] 
                for prediction_vector in data_matrix
        ])
        
    '''
    Izracunavanje matrice konfuzije nad pripremljenim podacima
    * true_labels - stvarne oznake klasa
    * predicted_labels - predvidjene oznake klasa
    '''
    def _compute_confusion_matrix(self, true_labels, predicted_labels):
        num_classes = np.max(true_labels) + 1
        num_labels = true_labels.shape[0]
        confusion_matrix = np.zeros((num_classes, num_classes))
        
        for i in range(num_labels):
            true_label = true_labels[i]
            predicted_label = predicted_labels[i]
            
            confusion_matrix[true_label, predicted_label] += 1
        
        return confusion_matrix
        
    '''
    Priprema podataka i izracunavanje matrice konfuzije
    * true_labels - stvarne oznake klasa
    * predicted_labels - predvidjene oznake klasa
    '''
    def get_confusion_matrix(self, true_labels, predicted_labels):
        if not self.model_trained:
            raise Exception('No training data provided')
                
        true_labels_matrix = None
        predicted_labels_matrix = None
        
        if isinstance(true_labels, pd.DataFrame):
            true_labels_matrix = true_labels.values.ravel()
        else:
            true_labels_matrix = np.array(true_labels).ravel()
            
        if isinstance(predicted_labels, pd.DataFrame):
            predicted_labels_matrix = predicted_labels.values.ravel()
        else:
            predicted_labels_matrix = np.array(predicted_labels).ravel()
        
        return self._compute_confusion_matrix(true_labels_matrix, predicted_labels_matrix)
      
    '''
    Izracunavanje tacnosti klasifikacije (accuracy)
    nad pripremljenim podacima
    * true_labels - stvarne oznake klasa
    * predicted_labels - predvidjene oznake klasa
    '''
    def _compute_accuracy(self, true_labels, predicted_labels):
        num_correct_labels = np.sum(
            [1 for i in range(true_labels.shape[0]) if true_labels[i] == predicted_labels[i]]
        ) 
        return num_correct_labels / true_labels.shape[0]

    '''
    Generisanje izvestaja o kvalitetu klasifikacije
    * true_labels - stvarne oznake klasa
    * predicted_labels - predvidjene oznake klasa
    '''
    def generate_report(self, true_labels, predicted_labels):
        if not self.model_trained:
            raise Exception('No training data provided')
                
        true_labels_matrix = None
        predicted_labels_matrix = None
        
        if isinstance(true_labels, pd.DataFrame):
            true_labels_matrix = true_labels.values.ravel()
        else:
            true_labels_matrix = np.array(true_labels).ravel()
            
        if isinstance(predicted_labels, pd.DataFrame):
            predicted_labels_matrix = predicted_labels.values.ravel()
        else:
            predicted_labels_matrix = np.array(predicted_labels).ravel()
        
        accuracy = self._compute_accuracy(true_labels_matrix, predicted_labels_matrix)
        
        return {
            'accuracy': accuracy
        }
    '''
    Predvidjanje oznaka klaza nepoznatih poznataka
    * data - matrica vrednosti (numpy matrica/pandas dataframe/2D niz/...)
    * metric - metrika za izracunavanje udaljenosti
    * y_true - stvarne oznake klasa ulaznih podataka
    * include_confusion_matrix - izracunavanje matrice konfuzije (samo ukoliko je y_true prosledjeno)
    '''
    def predict(self, data, metric='euclidean', y_true=None, include_confusion_matrix=False):
        # Provera da li su ucitani trening podaci
        if not self.model_trained:
            raise Exception('No training data provided')
        
        X = self.X
        y = self.y
        
        data_matrix = self._prepare_data_matrix(data)
            
        '''
        Racunanje matrice udaljenosti svakog ulaznog podatka u odnosu na svaki
        podatak iz trening skupa
        
        Matrica distance_matrix je oblika: 
        - red (i): ulazni podatak / prediction_vector 
        - kolona (j): trening podatak / train_vector,
        - polje(i, j) rastojanje 
                      izmedju ulaznog podatka (i) 
                      i trening podatka (j)
        '''
        
        distance_matrix = self._compute_distance_matrix(data_matrix)

        '''
         Sortiranje po kolonama matrice (axis = 1)
         - Sortiranje se vrsi u rastucem poretku
         - Za svaki red (ulazni  podatak) sortiraju se 
           vrednosti rastojanja izmedju ulaznog podatka i trening podataka
         - metod argsort ne vrsi konacno sortiranje vec vraca redosled KLJUCEVA (indeksa)
           kolona polazne matrice u redosledu u kome bi se nasli u sortiranoj matrici
           
           Primer:
           distance_matrix = [#  0  1   2  -- indeksi kolona
                               [10, 4,  1],
                               [ 3, 2, 14]
                             ]
                             
           sorted_keys = argsort(distance_matrix) -> [
                                         [2, 1, 0],
                                         [1, 0, 2]
                                       ]
        '''
        
        sorted_keys = distance_matrix.argsort(axis = 1)
        
        '''
        "Primenom" matrice kljuceva na matricu istog oblika dobijaju se vrednosti
        matrice u redosledu sortiranih kljuceva.
        
        Primer: 
            distance_matrix[sorted_keys] -> [
                                              [1, 4, 10],
                                              [2, 3, 14]
                                            ]
        '''
        
        sorted_classes = y[sorted_keys[:, :self.n_neighbors]]
        
        class_labels = np.array([mode(row).mode for row in sorted_classes]).ravel()
        return class_labels
        
    '''
    Pronalazenje k najblizih suseda zadatih podataka
    * data - matrica vrednosti ulaznih podataka
    * metric - metrika za racunanje udaljenosti
    * return_indices - True: povratna vrednost ce biti indeksi najblizih suseda, 
                       False: povratna vrednost ce biti vrednosti najblizih suseda
    '''
    def k_neighbors(self, data, metric='euclidean', return_indices=False):
        data_matrix = self._prepare_data_matrix(data)
        distance_matrix = self._compute_distance_matrix(data_matrix)
        sorted_keys = distance_matrix.argsort(axis = 1)
        
        if return_indices:
            return sorted_keys[:,:self.n_neighbors]
        else:
            return self.X[sorted_keys][:,:self.n_neighbors,:]

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)

k = 3
knn = KNN(n_neighbors=k)

knn.fit(X_train, y_train)
predicted_labels = knn.predict(X_test)
confusion_matrix = knn.get_confusion_matrix(y_test, predicted_labels)
report = knn.generate_report(y_test, predicted_labels)

print(f'K = {k}')
print('Oznaka klase: ', predicted_labels)
print('Matrica konfuzije:\n', confusion_matrix)
print('K najblizih suseda:', knn.k_neighbors(X_test))
print('K najblizih suseda (indeksi):', knn.k_neighbors(X_test, return_indices=True))
print('Tacnost (accuracy): ',report['accuracy'])

K = 3
Oznaka klase:  [0 0 2 2 0 1 1 1 0 1 0 0 2 1 1 1 2 2 0 0 1 2 0 0 0 2 0 0 1 1 2 2 0 0 2 2 2
 1 1 0 2 0 1 0 0]
Matrica konfuzije:
 [[19.  0.  0.]
 [ 0. 13.  1.]
 [ 0.  0. 12.]]
K najblizih suseda: [[[4.4 3.  1.3 0.2]
  [4.4 3.2 1.3 0.2]
  [4.6 3.2 1.4 0.2]]

 [[5.2 3.5 1.5 0.2]
  [5.5 3.5 1.3 0.2]
  [5.1 3.5 1.4 0.3]]

 [[7.2 3.  5.8 1.6]
  [7.1 3.  5.9 2.1]
  [7.3 2.9 6.3 1.8]]

 [[5.8 2.7 5.1 1.9]
  [5.7 2.5 5.  2. ]
  [5.6 2.8 4.9 2. ]]

 [[4.6 3.2 1.4 0.2]
  [4.7 3.2 1.6 0.2]
  [4.8 3.1 1.6 0.2]]

 [[5.7 2.9 4.2 1.3]
  [6.  2.9 4.5 1.5]
  [5.7 2.8 4.1 1.3]]

 [[5.7 2.9 4.2 1.3]
  [5.6 2.7 4.2 1.3]
  [6.  2.9 4.5 1.5]]

 [[6.1 3.  4.6 1.4]
  [6.  2.9 4.5 1.5]
  [6.1 2.9 4.7 1.4]]

 [[5.2 3.5 1.5 0.2]
  [5.4 3.7 1.5 0.2]
  [5.5 3.5 1.3 0.2]]

 [[5.  2.3 3.3 1. ]
  [5.  2.  3.5 1. ]
  [5.5 2.4 3.7 1. ]]

 [[5.5 4.2 1.4 0.2]
  [5.4 3.9 1.3 0.4]
  [5.7 3.8 1.7 0.3]]

 [[5.1 3.8 1.6 0.2]
  [5.1 3.5 1.4 0.3]
  [5.  3.6 1.4 0.2]]

 [[6.4 2.8 5.6 2.1]
  [6.3 2.9 5.6 1.8]
  [6.5 3.  5.5 1.8]]

 [[5.7 2.9 4.2 1.3]
  [5.6 3.  4.1 1.3]
  [5.7 2.8 4.1 1.3]]

 [[6.6 3.  4.4 1.4]
  [6.6 2.9 4.6 1.3]
  [6.5 2.8 4.6 1.5]]

 [[5.8 2.6 4.  1.2]
  [5.8 2.7 3.9 1.2]
  [5.7 2.8 4.1 1.3]]

 [[6.8 3.2 5.9 2.3]
  [6.7 3.1 5.6 2.4]
  [6.8 3.  5.5 2.1]]

 [[6.  2.7 5.1 1.6]
  [6.3 2.7 4.9 1.8]
  [6.2 2.8 4.8 1.8]]

 [[5.5 4.2 1.4 0.2]
  [5.4 3.9 1.7 0.4]
  [5.4 3.9 1.3 0.4]]

 [[4.4 3.  1.3 0.2]
  [4.9 3.  1.4 0.2]
  [4.8 3.1 1.6 0.2]]

 [[6.3 3.3 4.7 1.6]
  [6.6 3.  4.4 1.4]
  [6.7 3.1 4.4 1.4]]

 [[6.3 2.7 4.9 1.8]
  [6.  2.7 5.1 1.6]
  [6.  2.2 5.  1.5]]

 [[5.4 3.7 1.5 0.2]
  [5.2 3.5 1.5 0.2]
  [5.1 3.8 1.6 0.2]]

 [[5.  3.4 1.5 0.2]
  [5.1 3.5 1.4 0.2]
  [5.2 3.5 1.5 0.2]]

 [[4.4 3.  1.3 0.2]
  [4.4 3.2 1.3 0.2]
  [4.6 3.2 1.4 0.2]]

 [[7.3 2.9 6.3 1.8]
  [7.1 3.  5.9 2.1]
  [7.2 3.  5.8 1.6]]

 [[4.9 3.  1.4 0.2]
  [4.9 3.1 1.5 0.1]
  [4.9 3.1 1.5 0.2]]

 [[4.7 3.2 1.6 0.2]
  [5.  3.4 1.5 0.2]
  [5.  3.4 1.6 0.4]]

 [[5.8 2.7 3.9 1.2]
  [5.8 2.6 4.  1.2]
  [5.7 2.8 4.1 1.3]]

 [[7.  3.2 4.7 1.4]
  [6.7 3.1 4.7 1.5]
  [6.7 3.  5.  1.7]]

 [[7.6 3.  6.6 2.1]
  [7.1 3.  5.9 2.1]
  [7.3 2.9 6.3 1.8]]

 [[6.2 3.4 5.4 2.3]
  [6.4 3.2 5.3 2.3]
  [6.3 3.3 6.  2.5]]

 [[5.1 3.8 1.6 0.2]
  [5.  3.6 1.4 0.2]
  [5.1 3.5 1.4 0.3]]

 [[4.9 3.  1.4 0.2]
  [4.9 3.1 1.5 0.2]
  [4.8 3.1 1.6 0.2]]

 [[6.7 3.1 5.6 2.4]
  [6.8 3.2 5.9 2.3]
  [6.5 3.  5.8 2.2]]

 [[6.8 3.2 5.9 2.3]
  [6.8 3.  5.5 2.1]
  [6.7 3.1 5.6 2.4]]

 [[6.3 2.7 4.9 1.8]
  [6.4 2.7 5.3 1.9]
  [6.2 2.8 4.8 1.8]]

 [[5.6 2.7 4.2 1.3]
  [5.7 2.9 4.2 1.3]
  [5.5 2.5 4.  1.3]]

 [[5.4 3.  4.5 1.5]
  [5.7 2.9 4.2 1.3]
  [6.  2.9 4.5 1.5]]

 [[5.2 3.5 1.5 0.2]
  [5.1 3.5 1.4 0.2]
  [5.1 3.5 1.4 0.3]]

 [[6.8 3.2 5.9 2.3]
  [7.1 3.  5.9 2.1]
  [7.7 3.8 6.7 2.2]]

 [[4.7 3.2 1.6 0.2]
  [5.  3.4 1.6 0.4]
  [4.8 3.1 1.6 0.2]]

 [[5.  2.3 3.3 1. ]
  [5.  2.  3.5 1. ]
  [5.7 2.6 3.5 1. ]]

 [[5.  3.4 1.5 0.2]
  [5.  3.2 1.2 0.2]
  [5.1 3.5 1.4 0.2]]

 [[4.6 3.2 1.4 0.2]
  [4.7 3.2 1.3 0.2]
  [4.4 3.2 1.3 0.2]]]
K najblizih suseda (indeksi): [[ 53  38  78]
 [ 99  59  35]
 [ 63  21  13]
 [ 25  24  50]
 [ 78  45  46]
 [  3  15  81]
 [  3  75  15]
 [ 23  15   7]
 [ 99  55  59]
 [ 49   6  87]
 [ 40  27  12]
 [ 11  35  26]
 [ 64  65  77]
 [  3  85  81]
 [ 43  57 103]
 [ 90   9  81]
 [ 73  10 102]
 [ 60  67  19]
 [ 40  41  27]
 [ 53  54  46]
 [ 36  43  80]
 [ 67  60  66]
 [ 55  99  11]
 [ 44  29  99]
 [ 53  38  78]
 [ 13  21  63]
 [ 54  93  97]
 [ 45  44  14]
 [  9  90  81]
 [104 100  39]
 [ 83  21  13]
 [ 82   0  91]
 [ 11  26  35]
 [ 54  97  46]
 [ 10  73  72]
 [ 73 102  10]
 [ 67  47  19]
 [ 75   3 101]
 [ 74   3  15]
 [ 99  29  35]
 [ 73  21  51]
 [ 45  14  46]
 [ 49   6  68]
 [ 44  32  29]
 [ 78  18  38]]
Tacnost (accuracy):  0.9777777777777777