import matplotlib.pyplot as plt
import numpy as np

class BlackBoxAlgo:
    def __init__(self, X, y): 
        print('Training.')
        
        n_class = 3
        n_feat = 2
        self.Mean = np.zeros((n_class, n_feat))
        self.Std = np.zeros((n_class, n_feat))

        for i in range(0, n_class):
            for j in range(0, n_feat):
                X_sub = np.extract(y == i, X[:, j])
                self.Mean[i, j] = np.mean(X_sub)
                self.Std[i, j] = np.std(X_sub)


    def validate(self, X):
        print('Validation.')


        n = X.shape[0]
        n_feat = 2
        n_class = 3
        y_dist = np.zeros((n, n_class))
        y_prob = np.zeros((n, n_class))
        y = np.zeros(n, dtype=np.uint)

        dist = np.zeros((n, n_feat, n_class))

        for i in range(0, n):
            for j in range(0, n_class):
                y_p = 0; #np.zeros(1)
                for k in range(0, n_feat):
                    dist[i, k, j] = X[i, k] - self.Mean[j, k]

                    y_p += np.square(X[i, k] - self.Mean[j, k])
                y_dist[i, j] = np.sqrt(y_p)
                y_prob[i, j] = getGaussianProbab(0, self.Std[j, k], y_dist[i, j])
                
                #if y_dist[i, j] < y_dist[i, y[i]]:
                #    y[i] = j
                if y_prob[i, j] > y_prob[i, y[i]]:
                    y[i] = j
        return y

def getGaussianProbab(mu, sigma, x):
    var = np.square(sigma)
    p = 1 / (2 * np.pi * var) * np.exp(-1 * np.square(x - mu) / (2 * var))
    return p