from matplotlib import pyplot as plt
import matplotlib.lines as mlines
import numpy as np
import seaborn as sns
import os

sns.set()


default_color = ["r", "g", "b", "c", "m", "y", "orange", "brown"]
default_marker = ["*", "o", "s", "v", "X", "*", ".", "P"]


def plot_figures(figures, nrows=1, ncols=1, figsize=(15, 15), titles=None):
    """Plot a dictionary of figures.

    Parameters
    ----------
    figures : <title, figure> dictionary
    ncols : number of columns of subplots wanted in the display
    nrows : number of rows of subplots wanted in the figure
    """

    fig, axeslist = plt.subplots(ncols=ncols, nrows=nrows, figsize=figsize)

    assert (titles is None) or (len(titles) == len(figures))

    for ind, title in enumerate(figures):
        axeslist.ravel()[ind].imshow(figures[ind])
        if titles is None:
            axeslist.ravel()[ind].set_title(ind)
        else:
            axeslist.ravel()[ind].set_title(titles[ind])
        axeslist.ravel()[ind].set_axis_off()
    plt.tight_layout()  # optional

    return fig


def self_balancing_plot(random_identity_pics):

    size = len(random_identity_pics)
    nrows = int(round(np.sqrt(size)))
    ncols = int(np.ceil(size / nrows))
    return plot_figures(random_identity_pics, nrows, ncols)


class CMC:
    def __init__(self, cmc_dict, color=default_color, marker=default_marker):
        self.color = color
        self.marker = marker
        self.cmc_dict = cmc_dict

    def plot(
        self,
        title,
        rank=20,
        xlabel="Rank",
        ylabel="Matching Rates (%)",
        show_grid=True,
        figsize=(10, 10),
    ):
        fig, ax = plt.subplots(figsize=figsize)
        fig.suptitle(title)
        x = np.arange
        #         plt.ylim(0, 1.0)
        #         plt.xlim(1, rank)

        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.xticks(np.arange(1.0, rank, step=10.0))
        plt.grid(show_grid)

        method_name = []
        i = 0
        for name in self.cmc_dict.keys():
            if rank < len(self.cmc_dict[name]):
                temp_cmc = self.cmc_dict[name][:rank]
                r = list(range(1, rank + 1))
            else:
                temp_cmc = self.cmc_dict[name]
                r = list(range(1, len(temp_cmc) + 1))

            if name == list(self.cmc_dict.keys())[-1]:
                globals()[name] = mlines.Line2D(
                    r,
                    temp_cmc,
                    color=self.color[0],
                    marker=self.marker[0],
                    label="{:.1f}% {}".format(self.cmc_dict[name][0] * 100, name),
                )
            else:
                globals()[name] = mlines.Line2D(
                    r,
                    temp_cmc,
                    color=self.color[i + 1],
                    marker=self.marker[i + 1],
                    label="{:.1f}% {}".format(self.cmc_dict[name][0] * 100, name),
                )
                i = i + 1
            ax.add_line(globals()[name])
            method_name.append(globals()[name])

        plt.legend(handles=method_name)
        plt.xscale("log")
        plt.show()

    def save(
        self,
        title,
        filename,
        rank=20,
        xlabel="Rank",
        ylabel="Matching Rates (%)",
        show_grid=True,
        save_path=os.getcwd(),
        format="png",
        **kwargs
    ):
        fig, ax = plt.subplots()
        fig.suptitle(title)
        x = list(range(0, rank + 1, 5))
        plt.ylim(0, 1.0)
        plt.xlim(1, rank)
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.xticks(x)
        plt.grid(show_grid)

        method_name = []
        i = 0
        for name in self.cmc_dict.keys():
            if rank < len(self.cmc_dict[name]):
                temp_cmc = self.cmc_dict[name][:rank]
                r = list(range(1, rank + 1))
            else:
                temp_cmc = self.cmc_dict[name]
                r = list(range(1, len(temp_cmc) + 1))

            if name == list(self.cmc_dict.keys())[-1]:
                globals()[name] = mlines.Line2D(
                    r,
                    temp_cmc,
                    color="r",
                    marker="*",
                    label="{:.1f}% {}".format(self.cmc_dict[name][0] * 100, name),
                )
            else:
                globals()[name] = mlines.Line2D(
                    r,
                    temp_cmc,
                    color=self.color[i],
                    marker=self.marker[i],
                    label="{:.1f}% {}".format(self.cmc_dict[name][0] * 100, name),
                )
                i = i + 1
            ax.add_line(globals()[name])
            method_name.append(globals()[name])

        plt.legend(handles=method_name)
        fig.savefig(
            os.path.join(save_path, filename + "." + format),
            format=format,
            bbox_inches="tight",
            pad_inches=0,
            **kwargs
        )
