from PIL.PngImagePlugin import PngImageFile

from sparclur.parsers.present_parsers import get_sparclur_renderers
from sparclur._renderer import Renderer
import itertools
import matplotlib.pyplot as plt
from matplotlib.cm import ScalarMappable
from matplotlib.colors import to_rgb
import warnings
import numpy as np
import os
from sparclur.prc._prc import _parse_viz_renderers

AVAILABLE_RENDERERS = {r.get_name(): r for r in get_sparclur_renderers()}


class PRCViz:
    """
    PDF Render Comparator

    Contains methods for comparing and visualizing the renders generated by SPARCLUR wrapped renderers.
    """

    def __init__(self, doc_path,
                 renderers=get_sparclur_renderers(),
                 parser_args=dict(),
                 dpi=200,
                 verbose=False):
        """

        Parameters
        ----------
        doc_path : str
            Full path to the document to run through the PRC
        renderers : List[Renderer] or List[str]
            The list of renderers to use for the PRC comparison
        parser_args : Dict[str, Dict[str, Any]]
            A dictionary of dictionaries containing any optional parameters to pass into the renderers. See each
            renderer for it's possible parameters.
        """
        self._doc_path = doc_path
        self._doc = doc_path.split('/')[-1]
        self._renderers = _parse_viz_renderers(renderers)
        self._renders = dict()
        self._sims = dict()
        self._ssims_fig = None
        if verbose:
            print('Rendering:')
        for name, renderer in self._renderers.items():
            if verbose:
                print('\t%s:' % name)
            # args = parser_args.get(name, dict())
            # args['cache_renders'] = True
            # args['dpi'] = dpi
            # self._renders[name] = AVAILABLE_RENDERERS[name](doc_path=doc_path, **args)
            if isinstance(renderer, Renderer):
                self._renders[name] = renderer
                self._renders[name].caching = True
            else:
                args = parser_args.get(name, dict())
                args['cache_renders'] = True
                args['dpi'] = dpi
                self._renders[name] = renderer(doc=doc_path, **args)
            # if mpg_path is not None:
            #     self._mpg_renders[name] = renderer(doc_path=doc_path, **args)
        assert len(set([renderer.doc for renderer in self._renders.values()])) == 1, \
            "Document paths do not match for all renderers"
        self._sim_keys = list(itertools.combinations(self._renders.keys(), 2))
        for combo in self._sim_keys:
            if verbose:
                print('SIMing %s/%s' % (combo[0], combo[1]))
            self._sims[combo] = self._renders[combo[0]].compare(self._renders[combo[1]], full=True)
        self._observed_pages = max(len(entry) for entry in self._sims.values())

    def get_observed_pages(self):
        """Return the number of observed pages from the renderers"""
        return self._observed_pages

    def plot_sims(self, cmap='tab10', height=10, width=10, save_path=None):
        """
        Plot the structural similarity score for each page and each combination of renderers.

        Parameters
        ----------
        cmap: str or List[str]
            Either a matplotlib cmap name or a specific list of desired colors.
        height: int
            The height of the matplotlib figure
        width: int
            The width of the matplotlib figure
        save_path: str or None
            If None returns the figure for display, otherwise saves the figure to the specified file path.

        """
        if self._ssims_fig is None:
            fig = plt.figure()
            ax = fig.add_subplot(1, 1, 1)

            if isinstance(cmap, list) and len(cmap) >= len(self._sim_keys):
                colors = [to_rgb(color) for color in cmap]
            else:
                if len(cmap) < len(self._sim_keys):
                    warnings.warn("Not enough colors specified. Defaulting to tab10 cmap")
                    cmap = 'tab10'
                scalar_mappable = ScalarMappable(cmap=cmap)
                colors = scalar_mappable.to_rgba(range(len(self._sim_keys)), alpha=1.0).tolist()

            for (clr_idx, combo) in enumerate(self._sim_keys):
                color = colors[clr_idx]
                label = '%s/%s' % (combo[0], combo[1])
                x = list(self._sims[combo].keys())
                y = [prc_sim.sim for (page, prc_sim) in self._sims[combo].items()]
                ax.scatter(x, y, color=color, label=label)
            fig.set_figwidth(width)
            fig.set_figheight(height)
            plt.title("SIM by page for %s" % self._doc)
            plt.legend(loc=0, prop={'size': 15})

            if save_path is not None:
                fig.savefig(save_path)
                plt.close(fig)
            plt.close(fig)
            self._ssims_fig = fig
        return self._ssims_fig

    def display(self, page, renderers=None, width=10, height=10, save_path=None):
        """
        Show the comparison between the specified renderers and the visual difference between them.

        Parameters
        ----------
        page: int
            The page to visualize
        renderers: List[tuple] or tuple
            If None, then visualizes all the combinations. Otherwise specify a list of combinations or just one single
            combination.
        width: int
            The width of each subplot
        height: int
            The height of each subplot
        save_path: str or None
            If None returns the figure for display, otherwise saves the figure to the specified file path.

        """
        if renderers is None:
            renderers = self._sim_keys
        if isinstance(renderers, tuple):
            renderers = [renderers]
        nrows = len(renderers)

        fig, axes = plt.subplots(nrows=nrows, ncols=3, figsize=(width, height))

        fig.suptitle('Comparisons for %s\nPage: %s' % (self._doc, page))

        for (row, combo) in enumerate(renderers):
            images = [self._renders[combo[0]].get_renders(page), self._renders[combo[1]].get_renders(page), self._sims[combo][page].diff]
            labels = ['', '', self._sims[combo][page].sim]
            titles = [combo[0], combo[1], 'diff']
            if nrows > 1:
                for col in range(3):
                    image = np.asarray(images[col])
                    axes[row, col].set_title(titles[col])
                    axes[row, col].set_xticklabels([])
                    axes[row, col].set_yticklabels([])
                    axes[row, col].set_xticks([])
                    axes[row, col].set_yticks([])
                    axes[row, col].imshow(image)
                    axes[row, col].set_xlabel(labels[col])
            else:
                for col in range(3):
                    image: PngImageFile = images[col]
                    axes[col].set_title(titles[col])
                    axes[col].set_xticklabels([])
                    axes[col].set_yticklabels([])
                    axes[col].set_xticks([])
                    axes[col].set_yticks([])
                    axes[col].imshow(image)
                    axes[col].set_xlabel(labels[col])

        if save_path is not None:
            fig.savefig(os.path.join(save_path, '%s_%s_prc.png' % (self._doc, page)))
            plt.close(fig)
        else:
            plt.close(fig)
            return fig
