import pickle

import mnist
import seaborn as sns
import mpmath as mp
import numpy as np
from numpy.testing import assert_array_almost_equal as aae
from matplotlib import pyplot as plt
from scipy.spatial.distance import cdist
from scipy.stats import norm
from scipy.stats import multivariate_normal as mn

import sklearn
from sklearn.utils import check_random_state
from sklearn.utils.extmath import row_norms
import sklearn.cluster
from sklearn.decomposition import PCA
from sklearn.metrics import silhouette_score
#import umap

from _k_init import _k_init
from unionfind import UnionFind

COLORS = [
    '#cc5151', '#51cccc', '#337f7f', '#8ecc51', '#7f3333', '#597f33', '#8e51cc',
    '#59337f', '#ccad51', '#7f6c33', '#51cc70', '#337f46', '#5170cc', '#33467f',
    '#cc51ad', '#7f336c', '#cc7f51', '#7f4f33', '#bccc51', '#757f33', '#60cc51',
    '#3c7f33', '#51cc9e', '#337f62', '#519ecc', '#33627f', '#6051cc', '#3c337f'
]


def init_centers(X, n_clusters):
    """
    Run k-means++ to initialize centroids.
    Since we will be comparing to k-means, it's fair for both methods to have
    the same initialization method.

    Taken from scikit-learn.
    """
    random_state = check_random_state(None)
    n_samples = X.shape[0]

    x_squared_norms = row_norms(X, squared=True)

    centers = _k_init(
        X, n_clusters,
        random_state=random_state,
        x_squared_norms=x_squared_norms)

    return centers


