import numpy as np
from kmeans import KMeans
from numpy.linalg import LinAlgError
from tqdm import tqdm

SIGMA_CONST = 1e-06
LOG_CONST = 1e-32
FULL_MATRIX = False


class GMM(object):

def __init__(self, X, K, max_iters=100):
"""
Args:
X: the observations/datapoints, N x D numpy array
K: number of clusters/components
max_iters: maximum number of iterations (used in EM implementation)
"""
self.points = X
self.max_iters = max_iters
self.N = self.points.shape[0]
self.D = self.points.shape[1]
self.K = K

def softmax(self, logit):
"""
Args:
logit: N x D numpy array
Return:
prob: N x D numpy array. See the above function.
Hint:
Add keepdims=True in your np.sum() function to avoid broadcast error.
"""
logit_sub = logit - logit.max(axis=1, keepdims=True)
exp_logit = np.exp(logit_sub)
sum_exp_logit = np.sum(exp_logit, axis=1, keepdims=True)
return exp_logit / sum_exp_logit

def logsumexp(self, logit):
"""
Args:
logit: N x D numpy array
Return:
s: N x 1 array where s[i,0] = logsumexp(logit[i,:]). See the above function
Hint:
The keepdims parameter could be handy
"""
logit_max = np.max(logit, axis=1, keepdims=True)
return logit_max + np.log(
np.sum(np.exp(logit - logit_max), axis=1, keepdims=True)
)

def normalPDF(self, points, mu_i, sigma_i):
"""
Args:
points: N x D numpy array
mu_i: (D,) numpy array, the center for the ith gaussian.
sigma_i: DxD numpy array, the covariance matrix of the ith gaussian.
Return:
pdf: (N,) numpy array, the probability density value of N data for the ith gaussian

Hint:
np.diagonal() should be handy.
"""
variances = np.diagonal(sigma_i)
variances = np.maximum(variances, SIGMA_CONST)
log_coeff = -0.5 * np.log(2 * np.pi * variances)
log_exp_part = -0.5 * ((points - mu_i) ** 2) / variances
log_pdf = np.sum(log_coeff + log_exp_part, axis=1)
pdf = np.exp(log_pdf)
return pdf

def multinormalPDF(self, points, mu_i, sigma_i):
"""
Args:
points: N x D numpy array
mu_i: (D,) numpy array, the center for the ith gaussian.
sigma_i: DxD numpy array, the covariance matrix of the ith gaussian.
Return:
normal_pdf: (N,) numpy array, the probability density value of N data for the ith gaussian

Hint:
1. np.linalg.det() and np.linalg.inv() should be handy.
2. Note the value in self.D may be outdated and not correspond to the current dataset.
3. You may wanna check if the matrix is singular before implementing calculation process.
"""
raise NotImplementedError

def create_pi(self):
"""
Initialize the prior probabilities
Args:
Return:
pi: numpy array of length K, prior
"""
return np.ones(self.K) / self.K

def create_mu(self):
"""
Intialize random centers for each gaussian
Args:
Return:
mu: KxD numpy array, the center for each gaussian.
"""
random_indices = np.random.choice(self.N, self.K, replace=False)
mu = self.points[random_indices, :]
return mu

def create_sigma(self):
"""
Initialize the covariance matrix with np.eye() for each k. For grads, you can also initialize the
by K diagonal matrices.
Args:
Return:
sigma: KxDxD numpy array, the diagonal standard deviation of each gaussian.
You will have KxDxD numpy array for full covariance matrix case
"""
return np.repeat(np.eye(self.D)[np.newaxis, :, :], self.K, axis=0)

def _init_components(self, **kwargs):
"""
Args:
kwargs: any other arguments you want
Return:
pi: numpy array of length K, prior
mu: KxD numpy array, the center for each gaussian.
sigma: KxDxD numpy array, the diagonal standard deviation of each gaussian.
You will have KxDxD numpy array for full covariance matrix case

Hint: np.random.seed(5) must be used at the start of this function to ensure consistent outputs.
"""
np.random.seed(5)
pi = self.create_pi()
mu = self.create_mu()
sigma = self.create_sigma()
return pi, mu, sigma

def _ll_joint(self, pi, mu, sigma, full_matrix=FULL_MATRIX, **kwargs):
"""
Args:
pi: np array of length K, the prior of each component
mu: KxD numpy array, the center for each gaussian.
sigma: KxDxD numpy array, the diagonal standard deviation of each gaussian. You will have KxDxD numpy
array for full covariance matrix case
full_matrix: whether we use full covariance matrix in Normal PDF or not. Default is True.

Return:
ll(log-likelihood): NxK array, where ll(i, k) = log pi(k) + log NormalPDF(points_i | mu[k], sigma[k])
"""
N, D = self.points.shape
ll = np.zeros((N, self.K))
for k in range(self.K):
if not full_matrix:
pdf = self.normalPDF(self.points, mu[k], sigma[k])
ll[:, k] = np.log(pi[k] + LOG_CONST) + np.log(pdf + LOG_CONST)
return ll

