from scipy.spatial import distance
from scipy import ndimage
import matplotlib.pyplot as plt
import torch
from scipy import stats
import numpy as np


def pairwise_distances_fig(embs):
    embs = embs.detach().cpu().numpy()
    similarity_matrix_cos = distance.cdist(embs, embs, 'cosine')
    similarity_matrix_euc = distance.cdist(embs, embs, 'euclidean')

    fig = plt.figure(figsize=(16,10))

    ax = fig.add_subplot(121)
    cax = ax.matshow(similarity_matrix_cos)
    fig.colorbar(cax)
    ax.set_title('Cosine')
    ax.axis('off')

    ax = fig.add_subplot(122)
    cax = ax.matshow(similarity_matrix_euc)
    fig.colorbar(cax)
    ax.set_title('Euclidian')
    ax.axis('off')

    fig.suptitle('Action pairwise distances')
    plt.close()
    return fig


def pairwise_distances(embs):
    fig = pairwise_distances_fig(embs)
    fig.show()


def smooth(scalars, weight):  # Weight between 0 and 1
    last = scalars[0]  # First value in the plot (first timestep)
    smoothed = list()
    for point in scalars:
        smoothed_val = last * weight + (1 - weight) * point  # Calculate smoothed value
        smoothed.append(smoothed_val)                        # Save it
        last = smoothed_val                                  # Anchor the last smoothed value

    return smoothed


def smooth_gauss(arr, var):
    return ndimage.gaussian_filter1d(arr, var)


class Plotter:
    def __init__(self, loss, style):
        self.loss = loss
        self.style = style
        self.smoothing = lambda x: smooth_gauss(x, 4)

    def set_smoothing_func(self, f):
        self.smoothing = f

    def plot_loss(self):
        for row in self.style:
            fig, axes = plt.subplots(1, len(row), figsize=(16, 6))
            if len(row) == 1: axes = [axes]
            for col in range(len(row)):
                key = row[col]
                axes[col].set_title(key)
                axes[col].plot(self.loss['train']['step'],
                               self.smoothing(self.loss['train'][key]), 'b-',
                               label='train')
                axes[col].plot(self.loss['test']['step'],
                               self.loss['test'][key], 'r-.',
                               label='test')
            plt.legend()
        plt.show()

    def log_loss(self, key, item, test=False):
        kind = 'train'
        if test:
            kind = 'test'
        self.loss[kind][key].append(item)

    def log_losses(self, losses, test=False):
        for key, val in losses.items():
            self.log_loss(key, val, test)

    @staticmethod
    def kde_reconstruction_error(ad, gen_actions, true_actions, device=torch.device('cpu')):

        def rec_score(actions):
            return ad.rec_error(torch.tensor(actions).to(device).float()).detach().cpu().numpy()

        true_scores = rec_score(true_actions)
        gen_scores = rec_score(gen_actions)

        true_kernel = stats.gaussian_kde(true_scores)
        gen_kernel = stats.gaussian_kde(gen_scores)

        x = np.linspace(0, 1000, 100)
        probs_true = true_kernel(x)
        probs_gen = gen_kernel(x)
        fig = plt.figure(figsize=(16, 10))
        ax = fig.add_subplot(111)
        ax.plot(x, probs_true, '-b', label='true dist')
        ax.plot(x, probs_gen, '-r', label='generated dist')
        ax.legend()
        return fig

    @staticmethod
    def plot_kde_reconstruction_error(*args, **kwargs):
        fig = Plotter.kde_reconstruction_error(*args, **kwargs)
        fig.show()


[((11, 28, 11, 64), 'scipy.spatial.distance.cdist', 'distance.cdist', ({(11, 43, 11, 47): 'embs', (11, 49, 11, 53): 'embs', (11, 55, 11, 63): '"""cosine"""'}, {}), "(embs, embs, 'cosine')", False, 'from scipy.spatial import distance\n'), ((12, 28, 12, 67), 'scipy.spatial.distance.cdist', 'distance.cdist', ({(12, 43, 12, 47): 'embs', (12, 49, 12, 53): 'embs', (12, 55, 12, 66): '"""euclidean"""'}, {}), "(embs, embs, 'euclidean')", False, 'from scipy.spatial import distance\n'), ((14, 10, 14, 37), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((29, 4, 29, 15), 'matplotlib.pyplot.close', 'plt.close', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((50, 11, 50, 46), 'scipy.ndimage.gaussian_filter1d', 'ndimage.gaussian_filter1d', ({(50, 37, 50, 40): 'arr', (50, 42, 50, 45): 'var'}, {}), '(arr, var)', False, 'from scipy import ndimage\n'), ((76, 8, 76, 18), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((89, 71, 89, 90), 'torch.device', 'torch.device', ({(89, 84, 89, 89): '"""cpu"""'}, {}), "('cpu')", False, 'import torch\n'), ((97, 22, 97, 53), 'scipy.stats.gaussian_kde', 'stats.gaussian_kde', ({(97, 41, 97, 52): 'true_scores'}, {}), '(true_scores)', False, 'from scipy import stats\n'), ((98, 21, 98, 51), 'scipy.stats.gaussian_kde', 'stats.gaussian_kde', ({(98, 40, 98, 50): 'gen_scores'}, {}), '(gen_scores)', False, 'from scipy import stats\n'), ((100, 12, 100, 37), 'numpy.linspace', 'np.linspace', ({(100, 24, 100, 25): '0', (100, 27, 100, 31): '1000', (100, 33, 100, 36): '100'}, {}), '(0, 1000, 100)', True, 'import numpy as np\n'), ((103, 14, 103, 42), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((75, 12, 75, 24), 'matplotlib.pyplot.legend', 'plt.legend', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((92, 32, 92, 53), 'torch.tensor', 'torch.tensor', ({(92, 45, 92, 52): 'actions'}, {}), '(actions)', False, 'import torch\n')]