class GMM:
    def __init__(self, n_clusters, n_steps, eps=1e-20):
        self.n_clusters = n_clusters
        self.n_steps = n_steps
        self.eps = eps

    def _initialize(self):
        """
        Initializes self.alpha, self.mu, self.sigma, self.w
        """
        self.alpha = np.ones((self.n_clusters)) / self.n_clusters
        self.mu = self.X[np.random.choice(np.arange(self.n), self.n_clusters)]
        self.sigma = np.ones((self.n_clusters, self.d))
        self.chunklet_w = np.zeros((self.n_chunklets, self.n_clusters))

        #centers = init_centers(X, self.n_clusters)
        #dists = cdist(X, centers)
        #labels = np.argmin(dists, axis=1)

        #unq_labels, self.alpha = np.unique(labels, return_counts=True)

        #self.alpha = np.zeros(self.n_clusters)
        #self.mu = np.zeros((self.n_clusters, d))
        # Using diagonal variance
        #self.sigma = np.zeros((self.n_clusters, d))

        # for i, lbl in enumerate(unq_labels):
        #    cur_pts = np.where(labels == lbl)
        #    self.alpha[i] = cur_pts[0].shape[0]
        #    # initialize means
        #    self.mu[i, :] = np.mean(X[cur_pts], axis=0)

        #    centered = (X[cur_pts] - self.mu[i])**2
        #    centered = np.sum(centered, axis=0) / centered.shape[0]
        #    # initialize vars
        #    self.sigma[i, :] = self.alpha[i] * centered

        #self.alpha /= n

        # self._validate_sigma()

        #self.chunklet_w = np.zeros((self.chunklets.shape[0], self.n_clusters))

    def _transitive_closure(self):
        self.uf = UnionFind(np.arange(self.n))
        for link in self.ml:
            self.uf.union(link[0], link[1])
        self.chunklets = np.array(
            [np.array(list(i)) for i in self.uf.components()])
        self.n_chunklets = self.chunklets.shape[0]
        self.chunklet_shapes = np.array([i.shape[0] for i in self.chunklets])
        self.chunklet_shapes = self.chunklet_shapes.reshape(-1, 1)
        self.chunklet_means = np.array(
            [np.mean(self.X[i], axis=0) for i in self.chunklets])
        assert self.chunklet_means.shape == (self.n_chunklets, self.d)

    def fit(self, X, ml):
        self.n = X.shape[0]
        self.d = X.shape[1]
        self.X = X.copy()
        self.ml = ml.copy()

        self._transitive_closure()
        self._initialize()

        self.scores = []
        self.lls = []

        for step in range(self.n_steps):
            self.e_step()
            self.m_step()
            self.scores.append(self.score())
            self.lls.append(self.ll)
            print(f"Step {step+1} :: LL {self.ll} :: Score {self.scores[-1]}")
            if len(self.lls) >= 2 and np.abs(self.lls[-1] - self.lls[-2]) < 1e-2:
                print("Converged")
                break

    def get_labels(self):
        chunk_labels = np.argmax(self.chunklet_w, axis=1).astype(np.int)
        labels = np.zeros(self.n)
        for i, chunk in enumerate(self.chunklets):
            labels[chunk] = chunk_labels[i]
        return labels.astype(np.int)

    def llhood(self):
        ll = 0
        for i, chunklet in enumerate(self.chunklets):
            for j in range(self.n_clusters):
                numerator = mn.pdf(
                    self.X[chunklet], self.mu[j], np.diag(self.sigma[j]))
                ll += np.sum(np.log(numerator + self.eps), axis=0) *\
                        self.chunklet_w[i,j]
                ll += np.log(self.alpha[j] + self.eps) * self.chunklet_w[i,j]
        return ll

    def e_step(self):
        self.ll = 0

        for i, chunklet in enumerate(self.chunklets):
            denominator = 0
            numerators = []
            for j in range(self.n_clusters):
                numerator = mn.pdf(
                    self.X[chunklet], self.mu[j], np.diag(self.sigma[j]))

                self.ll += np.sum(np.log(numerator + self.eps), axis=0) *\
                        self.chunklet_w[i,j]
                self.ll += np.log(self.alpha[j] + self.eps) *\
                        self.chunklet_w[i,j]

                numerator = np.prod(numerator, axis=0)
                numerator *= self.alpha[j]
                denominator += numerator
                self.chunklet_w[i, j] = numerator
            self.chunklet_w[i, :] /= (denominator + self.eps)
            #assert np.abs(self.chunklet_w[i, :].sum() - 1) < eps,\
            #    np.abs(self.chunklet_w[i, :].sum())

    def m_step(self):
        self.alpha = self.chunklet_w.sum(axis=0) / self.n_chunklets

        for j in range(self.n_clusters):
            den = 0
            temp_mu = np.zeros((1, self.d))

            numfrac = self.chunklet_w[:, j, np.newaxis] * self.chunklet_shapes
            den = np.sum(numfrac, axis=0, keepdims=True)
            temp_mu = np.sum(self.chunklet_means * numfrac, axis=0)
            self.mu[j] = temp_mu / den

            diff_sq = (self.X - self.mu[j])**2
            temp_sigma = np.zeros((1, self.d))
            for i in range(self.n_chunklets):
                # calc sigmanew
                signew = diff_sq[self.chunklets[i]]
                signew = np.sum(signew, axis=0, keepdims=True)
                signew /= self.chunklet_shapes[i]
                temp_sigma += signew * numfrac[i]

            self.sigma[j] = temp_sigma / den

    def score(self):
        labels = self.get_labels()
        return silhouette_score(self.X, labels)


def scatter_points(x1, x2, labels=None, path=None):
    if labels is not None:
        unq = len(np.unique(labels))
        pal = COLORS[:unq]

        sns.scatterplot(x=x1, y=x2, hue=labels, palette=pal, linewidth=0,
                        s=10, legend='full')
        lgd = plt.legend(bbox_to_anchor=(1.05, 1), loc=2,
                         borderaxespad=0, labels=np.unique(labels))
        sns.despine(left=True, bottom=True)

        if path is not None:
            plt.savefig(path, bbox_extra_artists=(lgd,), bbox_inches='tight')

        plt.xticks([])
        plt.yticks([])
        plt.show()
    else:
        plt.scatter(x1, x2, s=1)
        plt.xticks([])
        plt.yticks([])
        plt.show()


