import anndata as ad
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

from anndata import AnnData

from matplotlib.axes import Axes
from matplotlib.figure import Figure
from matplotlib.patches import Patch
import matplotlib.cm

import scgenome.cnplot
import scgenome.refgenome


# TODO: refactor color functionality elsewhere
color_reference = {
    0:'#3182BD',
    1:'#9ECAE1',
    2:'#CCCCCC',
    3:'#FDCC8A',
    4:'#FC8D59',
    5:'#E34A33',
    6:'#B30000',
    7:'#980043',
    8:'#DD1C77',
    9:'#DF65B0',
    10:'#C994C7',
    11:'#D4B9DA',
}


def hex_to_rgb(h):
    if h is None:
        return np.array((0, 0, 0), dtype=int)
    h = h.lstrip('#')
    return np.array(tuple(np.uint8(int(h[i:i+2], 16)) for i in (0, 2 ,4)), dtype=int)


def plot_cell_cn_matrix(
        adata: AnnData,
        layer_name='state',
        cell_order_fields=(),
        ax=None,
        ax_legend=None,
        raw=False,
        max_cn=13,
        show_cell_ids=False):
    """ Plot a copy number matrix

    Parameters
    ----------
    adata : AnnData
        copy number data
    layer_name : str, optional
        layer with copy number data to plot, None for X, by default 'state'
    cell_order_fields : list, optional
        columns of obs on which to sort cells, by default None
    ax : matplotlib.axes.Axes, optional
        existing axis to plot into, by default None
    raw : bool, optional
        raw plotting, no integer color map, by default False
    max_cn : int, optional
        clip cn at max value, by default 13
    show_cell_ids : bool, optional
        show cell ids on heatmap axis, by default False

    Returns
    -------
    dict
        Dictionary of plot and data elements

    Examples
    -------

    .. plot::
        :context: close-figs

        import scgenome
        adata = scgenome.datasets.OV2295_HMMCopy_reduced()
        scgenome.pl.plot_cell_cn_matrix(adata)

    """

    if ax is None:
        ax = plt.gca()

    # Order the chromosomes
    chr_start = adata.var.reset_index().merge(scgenome.refgenome.info.chromosome_info[['chr', 'chr_index']], how='left')
    if chr_start['chr_index'].isnull().any():
        chromosomes = adata.var['chr'].astype(str).values
        raise ValueError(f'mismatching chromosomes {chromosomes} and {scgenome.refgenome.info.chromosomes}')
    chr_start = chr_start[['start', 'chr_index']].values
    genome_ordering = np.lexsort(chr_start.transpose())

    # Order the cells if requested
    if len(cell_order_fields) > 0:
        cell_order_fields = reversed(list(cell_order_fields))
        cell_order_values = adata.obs[cell_order_fields].values.transpose()
        cell_ordering = np.lexsort(cell_order_values)

    else:
        cell_ordering = range(adata.shape[0])

    adata = adata[cell_ordering, genome_ordering]

    if layer_name is not None:
        X = adata.layers[layer_name].copy()
    else:
        X = adata.X.copy()

    X = np.nan_to_num(X, nan=0)

    if max_cn is not None:
        X[X > max_cn] = max_cn

    cmap = None
    if not raw:
        X_colors = np.zeros(X.shape + (3,), dtype=int)
        X_colors[X < 0, :] = 0
        X_colors[X > max(color_reference.keys()), :] = max(color_reference.keys())
        for state, hex in color_reference.items():
            X_colors[X == state, :] = hex_to_rgb(hex)
        im = ax.imshow(X_colors, aspect='auto', cmap=cmap, interpolation='none')

    else:
        im = sns.heatmap(X, ax=ax, cbar=False)

    mat_chrom_idxs = chr_start[genome_ordering][:, 1]
    chrom_boundaries = np.array([0] + list(np.where(mat_chrom_idxs[1:] != mat_chrom_idxs[:-1])[0]) + [mat_chrom_idxs.shape[0] - 1])
    chrom_sizes = chrom_boundaries[1:] - chrom_boundaries[:-1]
    chrom_mids = chrom_boundaries[:-1] + chrom_sizes / 2
    ordered_mat_chrom_idxs = mat_chrom_idxs[np.where(np.array([1] + list(np.diff(mat_chrom_idxs))) != 0)]
    chrom_names = np.array(scgenome.refgenome.info.plot_chromosomes)[ordered_mat_chrom_idxs]

    ax.set_xticks(chrom_mids)
    ax.set_xticklabels(chrom_names, fontsize='6')

    if show_cell_ids:
        ax.set(yticks=range(len(adata.obs.index)))
        ax.set(yticklabels=adata.obs.index.values)
    else:
        ax.set(yticklabels=[])

    for val in chrom_boundaries[:-1]:
        ax.axvline(x=val, linewidth=1, color='black', zorder=100)

    legend = None
    if ax_legend is not None:
        states = []
        patches = []
        for s, h in color_reference.items():
            states.append(s)
            patches.append(Patch(facecolor=h, edgecolor=h))
        legend = ax_legend.legend(patches, states, ncol=3,
            frameon=True, loc=2, bbox_to_anchor=(0., 1.),
            facecolor='white', edgecolor='white', fontsize='4',
            title='Copy Number', title_fontsize='6')

    return {
        'ax': ax,
        'ax_legend': ax_legend,
        'im': im,
        'legend': legend,
        'adata': adata,
    }


