# AUTOGENERATED! DO NOT EDIT! File to edit: 33_EM.ipynb (unless otherwise specified).

__all__ = ['KMeansInit', 'EM', 'MultiEM', 'generateBags']

# Cell
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import scipy.stats as ss
from easydict import EasyDict
from scipy.spatial.distance import cdist
from sklearn.metrics import roc_auc_score,adjusted_mutual_info_score, adjusted_rand_score
from tqdm.notebook import tqdm

# Cell


# export
def KMeansInit(X, k):
    means = [X[np.random.choice(np.arange(X.shape[0]))]]
    for iteration in range(k-1):
        dists = cdist(X, np.stack(means)).min(1)
        probs = dists / dists.sum()
        nextIndex = np.random.choice(np.arange(X.shape[0]),p=probs)
        means.append(X[nextIndex])
    return np.stack(means)

class EM:
    def __init__(self, X, k):
        self.X = X
        self.k = k
        self.N,self.dim = self.X.shape
        self.initializeComponents()

    def initializeComponents(self):
        mins = self.X.min(0)
        maxs = self.X.max(0)
        # USE K-MEANS++
#         self.mus = KMeansInit(self.X, self.k)
        self.mus = self.X[np.random.choice(np.arange(self.X.shape[0]),size=self.k)]
        self.covs = np.stack([np.eye(self.dim) for _ in range(self.k)])
        self.pi = np.random.dirichlet(np.ones(self.k))

    def E_Step(self):
        gammas = np.zeros((self.N,self.k))
        for n in range(self.N):
            for k in range(self.k):
                gammas[n,k] = self.pi[k] * ss.multivariate_normal.pdf(self.X[n],
                                                                      mean=self.mus[k],
                                                                      cov=self.covs[k])
            gammas[n] = gammas[n] / gammas[n].sum()
        self.gammas = gammas

    def M_Step(self):
        Nk = np.sum(self.gammas, axis=0)[:,np.newaxis]
        self.mus = np.dot(self.gammas.T, self.X) / Nk
        self.pi = np.mean(self.gammas,axis=0)
        for k in range(self.k):
            x = np.matrix(self.X - self.mus[k,:])
            gamma_diag = np.matrix(np.diag(self.gammas[:,k]))
            sigma_k = x.T * gamma_diag * x
            self.covs[k,:,:] = sigma_k / Nk[k]


    def log_likelihood(self):
        ll = 0
        for xn in self.X:
            ll += np.log(np.sum([pik * ss.multivariate_normal.pdf(xn,
                                                                  mean=muk,
                                                                  cov=covk) for pik,muk,covk in zip(self.pi,
                                                                                                    self.mus,
                                                                                                    self.covs)]))
        return ll

# Cell
class MultiEM:
    def __init__(self, bags, n_components):
        self.em_instances = []
        self.labels = []
        self.bags = bags
        # Initialize each EM instance
        for bag in self.bags:
            self.em_instances.append(EM(bag.X_pos, n_components))
        # Override the initialization

    def run_positive_em(self, n_iters, multi=True, KMeansPPInit=True):
        if KMeansPPInit:
            if multi:
                for inst in range(len(self.em_instances)):
                    self.em_instances[inst].mus = KMeansInit(np.concatenate([e.X for e in self.em_instances]),
                                                      self.em_instances[0].k)
            else:
                self.em_instances[0].mus = KMeansInit(np.concatenate([e.X for e in self.em_instances]),
                                                      self.em_instances[0].k)
        self.logLikelihoods = np.zeros((len(self.em_instances),
                                        n_iters+2))
        for inst in range(len(self.em_instances)):
            self.logLikelihoods[inst,0] = self.em_instances[inst].log_likelihood()
        for iteration in tqdm(range(1, n_iters+1),total=n_iters,leave=False):
            for inst in range(len(self.em_instances)):
                self.em_instances[inst].E_Step()
                self.em_instances[inst].M_Step()
                if multi:
                    self.em_instances[(inst+1) % len(self.em_instances)].mus = self.em_instances[inst].mus
                    self.em_instances[(inst+1) % len(self.em_instances)].covs = self.em_instances[inst].covs
                self.logLikelihoods[inst,iteration] = self.em_instances[inst].log_likelihood()
        # After final iteration, pass mu and cov to all instances and update gammas through the E_Step
        for j in range(len(self.em_instances)):
            if multi:
                self.em_instances[j].mus = self.em_instances[0].mus
                self.em_instances[j].covs = self.em_instances[0].covs
                # update gamma with new mus and covs values
                self.em_instances[j].E_Step()
            self.logLikelihoods[j, -1] = self.em_instances[j].log_likelihood()
        self.compute_positive_clustering_performance()

    def compute_positive_clustering_performance(self):
        adjusted_rand_scores = []
        for em,bag in zip(self.em_instances, self.bags):
            adjusted_rand_scores.append(adjusted_rand_score(bag.positive_component_labels,
                                                            np.argmax(em.gammas,axis=1)))
        self.score = np.mean(adjusted_rand_scores)

    def cluster_unlabeled(self):
        for bagnum, (bag, em) in enumerate(zip(self.bags, self.em_instances)):
            assignments = np.zeros(bag.x_unlabeled.shape[0]).astype(int)
            for instnum,inst in enumerate(bag.x_unlabeled):
                lls = np.zeros(em.k)
                for i,(mu,cov) in enumerate(zip(em.mus, em.covs)):
                    lls[i] = ss.multivariate_normal.logpdf(inst,mean=mu,cov=cov)
                assignments[instnum] = np.argmax(lls)
            self.bags[bagnum].unlabeled_assignments = assignments

