import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from matplotlib.lines import Line2D
import matplotlib.colors as colors
import pandas as pd
import numpy as np
import seaborn as sns
from geospace.plot.plot import spatial, scale


def hotspots(adata, hotspot_key, ax=None, show=True, correct_categorical=False, labels=True,
             **kwargs):
    if ax is None:
        ax = plt.gca()
    if "title" not in kwargs:
        kwargs["title"] = " - ".join([v.capitalize() for v in hotspot_key.split("_")])

    key = "hotspots_%s" % hotspot_key

    if correct_categorical and not hasattr(adata.obs[key], 'cat'):
        v = pd.Categorical(adata.obs[key])
        v.categories = v.categories.astype(np.int_)
        adata.obs[key] = v

    if labels:
        legend_loc = "on data"
    else:
        legend_loc = None
    spatial(adata, color=key,
            ax=ax, show=False, legend_loc=legend_loc, na_in_legend=False,
            **kwargs)

    if show:
        plt.show()


def hotspots_multiple(adata, nrows, ncols, hotspot_keys=None, figsize=None):
    # interactions = ["WNT5A_FZD2", "WNT5B_FZD2", "WNT11_FZD2", "APLN_APLNR",
    #                "FGF5_FGFR1", "FGF10_FGFR1", "FGF15_FGFR1", "FGF5_FGFR2"]

    if hotspot_keys is None:
        obs_keys = list(adata.obs)
        hotspot_keys = [v[9:] for v in obs_keys if "hotspots_" in v]
        nrows = int(np.ceil(len(hotspot_keys) / ncols))

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

    for i, hotspot_key in enumerate(hotspot_keys):
        cur_row = i // ncols
        cur_col = i % ncols
        hotspots(adata, hotspot_key, ax=axs[cur_row, cur_col], show=False)

    for i in range(nrows):
        for j in range(ncols):
            axs[i, j].set_xlabel('')
            axs[i, j].set_ylabel('')
    plt.tight_layout()


def multi_hotspots(adata, show=True, inds=None, spotcolor=None, ax=None, num=None, linewidth=0.5,
                   legend_ncol=None, alpha_img=1.0, color_type=None, alpha=150, base_color=None,
                   frameon=False,
                   **kwargs):
    if color_type is None:
        color_type = "index"

    cm1 = mpl.cm.get_cmap('tab20')
    legend_handles = []
    legend_labels = []
    create_legend = True

    if color_type == "genes":
        create_legend = False
        vals = [len(v) for k, v in adata.uns["multi_hotspots"].items()]
        min_genes = min(vals)
        max_genes = max(vals)
        norm = colors.LogNorm(vmin=min_genes, vmax=max_genes)
        cmap = mpl.cm.YlOrRd

        sm = mpl.cm.ScalarMappable(norm=norm, cmap=cmap)


    try:
        scale_factor = list(adata.uns['spatial'].values())[0]['scalefactors']['tissue_hires_scalef']
    except KeyError:
        if scale is not None:
            scale_factor = scale
        else:
            scale_factor = 1

    if ax is None:
        fig, ax = plt.subplots()

    spatial(adata, color=base_color, ax=ax, alpha_img=alpha_img, frameon=frameon,
            na_color=spotcolor,
            show=False, **kwargs)

    hotspot_keys = [v for v in adata.obs if "hotspots_multi" in v]
    if inds is None:
        # Draw boundaries around all considered regions
        inds = list(range(len(hotspot_keys)))
    for c, idx in enumerate(inds):
        if num is not None and idx >= num:
            break
        if color_type == "index":
            color = cm1(c)
        elif color_type == "genes":
            color = sm.to_rgba(len(adata.uns["multi_hotspots"][str(idx)]))
        else:
            # assume it is a vector with colors
            color = colors.to_rgba(color_type[idx])

        color_int = [int(255 * v) for v in color]
        color_string = "#%02x%02x%02x%02x" % (color_int[0], color_int[1], color_int[2], alpha)
        color_string_noalpha = "#%02x%02x%02x" % (color_int[0], color_int[1], color_int[2])

        if True or idx < 21:
            legend_labels.append(idx)
            #legend_handles.append(Patch(facecolor=color_string, edgecolor='k',
            #                            label=idx))
            legend_handles.append(Line2D([0], [0], marker='o', color='w', label=idx,
                   markerfacecolor=color_string_noalpha, markersize=8))


        try:
            boundary = scale_factor*adata.uns["multi_boundaries"][str(idx)]
        except KeyError:
            raise ValueError("Need to compute boundaries first")
        ax.fill(boundary[:, 0], boundary[:, 1], color_string)
        ax.plot(boundary[:, 0], boundary[:, 1], linewidth=linewidth, color="black")
    if legend_ncol is not None and create_legend:
        ax.legend(legend_handles, legend_labels, loc="center left", bbox_to_anchor=(1, 0.5),
                  ncol=legend_ncol, frameon=False)
        #ax.legend(legend_handles, legend_labels, loc="upper center", bbox_to_anchor=(0.5, 0.0),
        #          ncol=legend_ncol)
    elif color_type == "genes":
        ax.get_figure().colorbar(sm, orientation='vertical', label='coexpressed genes',
                                 shrink=0.8)


    #plt.tight_layout()
    if show:
        plt.show()


