import numpy as np
import matplotlib.pyplot as plt

class1_points = np.array([[1.9, 1.2],
                          [1.5, 2.1],
                          [1.9, 0.5],
                          [1.5, 0.9],
                          [0.9, 1.2],
                          [1.1, 1.7],
                          [1.4, 1.1]])

class2_points = np.array([[3.2, 3.2],
                          [3.7, 2.9],
                          [3.2, 2.6],
                          [1.7, 3.3],
                          [3.4, 2.6],
                          [4.1, 2.3],
                          [3.0, 2.9]])

class3_points = np.array([[3.3, 1.2],
                          [3.8, 0.9],
                          [3.3, 0.6],
                          [2.8, 1.3],
                          [3.5, 0.6],
                          [4.2, 0.3],
                          [3.1, 0.9]])

X=np.concatenate((class1_points,class2_points,class3_points),axis=0)

Y=np.concatenate((np.zeros(len(class1_points)),np.ones(len(class1_points)),np.ones(len(class1_points))+1),axis=0)

print(Y)

prior_prob=[np.sum(Y==0)/len(Y),np.sum(Y==1)/len(Y),np.sum(Y==2)/len(Y)]

class_u=[np.mean(X[Y==0],axis=0),np.mean(X[Y==1],axis=0),np.mean(X[Y==2],axis=0)]

class_cov=[np.cov(X[Y==0],rowvar=False),np.cov(X[Y==1],rowvar=False),np.cov(X[Y==2],rowvar=False)]

def pdf(x, mean, cov):
    n = len(mean)
    coff = 1 / (2 * np.pi) ** (n / 2) * np.sqrt(np.linalg.det(cov))
    exponent = np.exp(-(1 / 2) * np.dot(np.dot((x - mean).T, np.linalg.inv(cov)), (x - mean)))
    return coff * exponent

xx, yy = np.meshgrid(np.arange(0, 5, 0.05), np.arange(0, 4, 0.05))

grid_points = np.c_[xx.ravel(), yy.ravel()]

grid_label = []

for point in grid_points:
    poster_prob = []
    for i in range(3):
        likelihood = pdf(point, class_u[i], class_cov[i])
        poster_prob.append(prior_prob[i] * likelihood)
    pre_class = np.argmax(poster_prob)
    grid_label.append(pre_class)

grid_label = np.array(grid_label)

pre_grid_label = grid_label.reshape(xx.shape)

plt.scatter(class1_points[:,0],class1_points[:,1],c="blue",label="class 1")
plt.scatter(class2_points[:,0],class2_points[:,1],c="red",label="class 2")
plt.scatter(class3_points[:,0],class3_points[:,1],c="yellow",label="class 3")

plt.legend()

contour=plt.contour(xx,yy,pre_grid_label,colors='green')

plt.show()