import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.cm as cmx
from matplotlib.animation import FuncAnimation
import sklearn
from tqdm import tqdm_notebook
from sklearn import manifold
from numpy import linalg

class tsneAnimate():
    def __init__(self,tsne):
        self.tsne = tsne
        self.isfit = False
    def getSteps(self,X,y):
        #based on https://github.com/oreillymedia/t-SNE-tutorial
        old_grad = sklearn.manifold.t_sne._gradient_descent
        positions = []
        def _gradient_descent(objective, p0, it, n_iter, objective_error=None,
                              n_iter_check=1, n_iter_without_progress=50,
                              momentum=0.5, learning_rate=1000.0, min_gain=0.01,
                              min_grad_norm=1e-7, min_error_diff=1e-7, verbose=0,
                              args=None, kwargs=None):
            if args is None:
                args = []
            if kwargs is None:
                kwargs = {}

            p = p0.copy().ravel()
            update = np.zeros_like(p)
            gains = np.ones_like(p)
            error = np.finfo(np.float).max
            best_error = np.finfo(np.float).max
            best_iter = 0

            for i in range(it, n_iter):
                # We save the current position.
                positions.append(p.copy())

                new_error, grad = objective(p, *args, **kwargs)
                grad_norm = linalg.norm(grad)

                inc = update * grad >= 0.0
                dec = np.invert(inc)
                gains[inc] += 0.05
                gains[dec] *= 0.95
                np.clip(gains, min_gain, np.inf)
                grad *= gains
                update = momentum * update - learning_rate * grad
                p += update

                if (i + 1) % n_iter_check == 0:
                    if new_error is None:
                        new_error = objective_error(p, *args)
                    error_diff = np.abs(new_error - error)
                    error = new_error

                    if verbose >= 2:
                        m = "[t-SNE] Iteration %d: error = %.7f, gradient norm = %.7f"
                        print(m % (i + 1, error, grad_norm))

                    if error < best_error:
                        best_error = error
                        best_iter = i
                    elif i - best_iter > n_iter_without_progress:
                        if verbose >= 2:
                            print("[t-SNE] Iteration %d: did not make any progress "
                                  "during the last %d episodes. Finished."
                                  % (i + 1, n_iter_without_progress))
                        break
                    if grad_norm <= min_grad_norm:
                        if verbose >= 2:
                            print("[t-SNE] Iteration %d: gradient norm %f. Finished."
                                  % (i + 1, grad_norm))
                        break
                    if error_diff <= min_error_diff:
                        if verbose >= 2:
                            m = "[t-SNE] Iteration %d: error difference %f. Finished."
                            print(m % (i + 1, error_diff))
                        break

                if new_error is not None:
                    error = new_error

            return p, error, i

        #Replace old gradient func
        sklearn.manifold.t_sne._gradient_descent = _gradient_descent
        X_proj = self.tsne.fit_transform(X)
        self.isfit = True
        #return old gradient descent back
        sklearn.manifold.t_sne._gradient_descent = old_grad
        return positions
    
    def animate(self,X,y,filename='animate.gif',useTqdm=0):
        pos = self.getSteps(X,y)
        y_mapping = {i:n for n,i in enumerate(set(y))}
        
        last_iter = pos[len(pos)-1].reshape(-1, 2)
        lims = np.max(last_iter,axis=0),np.min(last_iter,axis=0)
        NCOLORS = len(y_mapping)
        fig = plt.figure()
        fig.set_tight_layout(True)
        ax = fig.add_subplot(111)
        jet = plt.get_cmap('jet') 
        cNorm  = colors.Normalize(vmin=0, vmax=NCOLORS)
        scalarMap = cmx.ScalarMappable(norm=cNorm, cmap=jet)
        
        A,B = np.array(list(zip(*pos[0].reshape(-1, 2))))
        dots_list = []
        for i in range(NCOLORS):
            colorVal = scalarMap.to_rgba(i)
            a,b = A[y == i],B[y == i]
            dots, = ax.plot(b,a,'o',color=colorVal)
            dots_list.append(dots)
        
        
        def init():
            ax.set_xlim([lims[0][0],lims[1][0]])
            ax.set_ylim([lims[0][1],lims[1][1]])
            return (*dots_list),

        def update(i):
            for j in range(len(dots_list)):
                a,b = np.array(list(zip(*pos[i].reshape(-1, 2))))
                a,b = a[y == j],b[y == j]
                dots_list[j].set_xdata(a)
                dots_list[j].set_ydata(b) 
            return (*dots_list), ax

        if useTqdm==0:
            frames = np.arange(0, len(pos)-1)
        elif useTqdm==1:
            from tqdm import tqdm
            frames = tqdm(np.arange(0, len(pos)-1))
        elif useTqdm==2:
            from tqdm import tqdm_notebook
            frames = tqdm_notebook(np.arange(0, len(pos)-1))
        
        anim = FuncAnimation(fig, update, frames=frames, init_func=init, interval=50)
        anim.save(filename, dpi=80, writer='imagemagick')