"""
This script generates embedding visualization for features produced by the apply_model script.
"""
import argparse
import inspect
import os
import pickle
import sys
from multiprocessing import Pool

import matplotlib as mpl
import torch

# To facilitate plotting on a headless server
mpl.use('Agg')
import matplotlib.pyplot as plt
from tensorboardX import SummaryWriter
import numpy as np
from sklearn.manifold import TSNE, Isomap, MDS
from sklearn.decomposition import PCA
from sklearn.preprocessing import LabelEncoder
from util.misc import load_numpy_image, save_numpy_image
from PIL import Image
########################################################################################################################
def tsne(features, n_components=2):
    """
    Returns the embedded points for TSNE.
    Parameters
    ----------
    features: numpy.ndarray
        contains the input feature vectors.
    n_components: int
        number of components to transform the features into

    Returns
    -------
    embedding: numpy.ndarray
        x,y(z) points that the feature vectors have been transformed into
    """
    embedding = TSNE(n_components=n_components).fit_transform(features)
    return embedding


def isomap(features, n_components=2):
    """
    Returns the embedded points for Isomap.
    Parameters
    ----------
    features: numpy.ndarray
        contains the input feature vectors.
    n_components: int
        number of components to transform the features into

    Returns
    -------
    embedding: numpy.ndarray
        x,y(z) points that the feature vectors have been transformed into
    """
    embedding = Isomap(n_components=n_components, n_jobs=-1).fit_transform(features)
    return embedding

def mds(features, n_components=2):
    """
    Returns the embedded points for MDS.
    Parameters
    ----------
    features: numpy.ndarray
        contains the input feature vectors.
    n_components: int
        number of components to transform the features into

    Returns
    -------
    embedding: numpy.ndarray
        x,y(z) points that the feature vectors have been transformed into
    """
    embedding = MDS(n_components=n_components, n_jobs=-1).fit_transform(features)
    return embedding

def pca(features, n_components=2):
    """
    Returns the embedded points for PCA.
    Parameters
    ----------
    features: numpy.ndarray
        contains the input feature vectors.
    n_components: int
        number of components to transform the features into

    Returns
    -------
    embedding: numpy.ndarray
        x,y(z) points that the feature vectors have been transformed into
    """
    embedding = PCA(n_components=n_components).fit_transform(features)
    return embedding

########################################################################################################################
def _make_embedding(features, labels, embedding, three_d=False):
    """
    Generate an embedding image using features from a model.

    Adapted from https://indico.io/blog/visualizing-with-t-sne/
    Parameters
    ----------
    features: numpy.ndarray
        contains the feature array generated by the apply_model runner class
    labels:  numpy.ndarray
        contains labels for corresponding feature vectors
    embedding: str
        type of embedding to use
    three_d: bool
        specify whether to generate 2d or 3d visualization

    Returns
    -------
    data: numpy.ndarray
        contains an image of the plotted visualization

    """

    plt.style.use(['seaborn-white', 'seaborn-paper'])
    fig = plt.figure(figsize=(8, 8))
    plt.tight_layout()
    mpl.rc("font", family="Times New Roman")

    X = features

    le = LabelEncoder()

    labels = le.fit_transform(labels)

    cmap = plt.cm.get_cmap('jet', len(np.unique(labels)))

    if three_d:
        ax = plt.axes(projection='3d')
        X_embedded = getattr(sys.modules[__name__], embedding)(X, n_components=3)
        ax.scatter3D(X_embedded[:, 0], X_embedded[:, 1], X_embedded[:, 2], c=labels, cmap=cmap)
    else:
        X_embedded = getattr(sys.modules[__name__], embedding)(X)
        plt.scatter(X_embedded[:, 0], X_embedded[:, 1], c=labels, cmap=cmap)

    # plt.colorbar(ticks=range(len(np.unique(labels))))

    fig.canvas.draw()
    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
    fig.clf()
    plt.close()
    return data


def _load_thumbnail(path):
    """
    Return a thumbnail version of any image
    Parameters
    ----------
    path: str
        path to an image

    Returns
    -------
    img: numpy.ndarray
        resized image of size 16x16
    """
    img = Image.open(path).thumbnail((16, 16))
    img = np.array(img)
    return img

def _make_folder_if_not_exists(path):
    if not os.path.exists(path):
        os.makedirs(path)

def _main(args):
    """
    Main routine of script to generate embeddings.
    Parameters
    ----------
    args : argparse.Namespace
        contains all arguments parsed from input

    Returns
    -------
    None

    """
    with open(args.results_file, 'rb') as f:
        results = pickle.load(f)

    features, preds, labels, filenames = results

    _make_folder_if_not_exists(os.path.dirname(args.output_file))

    if args.tensorboard:
        if args.output.endwith('.png'):
            output_loc = os.path.dirname(args.output)
        else:
            output_loc = args.output
        writer = SummaryWriter(log_dir=output_loc)
        # with Pool(16) as pool:
        # images = pool.map(_load_thumbnail, filenames)
        writer.add_embedding(torch.from_numpy(features), metadata=labels,
                             # label_img=torch.from_numpy(np.array(images)).unsqueeze(1))
                             label_img=None)
        return
    else:
        viz_img = _make_embedding(features=features, labels=labels, embedding=args.embedding, three_d=args.three_d)
        save_numpy_image(args.output_file, viz_img)
        return


if __name__ == "__main__":
    # Embedding options:
    embedding_options = [name[0] for name in inspect.getmembers(sys.modules[__name__], inspect.isfunction)]

    parser = argparse.ArgumentParser()

    parser.add_argument('--results-file',
                        type=str,
                        help='path to a results pickle file')

    parser.add_argument('--embedding',
                        help='which embedding to use for the features',
                        choices=embedding_options,
                        type=str)
    parser.add_argument('--output',
                        type=str,
                        default='./output.png',
                        help='path to generate output image')
    parser.add_argument('--3d',
                        dest='three_d',
                        action='store_true',
                        default=False,
                        help='enable 3d plots')
    parser.add_argument('--tensorboard',
                        action='store_true',
                        default=False,
                        help='store embeddings to tensorboard')

    args = parser.parse_args()

    _main(args)
