import warnings
import streamlit as st
from streamlit import session_state as st_state
from streamlit_option_menu import option_menu
from pyfurnace.design.motifs import kissing_loops
from .motif_command import MotifCommand
from .. import second_menu_style
from ..motifs_icons import MOTIF_ICONS

kl_name_map = {
    "Dimer": "KissingDimer",
    "180°": "KissingLoop180",
    "Dimer 120°": "KissingDimer120",
    "120°": "KissingLoop120",
    "Branched Dimer": "BranchedDimer",
    "Branched KL": "BranchedKissingLoop",
}


class KissingLoopsCommand(MotifCommand):

    def execute(self, motif=None):
        warnings.filterwarnings("ignore")  # ignore kl energy warning
        ### Modify motif
        if motif:
            # interface data
            # flip, seq, energy, tolerance, pk_index
            data = self.interface(
                "mod",
                motif.get_kissing_sequence(),
                motif.energy,
                motif.energy_tolerance,
                motif.pk_index,
            )
            flip, seq, energy, tolerance, pk_index = data
            if flip:
                st_state.modified_motif_text += "\nmotif.flip()"
                motif.flip()

            if seq and motif.get_kissing_sequence() != seq:
                st_state.modified_motif_text += f"\nmotif.set_sequence('{seq}')"
                motif.set_sequence(seq)

            if energy and energy != motif.energy:
                st_state.modified_motif_text += f"\nmotif.energy = {energy}"
                motif.energy = energy

            if tolerance and tolerance != motif.energy_tolerance:
                st_state.modified_motif_text += (
                    f"\nmotif.energy_tolerance = " f"{tolerance}"
                )
                motif.energy_tolerance = tolerance

            if pk_index and pk_index != motif.pk_index:
                st_state.modified_motif_text += f'\nmotif.pk_index = "{pk_index}"'
                motif.pk_index = pk_index

        ### Create new motif
        else:
            kl_selction = option_menu(
                None,
                list(kl_name_map.keys()),
                icons=[MOTIF_ICONS[kl_name_map[name]] for name in kl_name_map],
                menu_icon="cast",
                orientation="horizontal",
                styles=second_menu_style,
                key="KLOption",
            )

            name = kl_name_map[kl_selction]
            kl_class = getattr(kissing_loops, name)

            flip_kl = st_state.current_line_occupied and "Dimer" not in name
            data = self.interface(flip=flip_kl)
            flip, top_seq, kl_energy, tolerance, pk_index = data

            if top_seq:
                motif = kl_class(open_left=flip, sequence=top_seq)
                st_state.motif_buffer = (
                    f"motif = pf.{name}("
                    f"open_left = {flip}, "
                    f"sequence = '{top_seq}')"
                )
            else:
                st_state.motif_buffer = (
                    f"motif = pf.{name}("
                    f"open_left = {flip}, "
                    f"energy = {kl_energy}, "
                    f"energy_tolerance = {tolerance}, "
                    f'pk_index = "{pk_index}")'
                )

                motif = kl_class(
                    open_left=flip,
                    energy=kl_energy,
                    energy_tolerance=tolerance,
                    pk_index=pk_index,
                )

            # save the motif in the session state
            st_state.motif = motif

        # reset the warning filter
        warnings.filterwarnings("default")

    def interface(
        self,
        key="",
        top_seq=None,
        kl_energy=-8.5,
        energy_tolerance=0.5,
        default_pk_index=None,
        flip=False,
    ):

        def pk_index_interface():
            all_pk_indexes = ["0"]
            last_ind = 0
            for last_ind in range(1, st_state.max_pk_index):
                all_pk_indexes.append(str(last_ind))
                all_pk_indexes.append(str(last_ind) + "'")
            all_pk_indexes.append(str(last_ind + 1))
            if default_pk_index is not None:
                index_for_list = all_pk_indexes.index(default_pk_index)
            else:
                index_for_list = len(all_pk_indexes) - 1
            pk_index = st.selectbox(
                "Pseudoknot id:",
                options=all_pk_indexes,
                index=index_for_list,
                key=f"pk_index_kl{key}",
                help="Select the pseudoknot id for the kissing loop motif. "
                "The id is used to identify the pseudoknot in the sequence."
                " Pseudoknots with the same id will have the same sequence."
                ' Pseudoknots with the same id, but with a quote "\'"'
                " will have a complementary sequence. Pseudoknots with id 0"
                " will not pair with any other pseudoknot.",
            )
            return pk_index

        col1, col2, col3 = st.columns([1, 1, 5], vertical_alignment="bottom")
        open_left = False
        with col1:
            specific_seq = st.toggle("Custom Sequence", key=f"seq_int_kl{key}")
        with col2:
            if key == "mod":
                open_left = st.button("Flip", key=f"open_left_kl{key}")
            else:
                open_left = st.toggle("Flip", value=flip, key=f"open_left_kl{key}")

        with col3:
            if specific_seq:
                subcol1, subcol2 = st.columns([5, 1])
                with subcol1:  # either sequence of energy info
                    top_seq = st.text_input(
                        "Sequence:", key=f"seq_top_int_kl{key}", value=top_seq
                    )
                with subcol2:
                    pk_index = pk_index_interface()

            else:
                subcol1, subcol2 = st.columns([4, 1])
                min_kl_E = abs(kl_energy) - abs(energy_tolerance)
                max_kl_E = abs(kl_energy) + abs(energy_tolerance)
                with subcol1:
                    kl_range = st.slider(
                        "Kissing binding energy range:",
                        min_value=2.0,
                        max_value=15.0,
                        value=(min_kl_E, max_kl_E),
                        step=0.1,
                        format="-%.1f",
                        key=f"energy_range_kl{key}",
                    )
                    kl_energy = round(-sum(kl_range) / 2, 2)
                    energy_tolerance = round(abs(kl_range[1] - kl_range[0]) / 2, 3)
                with subcol2:
                    pk_index = pk_index_interface()

        return open_left, top_seq, kl_energy, energy_tolerance, pk_index
