import os

import numpy as np
import seaborn as sns
from matplotlib import pyplot as plt

from cospar import tool as tl
from cospar.plotting import _utils as pl_util

from .. import help_functions as hf
from .. import logging as logg
from .. import settings


def barcode_heatmap(
    adata,
    selected_times=None,
    selected_fates=None,
    color_bar=True,
    rename_fates=None,
    normalize=False,
    binarize=False,
    log_transform=False,
    fig_width=4,
    fig_height=6,
    figure_index="",
    plot=True,
    pseudocount=10 ** (-10),
    order_map_x=False,
    order_map_y=False,
    fate_normalize_source="X_clone",
    select_clones_with_fates: list = None,
    select_clones_without_fates: list = None,
    select_clones_mode: str = "or",
    **kwargs,
):
    """
    Plot barcode heatmap among different fate clusters.

    We clonal measurement at selected time points and show the
    corresponding heatmap among selected fate clusters.

    Parameters
    ----------
    adata: :class:`~anndata.AnnData` object
    selected_times: `list`, optional (default: None)
        Time points to select the cell states.
    selected_fates: `list`, optional (default: all)
        List of fate clusters to use. If set to be [], use all.
    color_bar: `bool`, optional (default: True)
        Plot color bar.
    rename_fates: `list`, optional (default: None)
        Provide new names in substitution of names in selected_fates.
        For this to be effective, the new name list needs to have names
        in exact correspondence to those in the old list.
    normalize:
        To perform cluster-wise then clone-wise normalization per time point
    binarize: `bool`
        Binarize the coarse-grained barcode count matrix, just for the purpose of plotting.
    log_transform: `bool`, optional (default: False)
        If true, perform a log transform. This is needed when the data
        matrix has entries varying by several order of magnitude.
    fig_width: `float`, optional (default: 4)
        Figure width.
    fig_height: `float`, optional (default: 6)
        Figure height.
    plot: `bool`
        True: plot the result. False, suppress the plot.
    pseudocount: `float`
        Pseudocount for the heatmap (needed for ordering the map)
    order_map_x: `bool`
        Whether to re-order the x coordinate of the matrix or not
    order_map_y: `bool`
        Whether to re-order the y coordinate of the matrix or not
    fate_normalize_source:
        Source for cluster-wise normalization: {'X_clone','state_info'}. 'X_clone': directly row-normalize coarse_X_clone; 'state_info': compute each cluster size directly, and then normalize coarse_X_clone. The latter method is useful if we have single-cell resolution for each fate.
    select_clones_with_fates: list = None,
        Select clones that labels fates from this list.
    select_clones_without_fates: list = None,
        Exclude clones that labels fates from this list.
    select_clones_mode: str = {'or','and'}
        Logic rule for selection.

    Returns:
    --------
    The coarse-grained X_clone matrix and the selected clusters are returned at
    adata.uns['barcode_heatmap']. The coarse-grained X_clone keeps all clones and maintains their ordering.
    """

    data_des = adata.uns["data_des"][-1]
    data_des = f"{data_des}_clonal"
    figure_path = settings.figure_path

    coarse_X_clone, mega_cluster_list = tl.coarse_grain_clone_over_cell_clusters(
        adata,
        selected_times=selected_times,
        selected_fates=selected_fates,
        normalize=normalize,
        fate_normalize_source=fate_normalize_source,
        select_clones_with_fates=select_clones_with_fates,
        select_clones_without_fates=select_clones_without_fates,
        select_clones_mode=select_clones_mode,
        **kwargs,
    )

    if rename_fates is None:
        rename_fates = mega_cluster_list

    if len(rename_fates) != len(mega_cluster_list):
        logg.warn(
            "rename_fates does not have the same length as selected_fates, thus not used."
        )
        rename_fates = mega_cluster_list

    if "x_ticks" not in kwargs.keys():
        kwargs["x_ticks"] = rename_fates

    coarse_X_clone_new = pl_util.custom_hierachical_ordering(
        np.arange(coarse_X_clone.shape[0]), coarse_X_clone
    )
    adata.uns["barcode_heatmap"] = {
        "coarse_X_clone": coarse_X_clone,
        "fate_names": rename_fates,
    }
    logg.info("Data saved at adata.uns['barcode_heatmap']")
    if plot:
        if binarize:
            final_matrix = coarse_X_clone_new > 0
            color_bar_label = "Binarized barcode count"
        else:
            final_matrix = coarse_X_clone_new
            color_bar_label = "Barcode count"

        if normalize:
            color_bar_label += " (normalized)"

        clone_idx = final_matrix.sum(0) > 0
        ax = pl_util.heatmap(
            final_matrix[:, clone_idx].T + pseudocount,
            order_map_x=order_map_x,
            order_map_y=order_map_y,
            color_bar_label=color_bar_label,
            log_transform=log_transform,
            fig_width=fig_width,
            fig_height=fig_height,
            color_bar=color_bar,
            **kwargs,
        )
        plt.title(f"{np.sum(clone_idx)} clones")

        plt.tight_layout()
        if figure_index != "":
            figure_index == f"_{figure_index}"
        plt.savefig(
            os.path.join(
                figure_path,
                f"{data_des}_barcode_heatmap{figure_index}.{settings.file_format_figs}",
            )
        )
        return ax


