import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import argparse
from . import utilities


def plot_energy_landscape(
    states_file,  # as generated by RRIkinDP
    figure_path=None,  # output file; if None a default path is set
    energy="E",  # free energy to plot: E, Ehybrid, ED1, ED2 or ED
    structure=None,  # interaction as string as returned by RRIkinDP.utilities.get_string_representations()
    annotate=None,  # plot energy values per structure; True/False; default: True if less than 16 bps in interaction
    e_min=None,  # min value represented by color scale (kcal/mol)
    e_max=None,  # max value represented by color scale (kcal/mol)
    figsize=(9, 7),  # figure size in inches
    seqId1="Seq1",  # name of first RNA
    seqId2="Seq2",  # name of second RNA
):
    """Plot energy landscape from RRIkinDP states output file as heatmap."""

    # read states file (zero based indices)
    df = pd.read_csv(states_file, sep="\t")

    # convert energies from dekacalories/mol to kcal/mol and add initiation correction
    df[energy] = df.apply(lambda row: row[energy] / 100, axis=1)

    # set color scale limits
    if e_min is None:
        e_min = df[energy].min()
    if e_max is None:
        e_max = df[energy].max()

    # states list to energy landscape
    energy_df = df.pivot(index="k", columns="l", values=energy)

    # only annotate energies if the interaction has fewer than 15 base pairs
    if annotate is None:
        if len(energy_df) > 15:
            annotate = False
        else:
            annotate = True

    # set up plot
    f, ax = plt.subplots(figsize=figsize)

    # generate heatmap
    ax = sns.heatmap(
        energy_df,
        vmin=e_min,
        vmax=e_max,
        center=0,
        cmap="RdBu_r",
        linewidths=0.2,
        linecolor="lightgrey",
        annot=annotate,
        fmt="2.1f",
        cbar_kws={
            "label": f"free energy term: {energy} in kcal/mol",
            "pad": 0.1,
            "shrink": 1.0,
        },
        square=True,
        ax=ax,
        # cbar=False,
    )

    # set title
    plt.title(f"Energy Landscape: {seqId1} + {seqId2}")

    # set axis lables
    ax.set_ylabel("start base pair")
    ax.set_xlabel("end base pair")

    # draw axis ticks on all four sides
    plt.tick_params(labeltop=True, labelright=True, top=True, right=True)

    # include structure information in plot
    if structure is not None:
        # plot full interaction structure
        props = dict(facecolor="white", alpha=0.7)  # boxstyle='round')
        ax.text(
            0.05,
            0.15,
            structure,
            transform=ax.transAxes,
            fontsize=8,
            verticalalignment="top",
            bbox=props,
            fontfamily="monospace",
        )

        # label axis with base pairs in addition to one based indices
        struct = structure.split("\n")
        labels = [
            [struct[0][e], struct[2][e]] for e, c in enumerate(struct[1]) if c == "|"
        ]
        xticklabels = [f"{i+1}\n" + "\nI\n".join(l) for i, l in enumerate(labels)]
        yticklabels = [f"{i+1} " + " ─ ".join(l) for i, l in enumerate(labels)]
        ax.set_xticks([t + 0.5 for t in range(0, len(energy_df))])
        ax.set_yticks([t + 0.5 for t in range(0, len(energy_df))])
        ax.set_xlim(0, len(energy_df))
        ax.set_ylim(len(energy_df), 0)
        ax.set_xticklabels(
            xticklabels,
            fontdict={
                "fontsize": 8,
                "horizontalalignment": "center",
            },
        )
        ax.set_yticklabels(
            yticklabels,
            rotation="horizontal",
            fontsize=8,
        )
    # if no structure information: change zero based to one based indices
    else:
        yticklabels = [
            int(item.get_text()) + 1
            for item in ax.get_yticklabels()
            if item.get_position()[0] == 0
        ]
        ax.set_yticklabels(yticklabels, rotation="horizontal")
        xticklabels = [
            int(item.get_text()) + 1
            for item in ax.get_xticklabels()
            if item.get_position()[1] == 0
        ]
        ax.set_xticklabels(xticklabels)

    # write plot to file
    if figure_path is not None:
        f.savefig(figure_path, bbox_inches="tight")
    else:
        figure_path = f"{seqId1}_{seqId2}_{energy}.pdf"
        f.savefig(figure_path)

    # close plot to avoid to many open plots if function is called multiple times
    plt.close(f)


