from tqdm import tqdm

import seaborn as sns

import pandas as pd

import os

import matplotlib.pyplot as plt

import numpy as np

from bpnet.plot.utils import simple_yaxis_format, strip_axis, spaced_xticks

from bpnet.modisco.utils import bootstrap_mean, nan_like, ic_scale

from bpnet.plot.utils import show_figure

# TODO - make it as a bar-plot with two standard colors:

# #B23F49 (pos), #045CA8 (neg)

def plot_stranded_profile(profile, ax=None, ymax=None, profile_std=None, flip_neg=True, set_ylim=True):

    """Plot the stranded profile

    """

    if ax is None:

        ax = plt.gca()

    if profile.ndim == 1:

        # also compatible with single dim

        profile = profile[:, np.newaxis]

    assert profile.ndim == 2

    assert profile.shape[1] <= 2

    labels = ['pos', 'neg']

    # determine ymax if not specified

    if ymax is None:

        if profile_std is not None:

            ymax = (profile.max() - 2 * profile_std).max()

        else:

            ymax = profile.max()

    if set_ylim:

        if flip_neg:

            ax.set_ylim([-ymax, ymax])

        else:

            ax.set_ylim([0, ymax])

    ax.axhline(y=0, linewidth=1, linestyle='--', color='black')

    # strip_axis(ax)

    xvec = np.arange(1, len(profile) + 1)

    for i in range(profile.shape[1]):

        sign = 1 if not flip_neg or i == 0 else -1

        ax.plot(xvec, sign * profile[:, i], label=labels[i])

        # plot also the ribbons

        if profile_std is not None:

            ax.fill_between(xvec,

                            sign * profile[:, i] - 2 * profile_std[:, i],

                            sign * profile[:, i] + 2 * profile_std[:, i],

                            alpha=0.1)

    # return ax

def multiple_plot_stranded_profile(d_profile, figsize_tmpl=(4, 3), normalize=False):

    fig, axes = plt.subplots(1, len(d_profile),

                             figsize=(figsize_tmpl[0] * len(d_profile), figsize_tmpl[1]),

                             sharey=True)

    if len(d_profile)==1: #If only one task, then can't zip axes

        ax = axes

        task = [*d_profile][0]

        arr = d_profile[task].mean(axis=0)

        if normalize:

            arr = arr / arr.max()

        plot_stranded_profile(arr, ax=ax, set_ylim=False)

        ax.set_title(task)

        ax.set_ylabel("Avg. counts")

        ax.set_xlabel("Position")

        fig.subplots_adjust(wspace=0)  # no space between plots

        return fig

    else:

        for i, (task, ax) in enumerate(zip(d_profile, axes)):

            arr = d_profile[task].mean(axis=0)

            if normalize:

                arr = arr / arr.max()

            plot_stranded_profile(arr, ax=ax, set_ylim=False)

            ax.set_title(task)

            if i == 0:

                ax.set_ylabel("Avg. counts")

                ax.set_xlabel("Position")

        fig.subplots_adjust(wspace=0)  # no space between plots

        return fig

def aggregate_profiles(profile_arr, n_bootstrap=None, only_idx=None):

    if only_idx is not None:

        return profile_arr[only_idx], None

    if n_bootstrap is not None:

        return bootstrap_mean(profile_arr, n=n_bootstrap)

    else:

        return profile_arr.mean(axis=0), None

def extract_signal(x, seqlets, rc_fn=lambda x: x[::-1, ::-1]):

    def optional_rc(x, is_rc):

        if is_rc:

            return rc_fn(x)

        else:

            return x

    return np.stack([optional_rc(x[s['example'], s['start']:s['end']], s['rc'])

                     for s in seqlets])