def clonal_fates_across_time(adata, selected_times, **kwargs):
    """
    Scatter plot for clonal fate number across time point

    Parameters
    ----------
    adata: :class:`~anndata.AnnData` object
    selected_times: `list`, optional (default: None)
        Time points to select the cell states.

    Returns
    -------
    Results updated at adata.uns["clonal_fates_across_time"]
    """
    if len(selected_times) != 2:
        raise ValueError("selected_times must be a list with two values")
    barcode_heatmap(
        adata,
        selected_times=selected_times[0],
        color_bar=True,
        log_transform=False,
        plot=False,
    )
    clonal_fates_t1 = (adata.uns["barcode_heatmap"]["coarse_X_clone"] > 0).sum(0)
    barcode_heatmap(
        adata,
        selected_times=selected_times[1],
        color_bar=True,
        log_transform=False,
        plot=False,
    )
    clonal_fates_t2 = (adata.uns["barcode_heatmap"]["coarse_X_clone"] > 0).sum(0)

    pl_util.jitter(clonal_fates_t1, clonal_fates_t2, **kwargs)
    plt.xlabel(f"Number of fates per clone (t={selected_times[0]})")
    plt.ylabel(f"Number of fates per clone (t={selected_times[1]})")
    data_des = adata.uns["data_des"][0]
    plt.savefig(
        os.path.join(
            settings.figure_path,
            f"{data_des}_barcode_coupling_across_time.{settings.file_format_figs}",
        )
    )
    adata.uns["clonal_fates_across_time"] = {
        "clonal_fates_t1": clonal_fates_t1,
        "clonal_fates_t2": clonal_fates_t2,
    }
    logg.info("Data saved at adata.uns['clonal_fates_across_time']")


def clones_on_manifold(
    adata,
    selected_clone_list=[0],
    color_list=["red", "blue", "purple", "green", "cyan", "black"],
    selected_annotation_key="time_info",
    selected_annotation_info=None,
    title=None,
    clone_markersize=12,
    clone_markeredgewidth=1,
    markeredgecolor="black",
    **kwargs,
):
    """
    Plot clones on top of state embedding.

    Parameters
    ----------
    adata: :class:`~anndata.AnnData` object
    selected_clone_list: `list`
        List of selected clone ID's.
    color_list: `list`, optional (default: ['red','blue','purple','green','cyan','black'])
        The list of color that defines color at respective time points.
    selected_annotation_info: `list`, optional (default: all)
        Select time points to show corresponding states. If set to be [], use all states.
    title: None
        Default, show the clone id as panel title. Otherwise, show the text string from input 'title'
    clone_markersize: `int`, optional (default: 12)
        Clone marker size
    clone_markeredgewidth: `int`, optional (default: 1)
        Edige size for clone marker
    """

    fig_width = settings.fig_width
    fig_height = settings.fig_height
    point_size = settings.fig_point_size
    x_emb = adata.obsm["X_emb"][:, 0]
    y_emb = adata.obsm["X_emb"][:, 1]
    data_des = adata.uns["data_des"][-1]
    # data_path=settings.data_path
    figure_path = settings.figure_path
    X_clone = adata.obsm["X_clone"]
    time_info = np.array(adata.obs[selected_annotation_key])

    # use only valid time points
    sp_idx = hf.selecting_cells_by_time_points(time_info, selected_annotation_info)
    selected_annotation_info = np.sort(list(set(time_info[sp_idx])))

    selected_clone_list = np.array(selected_clone_list)
    full_id_list = np.arange(X_clone.shape[1])
    valid_idx = np.isin(full_id_list, selected_clone_list)
    if np.sum(valid_idx) < len(selected_clone_list):
        logg.error(
            f"Valid id range is (0,{X_clone.shape[1]-1}). Please use a smaller ID!"
        )
        selected_clone_list = full_id_list[valid_idx]

    if len(selected_clone_list) == 0:
        logg.error("No valid states selected.")
    else:
        # using all data
        for my_id in selected_clone_list:
            fig = plt.figure(figsize=(fig_width, fig_height))
            ax = plt.subplot(1, 1, 1)
            idx_t = np.zeros(len(time_info), dtype=bool)
            for j, xx in enumerate(selected_annotation_info):
                idx_t0 = time_info == selected_annotation_info[j]
                idx_t = idx_t0 | idx_t

            pl_util.customized_embedding(
                x_emb[idx_t],
                y_emb[idx_t],
                np.zeros(len(y_emb[idx_t])),
                ax=ax,
                point_size=point_size,
            )
            for j, xx in enumerate(selected_annotation_info):
                idx_t = time_info == selected_annotation_info[j]
                idx_clone = X_clone[:, my_id].toarray().flatten() > 0
                idx = idx_t & idx_clone
                ax.plot(
                    x_emb[idx],
                    y_emb[idx],
                    ".",
                    color=color_list[j % len(color_list)],
                    markersize=clone_markersize,
                    markeredgecolor=markeredgecolor,
                    markeredgewidth=clone_markeredgewidth,
                    **kwargs,
                )

                if title is None:
                    ax.set_title(f"ID: {my_id}")
                else:
                    ax.set_title(title)

            fig.savefig(
                os.path.join(
                    figure_path,
                    f"{data_des}_different_clones_{my_id}.{settings.file_format_figs}",
                )
            )
    return ax


