import glob
import re

import numpy as np
import pandas as pd
import plotly.graph_objs as go
import sklearn.manifold
from loguru import logger
from plotly.offline import iplot


class DashboardP2V(object):
    """
    Class for plotting networks weights and product maps
        plot_product_embedding: produce product embedding heat map plot (after P2V)
        plot_tsne_map: product product map (after t-SNE)
    """

    def __init__(self, path: str, n_steps: int = 20):
        """
        Class for plotting networks weights and product maps
            path: path that contains weights saved by pytorch trainer
            n_steps: number of steps in plot

        Note: The dashboard class automatically pick `n_steps` batches (linear split)
        unless the specific iterations are explicitly provided to
        `plot_product_embedding`. The method `plot_tsne_map` uses the same batches that
        `plot_product_embedding` uses.
        """
        logger.info(f"initialize P2V dashboard")
        self.path = path
        logger.info(f"path = {self.path}")
        self.n_steps = n_steps
        logger.info(f"n_steps = {self.n_steps}")
        self.files = {
            "wi": "wi_*.npy",
        }
        self.file_df = None
        self.plot_data = {}
        self.plot_data["steps"] = []
        self.plot_data["wi"] = []
        self.plot_data["tsne"] = []
        self.raw_files = glob.glob(f"{self.path}/{self.files['wi']}")

    def plot_product_embedding(
        self,
        epochs: list = None,
        l2norm: bool = True,
        transpose: bool = True,
        reload: bool = True,
    ) -> iplot:
        """
        Plot product embedding as heat map. Before plotting, several (optional)
        transformations can be applied to the data (see arguments).
            epochs: list of epochs to use for plot (see `__init__` for what happens if
            epochs is set to `None`)
            l2norm: Apply L2 normalization to embeddings?
            transpose: Transpose loaded embeddings?
            reload: Reload data (i.e., embeddings) before plotting?
        """
        if reload:
            self._load_data(epochs=epochs)
        data = []
        steps = []
        for i in range(self.n_steps):
            data_i = self.plot_data["wi"][i]
            # data
            if l2norm:
                data_i /= np.linalg.norm(data_i, axis=1)[:, np.newaxis]
            if transpose:
                data_i = data_i.T
            data.append(go.Heatmap(z=data_i, colorscale="Jet", zmin=-0.6, zmax=0.6))
            # step
            step = dict(
                method="restyle",
                label=self.plot_data["steps"][i],
                args=["visible", [False] * self.n_steps],
            )
            step["args"][1][i] = True
            steps.append(step)
        sliders = dict(
            active=0, currentvalue={"visible": False}, pad={"t": 50}, steps=steps
        )
        layout = go.Layout(
            height=700,
            width=1000,
            sliders=[sliders],
            margin=go.layout.Margin(l=50, r=50, b=150, t=20, pad=4),
            template="plotly_white",
        )
        return iplot(dict(data=data, layout=layout))

    def plot_tsne_map(
        self,
        product: pd.DataFrame,
    ) -> iplot:
        """
        Create product map (for multiple batches)
            product: Used to encode categories as color of bubbles in product map. The
            DataFrame must contain the columns `j` (product ID) and `c` (category ID).
        """
        if len(self.plot_data["tsne"]) < self.n_steps:
            self._tsne(product)
        data = []
        steps = []
        for i in range(self.n_steps):
            data_tsne_i = self.plot_data["tsne"][i]
            trace = go.Scatter(
                x=data_tsne_i["x"].values,
                y=data_tsne_i["y"].values,
                text=[
                    f"c: {c} <br> j: {c}"
                    for (c, j) in zip(
                        data_tsne_i["j"].values,
                        data_tsne_i["c"].values,
                    )
                ],
                hoverinfo="text",
                mode="markers",
                marker=dict(
                    size=12,
                    color=data_tsne_i["c"].values,
                    colorscale="Jet",
                    showscale=False,
                ),
            )
            data.append(trace)
            # step
            step = dict(
                method="restyle",
                label=self.plot_data["steps"][i],
                args=["visible", [False] * self.n_steps],
            )
            step["args"][1][i] = True
            steps.append(step)
        sliders = dict(
            active=0, currentvalue={"visible": False}, pad={"t": 50}, steps=steps
        )
        layout = go.Layout(
            height=700,
            width=800,
            sliders=[sliders],
            margin=go.layout.Margin(l=0, r=0, b=0, t=0, pad=4),
            # margin=go.layout.Margin(l=50, r=50, b=150, t=20, pad=4),
            hovermode="closest",
            template="plotly_white",
            xaxis=dict(range=[-2.2, 2.2]),
            yaxis=dict(range=[-2.2, 2.2]),
            showlegend=False,
        )
        return iplot(dict(data=data, layout=layout))

    def _tsne(self, product: pd.DataFrame) -> None:
        """
        Run t-SNE on product embedding
            product: see `plot_tsne_map`
        """
        for i in range(self.n_steps):
            if i == 0:
                init = "pca"
            else:
                init = self.plot_data["tsne"][i - 1][["x", "y"]].values
            tsne = sklearn.manifold.TSNE(n_components=2, perplexity=4, init=init)
            X = tsne.fit_transform(self.plot_data["wi"][i])
            X = X - X.mean(axis=0)
            X = X / X.std(axis=0)
            df = product.copy()
            df["x"] = X[:, 0]
            df["y"] = X[:, 1]
            self.plot_data["tsne"].append(df)

    def _build_file_df(self, epochs: list = None) -> None:
        """
        Build DataFrame that contains weight files (produced by trainer).
            epochs: see `plot_product_embedding`
        """
        files = [f for f in self.raw_files if re.search(r"(\d+)_(\d+).npy", f)]
        if not files:
            return None
        df = pd.DataFrame({"file": files})
        epoch_batch = df["file"].str.extract(r"(\d+)_(\d+).npy").astype(np.int32)
        epoch_batch.rename(columns={0: "epoch", 1: "batch"}, inplace=True)
        df = pd.concat([df, epoch_batch], axis=1)
        df = df.sort_values(["epoch", "batch"]).reset_index(drop=True)
        if self.n_steps < df.shape[0]:
            if epochs is None:
                rows_keep = (
                    np.linspace(0, 1, self.n_steps) ** 2 * (df.shape[0] - 1)
                ).astype(int)
            else:
                self.n_steps = df.shape[0]
                rows_keep = [i for i in epochs if i <= df.shape[0]]
                self.n_steps = len(rows_keep)
            df = df.iloc[rows_keep]
        self.file_df = df
        self.n_steps = np.min([self.file_df.shape[0], self.n_steps])

    def _load_data(self, epochs: list = None) -> None:
        """
        Load embedding matrix for selected epochs
            epochs: see `_build_file_df`
        """
        self._build_file_df(epochs=epochs)
        if self.file_df is not None:
            self.plot_data["wi"] = [np.load(f) for f in self.file_df.file.values]
            self.plot_data["steps"] = [
                f"e{x:02d}-b{y:06d}"
                for (x, y) in zip(self.file_df.epoch.values, self.file_df.batch.values)
            ]
        self.plot_data["tsne"] = []