def plot_profiles(seqlets_by_pattern,

                  x,

                  tracks,

                  contribution_scores={},

                  figsize=(20, 2),

                  start_vec=None,

                  width=20,

                  legend=True,

                  rotate_y=90,

                  seq_height=1,

                  ymax=None,  # determine y-max

                  n_limit=35,

                  n_bootstrap=None,

                  flip_neg=False,

                  patterns=None,

                  fpath_template=None,

                  only_idx=None,

                  mkdir=False,

                  rc_fn=lambda x: x[::-1, ::-1]):

    """

    Plot the sequence profiles

    Args:

      x: one-hot-encoded sequence

      tracks: dictionary of profile tracks

      contribution_scores: optional dictionary of contribution scores

    """

    import matplotlib.pyplot as plt

    from concise.utils.plot import seqlogo_fig, seqlogo

    # Setup start-vec

    if start_vec is not None:

        if not isinstance(start_vec, list):

            start_vec = [start_vec] * len(patterns)

    else:

        start_vec = [0] * len(patterns)

        width = len(x)

    if patterns is None:

        patterns = list(seqlets_by_pattern)

    # aggregated profiles

    d_signal_patterns = {pattern:

                         {k: aggregate_profiles(

                             extract_signal(y, seqlets_by_pattern[pattern])[:, start_vec[ip]:(start_vec[ip] + width)],

                             n_bootstrap=n_bootstrap, only_idx=only_idx)

                          for k, y in tracks.items()}

                         for ip, pattern in enumerate(patterns)}

    if ymax is None:

        # infer ymax

        def take_max(x, dx):

            if dx is None:

                return x.max()

            else:

                # HACK - hard-coded 2

                return (x + 2 * dx).max()

        ymax = [max([take_max(*d_signal_patterns[pattern][k])

                     for pattern in patterns])

                for k in tracks]  # loop through all the tracks

    if not isinstance(ymax, list):

        ymax = [ymax] * len(tracks)

    figs = []

    for i, pattern in enumerate(tqdm(patterns)):

        j = i

        # --------------

        # extract signal

        seqs = extract_signal(x, seqlets_by_pattern[pattern])[:, start_vec[i]:(start_vec[i] + width)]

        ext_contribution_scores = {s: extract_signal(contrib, seqlets_by_pattern[pattern])[:, start_vec[i]:(start_vec[i] + width)]

                                   for s, contrib in contribution_scores.items()}

        d_signal = d_signal_patterns[pattern]

        # --------------

        if only_idx is None:

            sequence = ic_scale(seqs.mean(axis=0))

        else:

            sequence = seqs[only_idx]

        n = len(seqs)

        if n < n_limit:

            continue

        fig, ax = plt.subplots(1 + len(contribution_scores) + len(tracks),

                               1, sharex=True,

                               figsize=figsize,

                               gridspec_kw={'height_ratios': [1] * len(tracks) + [seq_height] * (1 + len(contribution_scores))})

        # signal

        ax[0].set_title(f"{pattern} ({n})")

        for i, (k, signal) in enumerate(d_signal.items()):

            signal_mean, signal_std = d_signal_patterns[pattern][k]

            plot_stranded_profile(signal_mean, ax=ax[i], ymax=ymax[i],

                                  profile_std=signal_std, flip_neg=flip_neg)

            simple_yaxis_format(ax[i])

            strip_axis(ax[i])

            ax[i].set_ylabel(f"{k}", rotation=rotate_y, ha='right', labelpad=5)

            if legend:

                ax[i].legend()

        # -----------

        # contribution scores (seqlogo)

        # -----------

        # average the contribution scores

        if only_idx is None:

            norm_contribution_scores = {k: v.mean(axis=0)

                                        for k, v in ext_contribution_scores.items()}

        else:

            norm_contribution_scores = {k: v[only_idx]

                                        for k, v in ext_contribution_scores.items()}

        max_scale = max([np.maximum(v, 0).sum(axis=-1).max() for v in norm_contribution_scores.values()])

        min_scale = min([np.minimum(v, 0).sum(axis=-1).min() for v in norm_contribution_scores.values()])

        for k, (contrib_score_name, logo) in enumerate(norm_contribution_scores.items()):

            ax_id = len(tracks) + k

            # Trim the pattern if necessary

            # plot

            ax[ax_id].set_ylim([min_scale, max_scale])

            ax[ax_id].axhline(y=0, linewidth=1, linestyle='--', color='grey')

            seqlogo(logo, ax=ax[ax_id])

            # style

            simple_yaxis_format(ax[ax_id])

            strip_axis(ax[ax_id])

            # ax[ax_id].set_ylabel(contrib_score_name)

            ax[ax_id].set_ylabel(contrib_score_name, rotation=rotate_y, ha='right', labelpad=5)  # va='bottom',

        # -----------

        # information content (seqlogo)

        # -----------

        # plot

        seqlogo(sequence, ax=ax[-1])

        # style

        simple_yaxis_format(ax[-1])

        strip_axis(ax[-1])

        ax[-1].set_ylabel("Inf. content", rotation=rotate_y, ha='right', labelpad=5)

        ax[-1].set_xticks(list(range(0, len(sequence) + 1, 5)))

        figs.append(fig)

        # save to file

        if fpath_template is not None:

            pname = pattern.replace("/", ".")

            basepath = fpath_template.format(pname=pname, pattern=pattern)

            if mkdir:

                os.makedirs(os.path.dirname(basepath), exist_ok=True)

            plt.savefig(basepath + '.png', dpi=600)

            plt.savefig(basepath + '.pdf', dpi=600)

            plt.close(fig)    # close the figure

            show_figure(fig)

            plt.show()

    return figs

