import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.colors as mcolors
import matplotlib.gridspec as gridspec
from Bio import Phylo
from .treeFromClusters import new_phylogeny

def show_hor_in_seq(
    hor_in_seq,
    reference_seq,
    ax_seq,
    seq_fg_color='black',
    seq_bg_color='#D3D3D3'
):
    bp_seq_length = len(reference_seq)

    ax_seq.set_xlim([0, bp_seq_length])
    ax_seq.get_yaxis().set_visible(False)

    # spans = hor_in_seq.spans_in_seq
    locations = hor_in_seq.locations

    # bp_start = spans[0].span_start * monomer_size
    # bp_end = (spans[-1].span_start + spans[-1].span_length) * monomer_size

    ax_seq.add_patch(patches.Rectangle(
        (0, 0), bp_seq_length, 1, facecolor=seq_bg_color))
    for location in locations:
        ax_seq.add_patch(patches.Rectangle(
            (location.start, 0 if location.strand == -1 else 0.5),
            location.end - location.start,
            0.5,
            facecolor=seq_fg_color))
    


def show_hor(
    hor_in_seq,
    hor_index,
    reference_seq,
    parent_gridspec,
    seq_fg_color='black',
    seq_bg_color='#D3D3D3',
    hor_item_size_ratio=(0.04, 0.5),
    h_space_in_hor=0.02
):
        clade_seq = hor_in_seq.hor.clade_seq
        hor_size_ratio = hor_item_size_ratio[0]*len(clade_seq)

        gs_hor = gridspec.GridSpecFromSubplotSpec(
            2, 2,
            subplot_spec=parent_gridspec[hor_index, 0],
            width_ratios=[
                hor_size_ratio, 1 - hor_size_ratio
            ],
            height_ratios=[1, 2],
            hspace=h_space_in_hor
        )

        # gs_hor = ax_hor_and_seq.add_gridspec(
        #     2, 2,
        #     width_ratios=[hor_size_ratio,1-hor_size_ratio],
        #     height_ratios=[1,1],
        #     hspace=h_space
        # )

        ax_hor = plt.subplot(gs_hor[0, 0])
        ax_seq = plt.subplot(gs_hor[1, :])

        ax_hor.set_xlim([0, len(clade_seq)])
        ax_hor.grid(True)
        ax_hor.set_xticks(np.arange(0, len(clade_seq) + 1, 1))
        ax_hor.get_yaxis().set_visible(False)
        for clade_pos, clade in enumerate(clade_seq):
            ax_hor.add_patch(patches.Rectangle(
                (clade_pos, 0), 1, 1, facecolor=clade.color.to_hex()))
            
        show_hor_in_seq(
            hor_in_seq=hor_in_seq,
            reference_seq=reference_seq,
            ax_seq=ax_seq,
            seq_fg_color=seq_fg_color,
            seq_bg_color=seq_bg_color
        )


def show_hors(hors_in_seq, reference_seq,
              tree=None,
              label='unnamed',
              color_palette=[mcolors.TABLEAU_COLORS[color_id]
                             for color_id in mcolors.TABLEAU_COLORS],
              seq_fg_color='black',
              seq_bg_color='#D3D3D3',
              tree_bg_color='#D3D3D3',
              hor_item_size_ratio=(0.04, 0.5),
              seq_height_ratio=0.5,
              fig_width=7,
              tree_height=5,
              hor_height=1,
              h_space=0.1,
              h_space_in_hor=0.02,
              label_func=(lambda clade_name: None)):

    clade_set = set(
        [clade for hor_in_seq in hors_in_seq for clade in hor_in_seq.hor.clade_seq])
    for clade_index, clade in enumerate(clade_set):
        # this is bad, but at least avoids error
        clade.color = color_palette[clade_index % 10]

    hors_height = hor_height * len(hors_in_seq)
    fig = plt.figure(constrained_layout=True, figsize=(
        fig_width, tree_height + hors_height))
    # fig.suptitle(f'HOR {label} ({bp_start}-{bp_end})')

    # print([hors_height for i in range(len(hors_in_seq))].append(tree_height))

    gs = fig.add_gridspec(
        len(hors_in_seq) + 1, 1,
        # width_ratios=[hor_size_ratio,1-hor_size_ratio],
        height_ratios=[hor_height for i in range(
            len(hors_in_seq))] + [tree_height],
        hspace=h_space
    )
    # ax_hors = [plt.subplot(gs[i, :]) for i in range(len(hors_in_seq))]
    ax_tree = plt.subplot(gs[len(hors_in_seq), :])

    for hor_index in range(len(hors_in_seq)):
        # ax_hor_and_seq = plt.subplot(gs[hor_index, 0])
        show_hor(
            hor_in_seq=hors_in_seq[hor_index],
            hor_index=hor_index,
            reference_seq=reference_seq,
            parent_gridspec=gs,
            seq_fg_color=seq_fg_color,
            seq_bg_color=seq_bg_color,
            hor_item_size_ratio=hor_item_size_ratio,
            h_space_in_hor=h_space_in_hor
        )

    ax_tree.get_yaxis().set_visible(False)
    ax_tree.get_xaxis().set_visible(False)
    # ax_tree.spines['top'].set_visible(False)
    # ax_tree.spines['right'].set_visible(False)
    # ax_tree.spines['bottom'].set_visible(False)
    # ax_tree.spines['left'].set_visible(False)
    ax_tree.get_yaxis().set_ticks([])
    # ax_tree.axis('off')
    if tree is not None:
        tree.root.color = tree_bg_color
        Phylo.draw(tree, axes=ax_tree, label_func=label_func)
        # Phylo.draw_graphviz(tree, axes=ax_tree)

    for clade in clade_set:
        clade.color = None
    if tree is not None:
        tree.root.color = None


def show_hor_tree(hor_tree_root, reference_seq, tree, path=[], level=0):
    if not hor_tree_root.sub_hors:
        return
    print(f'Subtree: {path}')
    show_hors(hor_tree_root.sub_hors, reference_seq, tree=tree)
    for branch_index, sub_hor in enumerate(hor_tree_root.sub_hors):
        show_hor_tree(sub_hor, reference_seq, tree=new_phylogeny(
            tree.common_ancestor(sub_hor.hor.clade_seq)), path=path + [branch_index + 1])