# Cell
def generateBags(NBags,
                 NPos=200,
                 NUnlabeled=1000,
                 pos_means=[[2,2], [-2,-2], [2,-2]],
                 neg_means=[[0,0]],
                 pos_covs=[np.eye(2)]*3,
                 neg_covs=[np.eye(2)]):
    bags = []
    K = len(pos_means)
    D = len(pos_means[0])
    for _ in tqdm(range(NBags)):
        bag = EasyDict()
        bag.alpha = np.random.uniform()
        bag.pi = np.random.dirichlet(np.ones(K)*1.5)
        bag.rho = np.random.dirichlet(np.ones(K)*1.5)
        bag.gamma = bag.alpha * bag.pi + (1 - bag.alpha) * bag.rho
        bag.eta = bag.alpha * bag.pi / bag.gamma
        bag.posClusterAssignment = np.random.choice(np.arange(K), p=bag.pi.ravel(),size=NPos)
        bag.unlabeledPosClusterAssignment = np.random.choice(np.arange(K),
                                                         p=bag.pi.ravel(),
                                                         size=np.round(NUnlabeled * bag.alpha).astype(int))
        upLabels, upCounts = np.unique(bag.unlabeledPosClusterAssignment,return_counts=True)
        bag.empiricalPi = np.zeros(K)
        for label,count in zip(upLabels,upCounts):
            bag.empiricalPi[label] = count / upCounts.sum()
        bag.unlabeledNegClusterAssignment = np.random.choice(np.arange(K),
                                                         p=bag.rho.ravel(),
                                                         size=np.round(NUnlabeled * (1 - bag.alpha)).astype(int))
        unLabels, unCounts = np.unique(bag.unlabeledNegClusterAssignment,return_counts=True)
        bag.empiricalRho = np.zeros(K)
        for label,count in zip(unLabels,unCounts):
            bag.empiricalRho[label] = count / unCounts.sum()
        bag.empiricalGamma = bag.alpha * bag.empiricalPi + (1-bag.alpha) * bag.empiricalRho
        bag.empiricalEta = bag.alpha * bag.empiricalPi / bag.empiricalGamma
        bag.X_pos = np.zeros((0,D))
        bag.hiddenLabels = np.concatenate((np.ones_like(bag.unlabeledPosClusterAssignment),
                                           np.zeros_like(bag.unlabeledNegClusterAssignment)))
        bag.x_unlabeled = np.zeros((0,D))
        # sample positive data
        for ci in tqdm(bag.posClusterAssignment,leave=False):
            xi = ss.multivariate_normal(pos_means[ci], pos_covs[ci]).rvs()
            if D==1:
                xi = xi[None,None]
            else:
                xi = xi[None]
            bag.X_pos = np.concatenate((bag.X_pos,xi))

        # sample unlabeled pos data
        for ci in tqdm(bag.unlabeledPosClusterAssignment,leave=False):
            xi = ss.multivariate_normal(pos_means[ci], pos_covs[ci]).rvs()
            if D==1:
                xi = xi[None,None]
            else:
                xi = xi[None]
            bag.x_unlabeled = np.concatenate((bag.x_unlabeled, xi))
        # sample unlabeled neg data
        for ci in tqdm(bag.unlabeledNegClusterAssignment,leave=False):
            xi = ss.multivariate_normal(neg_means[ci], neg_covs[ci]).rvs()
            if D==1:
                xi = xi[None,None]
            else:
                xi = xi[None]
            bag.x_unlabeled = np.concatenate((bag.x_unlabeled, xi))
        bags.append(bag)
    return bags