if __name__ == '__main__':
    #x = mnist.train_images()
    #y = mnist.train_labels()
    from load_cifar10 import load_cifar10
    X, y = load_cifar10()

    X = X.reshape(X.shape[0], -1)[:200] / 255
    pca = PCA(n_components=50)
    X = pca.fit_transform(X)
    # with open("dumps/pca_emb_200.pkl", "rb") as f:
    #    X = pickle.load(f)
    #X_emb = umap.UMAP(n_components=2).fit_transform(X)
    # with open("dumps/2d_emb.pkl", "rb") as f:
    #    X_emb = pickle.load(f)

    gmm = GMM(n_clusters=10, n_steps=5)
    ml = np.array([(1, 2), (2, 3), (3, 4), (4, 5), (7, 9), (50, 51)])
    gmm.fit(X, ml)

    #scatter_points(X_emb[:, 0], X_emb[:, 1], gmm.labels)

import pickle

import mnist
import seaborn as sns
import mpmath as mp
import numpy as np
from numpy.testing import assert_array_almost_equal as aae
from matplotlib import pyplot as plt
from scipy.spatial.distance import cdist
from scipy.stats import norm
from scipy.stats import multivariate_normal as mn

import sklearn
from sklearn.utils import check_random_state
from sklearn.utils.extmath import row_norms
import sklearn.cluster
from sklearn.decomposition import PCA
from sklearn.metrics import silhouette_score
#import umap

from _k_init import _k_init
from unionfind import UnionFind

COLORS = [
    '#cc5151', '#51cccc', '#337f7f', '#8ecc51', '#7f3333', '#597f33', '#8e51cc',
    '#59337f', '#ccad51', '#7f6c33', '#51cc70', '#337f46', '#5170cc', '#33467f',
    '#cc51ad', '#7f336c', '#cc7f51', '#7f4f33', '#bccc51', '#757f33', '#60cc51',
    '#3c7f33', '#51cc9e', '#337f62', '#519ecc', '#33627f', '#6051cc', '#3c337f'
]


def init_centers(X, n_clusters):
    """
    Run k-means++ to initialize centroids.
    Since we will be comparing to k-means, it's fair for both methods to have
    the same initialization method.

    Taken from scikit-learn.
    """
    random_state = check_random_state(None)
    n_samples = X.shape[0]

    x_squared_norms = row_norms(X, squared=True)

    centers = _k_init(
        X, n_clusters,
        random_state=random_state,
        x_squared_norms=x_squared_norms)

    return centers


