from bioat.logger import set_logging_level
from bioat.lib.libcrispr import (
    TARGET_REGIONS_LIB,
    run_no_PAM_sgRNA_alignment_no_chop,
    run_sgRNA_alignment
)
from bioat.lib.libcolor import (
    convert_hex_to_rgb,
    convert_rgb_to_hex,
    map_color,
    make_color_list
)

from Bio.Align import PairwiseAligner
from Bio.Seq import Seq

from matplotlib.collections import PatchCollection
from matplotlib.patches import Rectangle
from matplotlib import pyplot as plt

import numpy as np
import pandas as pd
import matplotlib
import logging

matplotlib.use('Agg')
np.set_printoptions(suppress=True)


class TargetedDeepSequencing:
    def __init__(self):
        pass

    def region_heatmap(
            self,
            input_table: str,
            output_fig: str,
            target_region: str = None,  # sgRNA
            input_table_header: bool = True,
            output_fig_fmt: str = 'pdf',
            output_fig_dpi: int = 100,
            show_indel: bool = True,
            show_index: bool = True,
            box_border: bool = False,
            box_space: float = 0.03,
            min_color: tuple = (250, 239, 230),
            max_color: tuple = (154, 104, 57),
            min_ratio: float = 0.001,
            max_ratio: float = 0.99,
            region_extend_length: int = 25,
            local_alignment_scoring_matrix: tuple = (5, -4, -24, -8),
            local_alignment_min_score: int = 15,
            get_built_in_target_region: bool = False,
            log_level: str = 'INFO'
    ):
        """Plot region mutation info using table generated by `bioat bam mpileup_to_table`.

        :param input_table: Path of table generated by `bioat bam mpileup_to_table`. This table should contain base
        mutation information from a short genome region (no more than 1k nt).
        :param output_fig: Path of output figure.
        :param target_region: The program will attempt to align the sequence information specified by target_region to
        the reference sequence of mpileup.table. The target_region sequence can be sgRNA (if you need to specify the
        PAM sequence you can use this format, $NGG$.). When PAM is specified, local alignment takes into account PAM
        information. If PAM is not specified, normal local alignment is performed. Formal format of target_region should
        be like:
        'GAGTCCGAGCAGAAGAAGAA$GGG$', $GGG$ refers to PAM of SpCas9-BE on `EMX1` site;\n
        '$TTTA$GCCCCAATAATCCCCACATGTCA', $TTTA$ refers to PAM of cpf1-BE on `CDKN2A` site;\n
        'TGCTAGTAACCACGTTCTCCTGATCAAATATCACTCTCCTACTTACAGGA' refers to no PAM on `ND4` site.
        :param input_table_header: whether input_table has header or not
        :param output_fig_fmt: pdf | png
        :param output_fig_dpi: output figure dpi
        :param show_indel: whether show indel info in the out figure or not
        :param show_index: whether show index info in the out figure or not
        :param box_border: whether show box border info in the out figure or not
        :param box_space: space size between two boxes
        :param min_color: min color to plot heatmap with RGB format
        :param max_color: max color to plot heatmap with RGB format
        :param min_ratio: mutation ratio of box lower than min_ratio will show as white
        :param max_ratio: mutation ratio of box higher than min_ratio will show as ??
        :param region_extend_length: expand how many bps for the two side of this region
        :param local_alignment_scoring_matrix: set <align_match_score>  <align_mismatch_score>
                                                <align_gap_open_score> <align_gap_extension_score>
        :param local_alignment_min_score: If alignment score lower than this, consider as no appropriate alignment
        :param get_built_in_target_region: specify to True to return built-in target_region sequence.
        :param log_level: 'CRITICAL', 'ERROR', 'WARNING', 'INFO', 'DEBUG', 'NOTSET'

        :demo:
        `samtools mpileup test_sorted.bam --reference HK4-AOut-1.ref.upper.fa | gzip > test_sorted.mpileup.gz`
        `bioat bam mpileup_to_table test_sorted.mpileup.gz > test_sorted.mpileup.info.tsv`
        ```
        bioat targeted_deep_sequencing region_heatmap \
            test_sorted.mpileup.info.tsv \
            > test_sorted.mpileup.info.pdf
        ```
        """
        set_logging_level(level=log_level)

        if get_built_in_target_region:
            logging.info(
                f'You can use <key> in built-in target_regions to represent your target_region:\n' +
                '\t<key>\t<target_region>\n' +
                ''.join([f'\t{k}\t{v}\n' for k, v in TARGET_REGIONS_LIB.items()])
            )
            exit(0)
        else:
            # set sgRNA info
            if target_region in TARGET_REGIONS_LIB:
                target_region = TARGET_REGIONS_LIB[target_region]

        if target_region is None:
            PAM_seq = None
        elif '$' in target_region:
            if target_region.find('$') == 0:
                _, PAM_seq, target_region = target_region.split('$')
                target_region = target_region.upper()
            else:
                target_region, PAM_seq, _ = target_region.split('$')
                target_region = target_region.upper()
        else:
            pass
        logging.debug(
            f'parse <target_region>: target_region={target_region}, PAM_seq={PAM_seq}')
        # load .bmat file
        bmat_table = pd.read_csv(input_table, sep="\t") if input_table_header \
            else pd.read_csv(input_table, sep="\t", header=None)

        ref_seq = "".join(bmat_table.ref_base)

        # set alignment
        align_match, align_mismatch, align_gap_open, align_gap_extension = local_alignment_scoring_matrix

        seq_aligner = PairwiseAligner()
        seq_aligner.match = align_match
        seq_aligner.mismatch = align_mismatch
        seq_aligner.open_gap_score = align_gap_open
        seq_aligner.extend_gap_score = align_gap_extension
        seq_aligner.query_left_gap_score = 0
        seq_aligner.query_right_gap_score = 0
        seq_aligner.mode = "global"

        logging.debug("-" * 80)
        logging.debug(str(seq_aligner))
        logging.debug("-" * 80)

        # ---------------------------------------------------------------->>>>>
        # alignment
        # ---------------------------------------------------------------->>>>>
        # sgRNA_align_extend_len = 3

        # make ref seq
        # ref_seq_BioPy = Seq(ref_seq, Alphabet.IUPAC.unambiguous_dna)
        ref_seq_BioPy = Seq(ref_seq)
        ref_seq_rc = str(ref_seq_BioPy.reverse_complement())

        # fwd alignment
        final_align_fwd = run_no_PAM_sgRNA_alignment_no_chop(
            ref_seq, target_region, seq_aligner)

        # rev alignment
        final_align_rev = run_no_PAM_sgRNA_alignment_no_chop(
            ref_seq_rc, target_region, seq_aligner)

        # fwd alignment
        print("Forward best alignment:")
        print(final_align_fwd[0][6][final_align_fwd[0]
                                    [4]: final_align_fwd[0][5] + 1])
        print(final_align_fwd[0][7][final_align_fwd[0]
                                    [4]: final_align_fwd[0][5] + 1])
        print(final_align_fwd[0][8][final_align_fwd[0]
                                    [4]: final_align_fwd[0][5] + 1])
        print(final_align_fwd[0][3])
        print("-" * 80)

        # rev alignment
        print("Reverse best alignment:")
        print(final_align_rev[0][6][final_align_rev[0]
                                    [4]: final_align_rev[0][5] + 1])
        print(final_align_rev[0][7][final_align_rev[0]
                                    [4]: final_align_rev[0][5] + 1])
        print(final_align_rev[0][8][final_align_rev[0]
                                    [4]: final_align_rev[0][5] + 1])
        print(final_align_rev[0][3])
        print("-" * 80)

        # make alignment info
        sgRNA_align = [""] * len(ref_seq)
        sgRNA_align_insert = [""] * len(ref_seq)

        final_align_direction = None
        DNA_rev_cmp_dict = {
            "A": "T",
            "T": "A",
            "C": "G",
            "G": "C",
            "N": "N",
            "-": "-"}

        # define align direction and final align res
        if final_align_fwd[0][3] >= final_align_rev[0][3]:
            if final_align_fwd[0][3] >= local_alignment_min_score:
                final_align_direction = "Forward Alignment"
                final_align = final_align_fwd[0]

        elif final_align_rev[0][3] > final_align_fwd[0][3]:
            if final_align_rev[0][3] >= local_alignment_min_score:
                final_align_direction = "Reverse Alignment"
                final_align = final_align_rev[0]

        if final_align_direction is None:
            raise IOError("Alignment Error!")

        # make sgRNA alignment info
        final_align_info = final_align[7][final_align[4]: final_align[5] + 1]
        final_align_ref = final_align[6][final_align[4]: final_align[5] + 1]
        final_align_sgRNA = final_align[8][final_align[4]: final_align[5] + 1]

        final_align_ref_gap_count = final_align_ref.count("-")
        final_align_sgRNA_gap_count = final_align_sgRNA.count("-")

        ref_align_gap_count = 0
        ref_del_str = ""

        # select bmat file
        if (final_align_direction == "Forward Alignment") or (
                final_align_direction == "Reverse Alignment"):
            sgRNA_start = final_align[4]

            for align_index, align_ref in enumerate(final_align_ref):
                if align_ref != "-":
                    sgRNA_align[sgRNA_start + align_index - \
                                ref_align_gap_count] = ref_del_str + final_align_sgRNA[align_index]
                    ref_del_str = ""
                else:
                    ref_align_gap_count += 1
                    ref_del_str += final_align_sgRNA[align_index]
                    sgRNA_align_insert[sgRNA_start +
                                       align_index] = final_align_sgRNA[align_index]

            # add PAM info
            if final_align_direction == "Reverse Alignment":
                sgRNA_align = sgRNA_align[::-1]
                sgRNA_align_insert = sgRNA_align_insert[::-1]

        # set possible_sgRNA_region
        sgRNA_align_start = final_align[4]
        sgRNA_align_end = final_align[5]

        possible_sgRNA_region_start = max(
            sgRNA_align_start - region_extend_length, 0)
        possible_sgRNA_region_end = min(
            sgRNA_align_end + region_extend_length,
            len(ref_seq) - 1)
        possible_sgRNA_region = [
            possible_sgRNA_region_start,
            possible_sgRNA_region_end]

        # select bmat_table
        bmat_table_select = bmat_table[possible_sgRNA_region[0]: possible_sgRNA_region[1]]

        # ---------------------------------------------------------------->>>>>
        # make plot
        # ---------------------------------------------------------------->>>>>

        # --------------------------------------------------->>>>>
        # set color
        # --------------------------------------------------->>>>>
        # show indel
        indel_plot_state = show_indel
        index_plot_state = show_index
        box_border_plot_state = box_border

        # set panel size
        panel_box_width = 0.4
        panel_box_heigth = 0.4
        panel_space = 0.05
        panel_box_space = box_space

        # color part
        base_color_dict = {
            "A": "#04E3E3",
            "T": "#F9B874",
            "C": "#B9E76B",
            "G": "#F53798",
            "N": "#AAAAAA",
            "-": "#AAAAAA"}

        # make color breaks
        color_break_num = 20
        break_step = 1.0 / color_break_num
        min_color_value = min_ratio
        max_color_value = max_ratio
        color_break = np.round(
            np.arange(
                min_color_value,
                max_color_value,
                break_step),
            5)

        try:
            color_list = make_color_list(
                min_color, max_color, len(color_break) - 1, "Hex")
            color_list = ["#FFFFFF"] + color_list
        except BaseException:
            print
            min_color, max_color
            print
            color_break

        # get plot info
        total_box_count = possible_sgRNA_region[1] - possible_sgRNA_region[0]

        # calculate base info and fix zero
        base_sum_count = bmat_table_select[["A", "G", "C", "T"]].apply(
            lambda x: x.sum(), axis=1)
        total_sum_count = bmat_table_select[[
            "A", "G", "C", "T", "del_count", "insert_count"]].apply(lambda x: x.sum(), axis=1)
        base_sum_count[np.where(base_sum_count == 0)[0]] = 1
        total_sum_count[np.where(total_sum_count == 0)[0]] = 1

        # make plot size
        plot_data_list = None
        panel_space_coef = None
        panel_height_coef = None

        if indel_plot_state:
            panel_height_coef = [0.5, 0.9, 0.9] + [0.5] * 6 + [0.5] * 6
            panel_space_coef = [1, 1, 1] + [0.3] * \
                               3 + [1, 0.3, 1] + [0.3] * 3 + [1, 0.3]
            plot_data_list = [
                ["Index", bmat_table_select.chr_index],
                ["On-target", sgRNA_align[possible_sgRNA_region[0]: possible_sgRNA_region[1]]],
                ["Ref", bmat_table_select.ref_base],
                ["A", np.array(bmat_table_select["A"])],
                ["G", np.array(bmat_table_select["G"])],
                ["C", np.array(bmat_table_select["C"])],
                ["T", np.array(bmat_table_select["T"])],
                ["Del", np.array(bmat_table_select["del_count"])],
                ["Ins", np.array(bmat_table_select["insert_count"])],
                ["A.ratio", list(bmat_table_select["A"] / base_sum_count)],
                ["G.ratio", list(bmat_table_select["G"] / base_sum_count)],
                ["C.ratio", list(bmat_table_select["C"] / base_sum_count)],
                ["T.ratio", list(bmat_table_select["T"] / base_sum_count)],
                ["Del.ratio", list(bmat_table_select["del_count"] / total_sum_count)],
                ["Ins.ratio", list(bmat_table_select["insert_count"] / total_sum_count)]
            ]
        else:
            panel_height_coef = [0.5, 0.9, 0.9] + [0.5] * 4 + [0.5] * 4
            panel_space_coef = [1, 1, 1] + [0.3] * 3 + [1] + [0.3] * 3
            plot_data_list = [
                ["Index", bmat_table_select.chr_index],
                ["On-target", sgRNA_align[possible_sgRNA_region[0]: possible_sgRNA_region[1]]],
                ["Ref", bmat_table_select.ref_base],
                ["A", np.array(bmat_table_select["A"])],
                ["G", np.array(bmat_table_select["G"])],
                ["C", np.array(bmat_table_select["C"])],
                ["T", np.array(bmat_table_select["T"])],
                ["A.ratio", list(bmat_table_select["A"] / base_sum_count)],
                ["G.ratio", list(bmat_table_select["G"] / base_sum_count)],
                ["C.ratio", list(bmat_table_select["C"] / base_sum_count)],
                ["T.ratio", list(bmat_table_select["T"] / base_sum_count)]
            ]

        # get box and space info
        box_height_list = np.array(panel_height_coef) * panel_box_heigth
        panel_space_list = np.array(panel_space_coef) * panel_space

        # calculate figure total width and height
        figure_width = total_box_count * panel_box_width + \
                       (total_box_count - 1) * panel_box_space + panel_box_width * 2
        figure_height = sum(box_height_list) + sum(panel_space_list)

        # make all box_x
        box_x_vec = np.arange(
            0,
            figure_width +
            panel_box_width,
            panel_box_width +
            panel_box_space)
        box_x_vec = box_x_vec[:(len(ref_seq) + 1)]

        # make box border
        if box_border_plot_state:
            box_edgecolor = "#AAAAAA"
            box_linestyle = "-"
            box_linewidth = 2
        else:
            box_edgecolor = "#FFFFFF"
            box_linestyle = "None"
            box_linewidth = 0

        # make box_y initialize
        current_y = 0

        # plot region
        # set new figure
        fig = plt.figure(figsize=(figure_width * 1.1, figure_height * 1.1))
        ax = fig.add_subplot(111, aspect="equal")
        plt.xlim([0, figure_width])
        plt.ylim([-figure_height, 0])
        plt.axis("off")

        # make plot
        text_list = []
        patches = []

        for panel_index in range(len(panel_height_coef)):

            # panel name
            panel_name = plot_data_list[panel_index][0]
            panel_name_x = box_x_vec[0]
            panel_name_y = current_y - box_height_list[panel_index] * 0.5
            text_list.append((panel_name_x, panel_name_y, panel_name, 10))

            # plot panel box
            if panel_name == "Index":
                # don't draw box, only add text
                for index, box_value in enumerate(
                        plot_data_list[panel_index][1]):
                    box_x = box_x_vec[index + 1]
                    text_list.append(
                        (box_x +
                         panel_box_width *
                         0.5,
                         current_y -
                         box_height_list[panel_index] *
                         0.5,
                         str(box_value),
                         7))

                # make next panel_y
                current_y = current_y - \
                            (box_height_list[panel_index] + panel_space_list[panel_index])

            elif panel_name in ["On-target", "Ref"]:
                for index, box_value in enumerate(
                        plot_data_list[panel_index][1]):
                    if box_value == "":
                        box_fill = False
                        box_color = "#FFFFFF"

                    else:
                        if "Reverse" in final_align_direction:
                            if panel_name == "Ref":
                                box_value = "".join(
                                    [DNA_rev_cmp_dict.get(x) for x in box_value])
                            else:
                                pass

                            box_color = base_color_dict.get(box_value[0])

                        else:
                            box_color = base_color_dict.get(box_value[-1])

                        if not box_color:
                            box_fill = False
                            box_color = "#FFFFFF"
                        else:
                            box_fill = True

                    box_x = box_x_vec[index + 1]

                    patches.append(Rectangle(
                        xy=(box_x, current_y - box_height_list[panel_index]),
                        width=panel_box_width,
                        height=box_height_list[panel_index],
                        fill=box_fill,
                        alpha=1,
                        linestyle=box_linestyle,
                        linewidth=box_linewidth,
                        edgecolor=box_edgecolor,
                        facecolor=box_color)
                    )

                    # text
                    text_list.append(
                        (box_x +
                         0.5 *
                         panel_box_width,
                         current_y -
                         0.5 *
                         box_height_list[panel_index],
                         str(box_value),
                         16))

                # make next panel_y
                current_y = current_y - \
                            (box_height_list[panel_index] + panel_space_list[panel_index])

            elif panel_name in ["A", "G", "C", "T", "Del", "Ins"]:
                if panel_name in ["Del", "Ins"]:
                    box_ratio = plot_data_list[panel_index][1] / \
                                total_sum_count
                else:
                    box_ratio = plot_data_list[panel_index][1] / base_sum_count

                box_color_list = map_color(box_ratio, color_break, color_list)
                for index, box_value in enumerate(
                        plot_data_list[panel_index][1]):
                    box_color = box_color_list[index]
                    box_x = box_x_vec[index + 1]

                    patches.append(Rectangle(
                        xy=(box_x, current_y - box_height_list[panel_index]),
                        width=panel_box_width,
                        height=box_height_list[panel_index],
                        fill=True,
                        alpha=1,
                        linestyle=box_linestyle,
                        linewidth=box_linewidth,
                        edgecolor=box_edgecolor,
                        facecolor=box_color)
                    )

                    # text
                    text_list.append(
                        (box_x +
                         0.5 *
                         panel_box_width,
                         current_y -
                         0.5 *
                         box_height_list[panel_index],
                         str(box_value),
                         6))

                # make next panel_y
                current_y = current_y - \
                            (box_height_list[panel_index] + panel_space_list[panel_index])

            else:
                box_color_list = map_color(
                    plot_data_list[panel_index][1], color_break, color_list)
                for index, box_value in enumerate(
                        plot_data_list[panel_index][1]):
                    box_color = box_color_list[index]
                    box_x = box_x_vec[index + 1]

                    patches.append(Rectangle(
                        xy=(box_x, current_y - box_height_list[panel_index]),
                        width=panel_box_width,
                        height=box_height_list[panel_index],
                        fill=True,
                        alpha=1,
                        linestyle=box_linestyle,
                        linewidth=box_linewidth,
                        edgecolor=box_edgecolor,
                        facecolor=box_color)
                    )

                    # text
                    text_list.append(
                        (box_x +
                         0.5 *
                         panel_box_width,
                         current_y -
                         0.5 *
                         box_height_list[panel_index],
                         round(
                             box_value *
                             100,
                             4),
                         6))

                # make next panel_y
                if panel_index < len(panel_space_list):
                    current_y = current_y - \
                                (box_height_list[panel_index] + panel_space_list[panel_index])

                    # plot box
        ax.add_collection(PatchCollection(patches, match_original=True))

        # add text
        for text_x, text_y, text_info, text_fontsize in text_list:
            plt.text(
                x=text_x,
                y=text_y,
                s=text_info,
                horizontalalignment='center',
                verticalalignment='center',
                fontsize=text_fontsize,
                fontname="Arial"
            )

        # output plot
        fig.savefig(
            fname=output_fig,
            bbox_inches='tight',
            dpi=output_fig_dpi,
            format=output_fig_fmt)

    # def region_heatmap_merge(self):
    #     pass


if __name__ == '__main__':
    TargetedDeepSequencing.region_heatmap()
    pass
