import numpy as np

from .base import BaseWidget, to_attr
from ..core.waveform_extractor import WaveformExtractor


class TemplateSimilarityWidget(BaseWidget):
    """
    Plots unit template similarity.

    Parameters
    ----------
    waveform_extractor : WaveformExtractor
        The object to compute/get template similarity from
    unit_ids : list
        List of unit ids default None
    display_diagonal_values : bool
        If False, the diagonal is displayed as zeros.
        If True, the similarity values (all 1s) are displayed, default False
    cmap : str
        The matplotlib colormap. Default 'viridis'.
    show_unit_ticks : bool
        If True, ticks display unit ids, default False.
    show_colorbar : bool
        If True, color bar is displayed, default True.
    """

    def __init__(
        self,
        waveform_extractor: WaveformExtractor,
        unit_ids=None,
        cmap="viridis",
        display_diagonal_values=False,
        show_unit_ticks=False,
        show_colorbar=True,
        backend=None,
        **backend_kwargs,
    ):
        self.check_extensions(waveform_extractor, "similarity")
        tsc = waveform_extractor.load_extension("similarity")
        similarity = tsc.get_data().copy()

        sorting = waveform_extractor.sorting
        if unit_ids is None:
            unit_ids = sorting.unit_ids
        else:
            unit_indices = sorting.ids_to_indices(unit_ids)
            similarity = similarity[unit_indices][:, unit_indices]

        if not display_diagonal_values:
            np.fill_diagonal(similarity, 0)

        plot_data = dict(
            similarity=similarity,
            unit_ids=unit_ids,
            cmap=cmap,
            show_unit_ticks=show_unit_ticks,
            show_colorbar=show_colorbar,
        )

        BaseWidget.__init__(self, plot_data, backend=backend, **backend_kwargs)

    def plot_matplotlib(self, data_plot, **backend_kwargs):
        import matplotlib.pyplot as plt
        from .utils_matplotlib import make_mpl_figure

        dp = to_attr(data_plot)

        self.figure, self.axes, self.ax = make_mpl_figure(**backend_kwargs)

        im = self.ax.matshow(dp.similarity, cmap=dp.cmap)

        if dp.show_unit_ticks:
            # Major ticks
            self.ax.set_xticks(np.arange(0, len(dp.unit_ids)))
            self.ax.set_yticks(np.arange(0, len(dp.unit_ids)))
            self.ax.xaxis.tick_bottom()

            # Labels for major ticks
            self.ax.set_yticklabels(dp.unit_ids, fontsize=12)
            self.ax.set_xticklabels(dp.unit_ids, fontsize=12)
        if dp.show_colorbar:
            self.figure.colorbar(im)

    def plot_sortingview(self, data_plot, **backend_kwargs):
        import sortingview.views as vv
        from .utils_sortingview import generate_unit_table_view, make_serializable, handle_display_and_url

        dp = to_attr(data_plot)

        # ensure serializable for sortingview
        unit_ids = make_serializable(dp.unit_ids)

        # similarity
        ss_items = []
        for i1, u1 in enumerate(unit_ids):
            for i2, u2 in enumerate(unit_ids):
                ss_items.append(
                    vv.UnitSimilarityScore(unit_id1=u1, unit_id2=u2, similarity=dp.similarity[i1, i2].astype("float32"))
                )

        self.view = vv.UnitSimilarityMatrix(unit_ids=list(unit_ids), similarity_scores=ss_items)

        self.url = handle_display_and_url(self, self.view, **backend_kwargs)