class GMM:
    def __init__(self, n_clusters, n_steps, eps=1e-20):
        self.n_clusters = n_clusters
        self.n_steps = n_steps
        self.eps = eps

    def _initialize(self):
        """
        Initializes self.alpha, self.mu, self.sigma, self.w
        """
        self.alpha = np.ones((self.n_clusters)) / self.n_clusters
        self.mu = self.X[np.random.choice(np.arange(self.n), self.n_clusters)]
        self.sigma = np.ones((self.n_clusters, self.d))
        self.chunklet_w = np.zeros((self.n_chunklets, self.n_clusters))

        #centers = init_centers(X, self.n_clusters)
        #dists = cdist(X, centers)
        #labels = np.argmin(dists, axis=1)

        #unq_labels, self.alpha = np.unique(labels, return_counts=True)

        #self.alpha = np.zeros(self.n_clusters)
        #self.mu = np.zeros((self.n_clusters, d))
        # Using diagonal variance
        #self.sigma = np.zeros((self.n_clusters, d))

        # for i, lbl in enumerate(unq_labels):
        #    cur_pts = np.where(labels == lbl)
        #    self.alpha[i] = cur_pts[0].shape[0]
        #    # initialize means
        #    self.mu[i, :] = np.mean(X[cur_pts], axis=0)

        #    centered = (X[cur_pts] - self.mu[i])**2
        #    centered = np.sum(centered, axis=0) / centered.shape[0]
        #    # initialize vars
        #    self.sigma[i, :] = self.alpha[i] * centered

        #self.alpha /= n

        # self._validate_sigma()

        #self.chunklet_w = np.zeros((self.chunklets.shape[0], self.n_clusters))

    def _transitive_closure(self):
        self.uf = UnionFind(np.arange(self.n))
        for link in self.ml:
            self.uf.union(link[0], link[1])
        self.chunklets = np.array(
            [np.array(list(i)) for i in self.uf.components()])
        self.n_chunklets = self.chunklets.shape[0]
        self.chunklet_shapes = np.array([i.shape[0] for i in self.chunklets])
        self.chunklet_shapes = self.chunklet_shapes.reshape(-1, 1)
        self.chunklet_means = np.array(
            [np.mean(self.X[i], axis=0) for i in self.chunklets])
        assert self.chunklet_means.shape == (self.n_chunklets, self.d)

    def fit(self, X, ml):
        self.n = X.shape[0]
        self.d = X.shape[1]
        self.X = X.copy()
        self.ml = ml.copy()

        self._transitive_closure()
        self._initialize()

        self.scores = []
        self.lls = []

        for step in range(self.n_steps):
            self.e_step()
            self.m_step()
            self.scores.append(self.score())
            self.lls.append(self.ll)
            print(f"Step {step+1} :: LL {self.ll} :: Score {self.scores[-1]}")
            if len(self.lls) >= 2 and np.abs(self.lls[-1] - self.lls[-2]) < 1e-2:
                print("Converged")
                break

    def get_labels(self):
        chunk_labels = np.argmax(self.chunklet_w, axis=1).astype(np.int)
        labels = np.zeros(self.n)
        for i, chunk in enumerate(self.chunklets):
            labels[chunk] = chunk_labels[i]
        return labels.astype(np.int)

    def llhood(self):
        ll = 0
        for i, chunklet in enumerate(self.chunklets):
            for j in range(self.n_clusters):
                numerator = mn.pdf(
                    self.X[chunklet], self.mu[j], np.diag(self.sigma[j]))
                ll += np.sum(np.log(numerator + self.eps), axis=0) *\
                        self.chunklet_w[i,j]
                ll += np.log(self.alpha[j] + self.eps) * self.chunklet_w[i,j]
        return ll

    def e_step(self):
        self.ll = 0

        for i, chunklet in enumerate(self.chunklets):
            denominator = 0
            numerators = []
            for j in range(self.n_clusters):
                numerator = mn.pdf(
                    self.X[chunklet], self.mu[j], np.diag(self.sigma[j]))

                self.ll += np.sum(np.log(numerator + self.eps), axis=0) *\
                        self.chunklet_w[i,j]
                self.ll += np.log(self.alpha[j] + self.eps) *\
                        self.chunklet_w[i,j]

                numerator = np.prod(numerator, axis=0)
                numerator *= self.alpha[j]
                denominator += numerator
                self.chunklet_w[i, j] = numerator
            self.chunklet_w[i, :] /= (denominator + self.eps)
            #assert np.abs(self.chunklet_w[i, :].sum() - 1) < eps,\
            #    np.abs(self.chunklet_w[i, :].sum())

    def m_step(self):
        self.alpha = self.chunklet_w.sum(axis=0) / self.n_chunklets

        for j in range(self.n_clusters):
            den = 0
            temp_mu = np.zeros((1, self.d))

            numfrac = self.chunklet_w[:, j, np.newaxis] * self.chunklet_shapes
            den = np.sum(numfrac, axis=0, keepdims=True)
            temp_mu = np.sum(self.chunklet_means * numfrac, axis=0)
            self.mu[j] = temp_mu / den

            diff_sq = (self.X - self.mu[j])**2
            temp_sigma = np.zeros((1, self.d))
            for i in range(self.n_chunklets):
                # calc sigmanew
                signew = diff_sq[self.chunklets[i]]
                signew = np.sum(signew, axis=0, keepdims=True)
                signew /= self.chunklet_shapes[i]
                temp_sigma += signew * numfrac[i]

            self.sigma[j] = temp_sigma / den

    def score(self):
        labels = self.get_labels()
        return silhouette_score(self.X, labels)