def clonal_fate_bias(adata, show_histogram=True, FDR=0.05):
    """
    Plot clonal fate bias towards a cluster.

    The results should be pre-computed from :func:`cospar.tl.clonal_fate_bias`

    Parameters
    ----------
    adata: :class:`~anndata.AnnData` object
    show_histogram: `bool`, optional (default: True)
        If true, show the distribution of inferred fate probability.
    FDR: `float`, optional (default: 0.05)
        False-discovery rate after the Benjamini-Hochberg correction.
    """

    if "clonal_fate_bias" not in adata.uns.keys():
        raise ValueError(
            "clonal_fate_bias has not been computed. Run cs.tl.clonal_fate_bias first"
        )
    else:
        df = adata.uns["clonal_fate_bias"]
        fate_bias = df["Fate_bias"]
        target_fraction_array = df["clonal_fraction_in_target_fate"]

    fig_width = settings.fig_width
    fig_height = settings.fig_height
    data_des = adata.uns["data_des"][-1]
    figure_path = settings.figure_path
    FDR_threshold = -np.log10(FDR)

    fig = plt.figure(figsize=(fig_width, fig_height))
    ax = plt.subplot(1, 1, 1)
    ax.plot(np.arange(len(fate_bias)), fate_bias, ".", color="blue", markersize=5)
    ax.plot(
        np.arange(len(fate_bias)),
        np.zeros(len(fate_bias)) + FDR_threshold,
        "-.",
        color="grey",
        markersize=5,
        label=f"FDR={FDR}",
    )

    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    # ax.set_ylabel('Fate bias ($-\\log_{10}P_{value}$)')
    ax.set_ylabel("Clonal fate bias")
    ax.set_xlabel("Clonal index")
    ax.legend()
    fig.tight_layout()
    fig.savefig(
        os.path.join(
            figure_path, f"{data_des}_clonal_fate_bias.{settings.file_format_figs}"
        )
    )

    if show_histogram:
        fig = plt.figure(figsize=(fig_width, fig_height))
        ax = plt.subplot(1, 1, 1)
        ax.hist(target_fraction_array, color="#2ca02c", density=True)
        ax.set_xlim([0, 1])
        ax.set_xlabel("Clonal fraction in selected fates")
        ax.set_ylabel("Density")
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)
        ax.set_title(f"Average: {np.mean(target_fraction_array):.2f}")
        fig.tight_layout()
        fig.savefig(
            os.path.join(
                figure_path,
                f"{data_des}_observed_clonal_fraction.{settings.file_format_figs}",
            )
        )


def clonal_reports(adata, selected_times=None, **kwargs):
    """
    Report the statistics of the clonal data.

    It includes the statistics for clone size , and the barcode number per cell.
    """

    time_info = np.array(adata.obs["time_info"])
    sp_idx = hf.selecting_cells_by_time_points(time_info, selected_times)
    adata_1 = adata[sp_idx]
    persistent_clone_ids = tl.identify_persistent_clones(adata_1)
    X_clone = adata_1.obsm["X_clone"]
    total_clone_N = X_clone.shape[1]
    print(
        f"  Clones observed across selected times: {len(persistent_clone_ids)} (out of {total_clone_N} clones)"
    )

    for x in set(adata_1.obs["time_info"]):
        print(f"---------t={x}---------")
        adata_sp = adata_1[adata_1.obs["time_info"] == x]
        X_clone = adata_sp.obsm["X_clone"]
        clone_size = X_clone.toarray().sum(0).flatten()
        clonal_bc_number = X_clone.toarray().sum(1).flatten()
        clonal_cells_N = np.sum(clonal_bc_number > 0)
        total_N = X_clone.shape[0]
        total_clone_N = X_clone.shape[1]
        useful_clone_N = np.sum(clone_size > 0)
        print(f"    Cells with barcode: {clonal_cells_N} (out of {total_N} cells)")
        print(
            f"    Barcodes with cells: {useful_clone_N} (out of {total_clone_N} clones)"
        )

        fig, axs = plt.subplots(1, 2, figsize=(8, 3.5))
        ax = sns.histplot(clone_size[clone_size > 0], ax=axs[0], **kwargs)
        ax.set_xlabel("Clone size")
        ax.set_ylabel("Count")

        ax = sns.histplot(clonal_bc_number[clonal_bc_number > 0], ax=axs[1], **kwargs)
        ax.set_xlabel("Clonal barcode number per cell")
        ax.set_ylabel("Count")
        fig.suptitle(f"Time={x}")