def _E_step(self, pi, mu, sigma, full_matrix=FULL_MATRIX, **kwargs):
"""
Args:
pi: np array of length K, the prior of each component
mu: KxD numpy array, the center for each gaussian.
sigma: KxDxD numpy array, the diagonal standard deviation of each gaussian.You will have KxDxD numpy
array for full covariance matrix case
full_matrix: whether we use full covariance matrix in Normal PDF or not. Default is True.
Return:
tau: NxK array, the posterior distribution (a.k.a, the soft cluster assignment) for each observation.

Hint:
You should be able to do this with just a few lines of code by using _ll_joint() and softmax() defined above.
"""
ll_joint = self._ll_joint(pi, mu, sigma, **kwargs)
tau = self.softmax(ll_joint)
return tau

def _M_step(self, tau, full_matrix=FULL_MATRIX, **kwargs):
"""
Args:
tau: NxK array, the posterior distribution (a.k.a, the soft cluster assignment) for each observation.
full_matrix: whether we use full covariance matrix in Normal PDF or not. Default is True.
Return:
pi: np array of length K, the prior of each component
mu: KxD numpy array, the center for each gaussian.
sigma: KxDxD numpy array, the diagonal standard deviation of each gaussian. You will have KxDxD numpy
array for full covariance matrix case

Hint:
There are formulas in the slides and in the Jupyter Notebook.
Undergrads: To simplify your calculation in sigma, make sure to only take the diagonal terms in your covariance matrix
"""
N_k = np.sum(tau, axis=0)
pi = N_k / self.N
mu = np.dot(tau.T, self.points) / N_k[:, np.newaxis]

sigma = np.zeros((self.K, self.D, self.D))
for k in range(self.K):
diff = self.points - mu[k]
weighted_diff = diff * tau[:, k][:, np.newaxis]
sigma[k] = np.diag(
np.sum(weighted_diff * diff, axis=0) / N_k[k]
) + SIGMA_CONST * np.eye(self.D)

return pi, mu, sigma

def __call__(
self, full_matrix=FULL_MATRIX, abs_tol=1e-16, rel_tol=1e-16, **kwargs
): # No need to change
"""
Args:
abs_tol: convergence criteria w.r.t absolute change of loss
rel_tol: convergence criteria w.r.t relative change of loss
kwargs: any additional arguments you want

Return:
tau: NxK array, the posterior distribution (a.k.a, the soft cluster assignment) for each observation.
(pi, mu, sigma): (1xK np array, KxD numpy array, KxDxD numpy array)

Hint:
You do not need to change it. For each iteration, we process E and M steps, then update the paramters.
"""
pi, mu, sigma = self._init_components(**kwargs)
pbar = tqdm(range(self.max_iters))

prev_loss = None
for it in pbar:
tau = self._E_step(pi, mu, sigma, full_matrix)
pi, mu, sigma = self._M_step(tau, full_matrix)
joint_ll = self._ll_joint(pi, mu, sigma, full_matrix)
loss = -np.sum(self.logsumexp(joint_ll))
if it:
diff = np.abs(prev_loss - loss)
if diff < abs_tol and diff / prev_loss < rel_tol:
break
prev_loss = loss
pbar.set_description("iter %d, loss: %.4f" % (it, loss))
return tau, (pi, mu, sigma)


def cluster_pixels_gmm(image, K, max_iters=10, full_matrix=True):
"""
Clusters pixels in the input image

Each pixel can be considered as a separate data point (of length 3),
which you can then cluster using GMM. Then, process the outputs into
the shape of the original image, where each pixel is its most likely value.

Args:
image: input image of shape(H, W, 3)
K: number of components
max_iters: maximum number of iterations in GMM. Default is 10
full_matrix: whether we use full covariance matrix in Normal PDF or not. Default is True.
Return:
clustered_img: image of shape(H, W, 3) after pixel clustering

Hints:
What do mu and tau represent?
"""
raise NotImplementedError


write code for cluster_pixels_gmm

## 3.3 Image Compression and pixel clustering [10pts]

Images typically need a lot of bandwidth to be transmitted over the network. In order to optimize this process, most image processors perform lossy compression of images (lossy implies some information is lost in the process of compression).

In this section, you will use your GMM algorithm implementation to do pixel clustering and compress the images. That is to say, you would develop a lossy image compression algorithm. This question is autograded based on your GMM implementation.
(Hint: you can adjust the number of clusters formed and justify your answer based on visual inspection of the resulting images or on a different metric of your choosing)

Implement the `cluster_pixels_gmm` function in gmm.py. Each pixel can be considered as a separate data point (of length 3), which you can then cluster using GMM. Then, process the outputs into the shape of the original image, where each pixel is its most likely value. What do $\mu$ and $\tau$ represent?

**Special Notes**

- Try to add a small value(e.g. SIGMA_CONST and LOG_CONST) before taking the operation if the output image is solid black.
- The output images may be slightly different due to different initialization methods in GMM() function.
- Sample with replacement in `create_mu`