def plot_profiles_single(seqlet,

                         x,

                         tracks,

                         contribution_scores={},

                         figsize=(20, 2),

                         legend=True,

                         rotate_y=90,

                         seq_height=1,

                         flip_neg=False,

                         rc_fn=lambda x: x[::-1, ::-1]):

    """

    Plot the sequence profiles

    Args:

      x: one-hot-encoded sequence

      tracks: dictionary of profile tracks

      contribution_scores: optional dictionary of contribution scores

    """

    import matplotlib.pyplot as plt

    from concise.utils.plot import seqlogo_fig, seqlogo

    # --------------

    # extract signal

    seq = seqlet.extract(x)

    ext_contribution_scores = {s: seqlet.extract(contrib) for s, contrib in contribution_scores.items()}

    fig, ax = plt.subplots(1 + len(contribution_scores) + len(tracks),

                           1, sharex=True,

                           figsize=figsize,

                           gridspec_kw={'height_ratios': [1] * len(tracks) + [seq_height] * (1 + len(contribution_scores))})

    # signal

    for i, (k, signal) in enumerate(tracks.items()):

        plot_stranded_profile(seqlet.extract(signal), ax=ax[i],

                              flip_neg=flip_neg)

        simple_yaxis_format(ax[i])

        strip_axis(ax[i])

        ax[i].set_ylabel(f"{k}", rotation=rotate_y, ha='right', labelpad=5)

        if legend:

            ax[i].legend()

    # -----------

    # contribution scores (seqlogo)

    # -----------

    max_scale = max([np.maximum(v, 0).sum(axis=-1).max() for v in ext_contribution_scores.values()])

    min_scale = min([np.minimum(v, 0).sum(axis=-1).min() for v in ext_contribution_scores.values()])

    for k, (contrib_score_name, logo) in enumerate(ext_contribution_scores.items()):

        ax_id = len(tracks) + k

        # plot

        ax[ax_id].set_ylim([min_scale, max_scale])

        ax[ax_id].axhline(y=0, linewidth=1, linestyle='--', color='grey')

        seqlogo(logo, ax=ax[ax_id])

        # style

        simple_yaxis_format(ax[ax_id])

        strip_axis(ax[ax_id])

        # ax[ax_id].set_ylabel(contrib_score_name)

        ax[ax_id].set_ylabel(contrib_score_name, rotation=rotate_y, ha='right', labelpad=5)  # va='bottom',

    # -----------

    # information content (seqlogo)

    # -----------

    # plot

    seqlogo(seq, ax=ax[-1])

    # style

    simple_yaxis_format(ax[-1])

    strip_axis(ax[-1])

    ax[-1].set_ylabel("Inf. content", rotation=rotate_y, ha='right', labelpad=5)

    ax[-1].set_xticks(list(range(0, len(seq) + 1, 5)))

    return fig

def hist_position(dfp, tasks):

    """Make the positional histogram

    Args:

      dfp: pd.DataFrame with columns: peak_id, and center

      tasks: list of tasks for which to plot the different peak_id columns

    """

    fig, axes = plt.subplots(1, len(tasks), figsize=(5 * len(tasks), 2),

                             sharey=True, sharex=True)

    if len(tasks) == 1:

        axes = [axes]

    for i, (task, ax) in enumerate(zip(tasks, axes)):

        ax.hist(dfp[dfp.peak_id == task].center, bins=100)

        ax.set_title(task)

        ax.set_xlabel("Position")

        ax.set_xlim([0, 1000])

        if i == 0:

            ax.set_ylabel("Frequency")

    plt.subplots_adjust(wspace=0)

    return fig

def bar_seqlets_per_example(dfp, tasks):

    """Make the positional histogram

    Args:

      dfp: pd.DataFrame with columns: peak_id, and center

      tasks: list of tasks for which to plot the different peak_id columns

    """

    fig, axes = plt.subplots(1, len(tasks), figsize=(5 * len(tasks), 2),

                             sharey=True, sharex=True)

    if len(tasks) == 1:

        axes = [axes]

    for i, (task, ax) in enumerate(zip(tasks, axes)):

        dfpp = dfp[dfp.peak_id == task]

        ax.set_title(task)

        ax.set_xlabel("Frequency")

        if i == 0:

            ax.set_ylabel("# per example")

        if not len(dfpp):

            continue

        dfpp.groupby("example_idx").\

            size().value_counts().plot(kind="barh", ax=ax)

    plt.subplots_adjust(wspace=0)

    return fig

def box_counts(total_counts, pattern_idx):

    """Make a box-plot with total counts in the region

    Args:

      total_counts: dict per task

      pattern_idx: array with example_idx of the pattern

    """

    dfs = pd.concat([total_counts.melt().assign(subset="all peaks"),

                     total_counts.iloc[pattern_idx].melt().assign(subset="contains pattern")])

    dfs.value = np.log10(1 + dfs.value)

    fig, ax = plt.subplots(figsize=(5, 5))

    sns.boxplot("variable", "value", hue="subset", data=dfs, ax=ax)

    ax.set_xlabel("Task")

    ax.set_ylabel("log10(1+counts)")

    ax.set_title("Total number of counts in the region")

    return fig