if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        description="Plot energylandscape from states provided by RRIkinDP."
    )

    parser.add_argument("states", help="path to states file from RRIkinDP", type=str)

    parser.add_argument(
        "-o", "--out", help="path to output figure", type=str, default=None
    )

    parser.add_argument("--id1", help="name of first RNA", type=str, default="Seq1")
    parser.add_argument("--id2", help="name of second RNA", type=str, default="Seq2")

    parser.add_argument(
        "--energy",
        help="free energy to plot: E, ED1, ED2 or Ehybrid",
        type=str,
        default="E",
    )

    parser.add_argument(
        "--seq1",
        help="sequence of first RNA",
        type=str,
    )
    parser.add_argument(
        "--seq2",
        help="sequence of second RNA",
        type=str,
    )

    parser.add_argument(
        "--bp_list",
        help="base pair list as string; example: (134,56):(135,55):(136,54):(137,53):(138,52):(139,51)",
        type=str,
    )

    parser.add_argument(
        "--e_min", help="path to output figure", type=float, default=None
    )
    parser.add_argument(
        "--e_max", help="path to output figure", type=float, default=None
    )

    parser.add_argument(
        "--figuresize",
        help="figure dimensions in inches; provide two arguments seperated by space; width hight",
        nargs=2,
        type=float,
        default=[8, 8],
    )

    annotate_group = parser.add_mutually_exclusive_group()
    annotate_group.add_argument(
        "--annotate",
        help="annotate energy per interaction structure; default: do not annotate if more than 15 base pairs in interaction",
        action="store_true",
    )
    annotate_group.add_argument(
        "--remove_annotations",
        help="do not annotate energy per interaction structure; default: annotate if less than 16 base pairs in interaction",
        action="store_true",
    )

    args = parser.parse_args()
    states_file = args.states

    figure_path = args.out

    seqId1 = args.id1
    seqId2 = args.id2

    seq1 = args.seq1
    seq2 = args.seq2
    bp_list_str = args.bp_list

    # only plot struture annotation if both sequences and the base pair list are provided
    structure = None
    if seq1 and seq2 and bp_list_str:
        bp_list = utilities.intarna_to_bplist(bp_list_str)
        structure = utilities.get_string_representations(
            seq1, seq2, bp_list, id1=seqId1, id2=seqId2
        )

    if args.energy in ["E", "ED1", "ED2", "ED", "Ehybrid"]:
        energy = args.energy
    else:
        raise SystemExit(
            f"error: argument --energy: {args.energy} invalide value; use 'E', 'ED1', 'ED2', 'ED' or 'Ehybrid'"
        )

    width, height = args.figuresize
    figsize = (width, height)

    e_min = args.e_min
    e_max = args.e_max

    annotate = None
    if args.annotate:
        annotate = True
    if args.remove_annotations:
        annotate = False

    plot_energy_landscape(
        states_file,
        figure_path=figure_path,
        energy=energy,
        structure=structure,
        annotate=annotate,
        e_min=e_min,
        e_max=e_max,
        figsize=figsize,
        seqId1=seqId1,
        seqId2=seqId2,
    )

# python landscape.py /home/maria/Work/Projects/RRI/data/intarna_directpaths_treekin/richter2012shuffled4/b1737_ChiX/dp_states_b1737_ChiX_maxlen-42.csv -o 'test.pdf' --energy E --seq1 'GUUUGUUACCCAACAAACCGGUUGAAGUAAUUGACUCGCUGCUUUAUGGCAAAGUCGAUGGUUUAGGCGUGCUUAAGGCUGCGGUUGCAGCGAUUAAAAAAGCCGCAGCAAAUUAAUUUAUUUUAAAUUUUCCCGUCAAAGAGUUAUUUCAUAAAUCAAUACCGCAAUAUUUAAAUUGCGGUUUUUAAGGGUAUUUUUCUAUGAGUAAUGUUAUUGCAUCGCUUGAAAAGGUACUCCUCCCUUUUGCAGUUAAAAUAGGAAAGCAGCCACACGUUAAUGCAAUCAAAAAUGGCUUUAUUC' --seq2 acaccgucgcuuaaagugacggcauaauaauaaaaaaaugaaauuccucuuugacgggccaauagcgauauuggccauuuuuuu --bp_list '(134,56):(135,55):(136,54):(137,53):(138,52):(139,51):(140,50):(141,49):(142,48):(143,47):(146,44):(147,43):(148,42):(149,41):(150,40):(151,39):(152,38)'