def scatter_points(x1, x2, labels=None, path=None):
    if labels is not None:
        unq = len(np.unique(labels))
        pal = COLORS[:unq]

        sns.scatterplot(x=x1, y=x2, hue=labels, palette=pal, linewidth=0,
                        s=10, legend='full')
        lgd = plt.legend(bbox_to_anchor=(1.05, 1), loc=2,
                         borderaxespad=0, labels=np.unique(labels))
        sns.despine(left=True, bottom=True)

        if path is not None:
            plt.savefig(path, bbox_extra_artists=(lgd,), bbox_inches='tight')

        plt.xticks([])
        plt.yticks([])
        plt.show()
    else:
        plt.scatter(x1, x2, s=1)
        plt.xticks([])
        plt.yticks([])
        plt.show()


if __name__ == '__main__':
    #x = mnist.train_images()
    #y = mnist.train_labels()
    from load_cifar10 import load_cifar10
    X, y = load_cifar10()

    X = X.reshape(X.shape[0], -1)[:200] / 255
    pca = PCA(n_components=50)
    X = pca.fit_transform(X)
    # with open("dumps/pca_emb_200.pkl", "rb") as f:
    #    X = pickle.load(f)
    #X_emb = umap.UMAP(n_components=2).fit_transform(X)
    # with open("dumps/2d_emb.pkl", "rb") as f:
    #    X_emb = pickle.load(f)

    gmm = GMM(n_clusters=10, n_steps=5)
    ml = np.array([(1, 2), (2, 3), (3, 4), (4, 5), (7, 9), (50, 51)])
    gmm.fit(X, ml)

    #scatter_points(X_emb[:, 0], X_emb[:, 1], gmm.labels)

#import umap Run k-means++ to initialize centroids.     Since we will be comparing to k-means, it's fair for both methods to have     the same initialization method.      Taken from scikit-learn. Initializes self.alpha, self.mu, self.sigma, self.w #centers = init_centers(X, self.n_clusters) #dists = cdist(X, centers) #labels = np.argmin(dists, axis=1) #unq_labels, self.alpha = np.unique(labels, return_counts=True) #self.alpha = np.zeros(self.n_clusters) #self.mu = np.zeros((self.n_clusters, d)) # Using diagonal variance #self.sigma = np.zeros((self.n_clusters, d)) # for i, lbl in enumerate(unq_labels): #    cur_pts = np.where(labels == lbl) #    self.alpha[i] = cur_pts[0].shape[0] #    # initialize means #    self.mu[i, :] = np.mean(X[cur_pts], axis=0) #    centered = (X[cur_pts] - self.mu[i])**2 #    centered = np.sum(centered, axis=0) / centered.shape[0] #    # initialize vars #    self.sigma[i, :] = self.alpha[i] * centered #self.alpha /= n # self._validate_sigma() #self.chunklet_w = np.zeros((self.chunklets.shape[0], self.n_clusters)) #assert np.abs(self.chunklet_w[i, :].sum() - 1) < eps,\ #    np.abs(self.chunklet_w[i, :].sum()) # calc sigmanew #x = mnist.train_images() #y = mnist.train_labels() # with open("dumps/pca_emb_200.pkl", "rb") as f: #    X = pickle.load(f) #X_emb = umap.UMAP(n_components=2).fit_transform(X) # with open("dumps/2d_emb.pkl", "rb") as f: #    X_emb = pickle.load(f) #scatter_points(X_emb[:, 0], X_emb[:, 1], gmm.labels)