def map_catagorigal_colors(values, cmap_name=None):
    levels = np.unique(values)
    n_levels = len(levels)

    if cmap_name is None:
        if n_levels <= 10:
            cmap_name = 'tab10'
        elif n_levels <= 20:
            cmap_name = 'tab20'
        else:
            cmap_name = 'hsv'

    cmap = matplotlib.cm.get_cmap(cmap_name)

    level_colors = dict(zip(levels, cmap(np.linspace(0, 1, n_levels))))

    value_colors = np.zeros(values.shape + (4,))
    for l, c in level_colors.items():
        value_colors[values == l, :] = c

    return level_colors, value_colors


def _plot_categorical_annotation(values, ax, ax_legend, title):
        level_colors, value_colors = map_catagorigal_colors(values)

        im = ax.imshow(value_colors, aspect='auto', interpolation='none')

        levels = []
        patches = []
        for s, h in level_colors.items():
            levels.append(s)
            patches.append(Patch(facecolor=h, edgecolor=h))
        legend = ax_legend.legend(patches, levels, ncol=3,
            frameon=True, loc=2, bbox_to_anchor=(0., 1.),
            facecolor='white', edgecolor='white', fontsize='4',
            title=title, title_fontsize='6')

        ax.grid(False)
        ax.set_xticks([0.], [title], rotation=90, fontsize='6')
        ax.set_yticks([])

        annotation_info = {}
        annotation_info = {}
        annotation_info['ax'] = ax
        annotation_info['im'] = im
        annotation_info['ax_legend'] = ax_legend
        annotation_info['legend'] = legend

        return annotation_info


def _plot_continuous_annotation(values, ax, ax_legend, title):

    im = ax.imshow(values, aspect='auto', interpolation='none', cmap='Reds')

    ax.grid(False)
    ax.set_xticks([0.], [title], rotation=90, fontsize='6')
    ax.set_yticks([])

    ax_legend.grid(False)
    ax_legend.set_xticks([])
    ax_legend.set_yticks([])

    axins = ax_legend.inset_axes([0.5, 0.1, 0.05, 0.6])

    cbar = plt.colorbar(im, cax=axins)
    axins.set_title(title, fontsize='6')
    cbar.ax.tick_params(labelsize='4')

    return {}


def plot_cell_cn_matrix_clusters_fig(
        adata: AnnData,
        layer_name='state',
        cell_order_fields=(),
        annotation_fields=(),
        fig=None,
        raw=False,
        max_cn=13,
        show_cell_ids=False):
    """ Plot a copy number matrix

    Parameters
    ----------
    adata : AnnData
        copy number data
    layer_name : str, optional
        layer with copy number data to plot, None for X, by default 'state'
    cell_order_fields : list, optional
        columns of obs on which to sort cells, by default None
    annotation_fields : list, optional
        column of obs to use as an annotation colorbar, by default 'cluster_id'
    fig : matplotlib.figure.Figure, optional
        existing figure to plot into, by default None
    raw : bool, optional
        raw plotting, no integer color map, by default False
    max_cn : int, optional
        clip cn at max value, by default 13
    show_cell_ids : bool, optional
        show cell ids on heatmap axis, by default False

    Returns
    -------
    dict
        Dictionary of plot and data elements

    Examples
    -------

    .. plot::
        :context: close-figs

        import scgenome
        adata = scgenome.datasets.OV2295_HMMCopy_reduced()

        g = scgenome.pl.plot_cell_cn_matrix_clusters_fig(
            adata,
            cell_order_fields=['cell_order'],
            annotation_fields=['cluster_id', 'sample_id', 'quality'])

    """

    if fig is None:
        fig = plt.figure()

    fig_main, fig_legends = fig.subfigures(nrows=2, ncols=1, height_ratios=[5, 1], squeeze=True)
    fig_legends.patch.set_alpha(0.0)

    width_ratios = [1] + [0.005] + [0.02] * len(annotation_fields)

    axes = fig_main.subplots(
        nrows=1, ncols=len(width_ratios), sharey=True, width_ratios=width_ratios,
        squeeze=False, gridspec_kw=dict(hspace=0.02, wspace=0.02))[0]
    axes[1].set_axis_off()

    axes_legends = fig_legends.subplots(
        nrows=1, ncols=1+len(annotation_fields), squeeze=False)[0]
    for ax in axes_legends:
        ax.set_axis_off()
        ax.set_alpha(0.0)
        ax.patch.set_alpha(0.0)

    ax = axes[0]
    ax_legend = axes_legends[0]
    g = plot_cell_cn_matrix(
        adata, layer_name=layer_name,
        cell_order_fields=cell_order_fields,
        ax=ax, ax_legend=ax_legend, raw=raw, max_cn=max_cn,
        show_cell_ids=show_cell_ids)

    adata = g['adata']

    annotation_info = {}

    for ax, ax_legend, annotation_field in zip(axes[2:], axes_legends[1:], annotation_fields):
        if adata.obs[annotation_field].dtype.name in ('category', 'object'):
            values = adata.obs[[annotation_field]].values
            annotation_info[annotation_field] = _plot_categorical_annotation(values, ax, ax_legend, annotation_field)

        else:
            values = adata.obs[[annotation_field]].values
            annotation_info[annotation_field] = _plot_continuous_annotation(values, ax, ax_legend, annotation_field)

    return {
        'fig': fig,
        'axes': axes,
        'adata': adata,
        'annotation_info': annotation_info,
    }
