import streamlit as st
from streamlit_option_menu import option_menu
import json
import os
import warnings

# App utilities
from utils import (load_logo, 
                   main_menu_style, 
                   copy_to_clipboard, 
                   write_format_text,
                   check_import_pyfurnace)
from utils.template_functions import symbols, reference, sanitize_input
from utils.design_functions import origami_build_view
check_import_pyfurnace()
from pyfurnace import prepare as prep

# Melting and primer logic
from pyfurnace.prepare.utils import (
    tm_methods, tm_models, default_values,
    calculate_gc, make_tm_calculator, check_dimer,
    auto_design_primers, annealing_temp
)


def upload_setting_button():
    """Allow to upload setting"""
    st.session_state['upload_setting'] = True
    return

def calculate_annealing(seq, mts, c_primer, nc_primer, tm_kwargs):

    col1, col2 = st.columns(2, gap='large')
    with col1:
        subcol1, subcol2 = st.columns([11,1], vertical_alignment='center')
        with subcol1:
            write_format_text(c_primer)
        with subcol2:
            copy_to_clipboard(c_primer, "")

        subcol1, subcol2, subcol3 = st.columns(3)
        with subcol1:
            st.markdown(f"GC content: {calculate_gc(c_primer)}%")
        with subcol2:
            st.markdown(f"**:green[Tm: {round(mts[0], 1)}°C]**")
        with subcol3:
            st.markdown(f"Length: {len(c_primer)}")

    with col2:
        
        subcol1, subcol2 = st.columns([11,1], vertical_alignment='center')
        with subcol1:
            write_format_text(nc_primer)
        with subcol2:
            copy_to_clipboard(nc_primer, "")

        subcol1, subcol2, subcol3 = st.columns(3)
        with subcol1:
            st.markdown(f"GC content: {calculate_gc(nc_primer)}%")
        with subcol2:
            st.markdown(f"**:green[Tm: {round(mts[1], 1)}°C]**")
        with subcol3:
            st.markdown(f"Length: {len(nc_primer)}")

    # show the primers preview, check the self-dimerization of the primer and 
    # check the dimer between the primers and the sequence
    c_bases = len(c_primer)
    nc_bases = len(nc_primer)
    with st.expander("Dimerization preview"):
        st.markdown(f'<div style="text-align: center;"><span style="color: #D00000">'
                    f'{c_primer}</span>{seq[c_bases]}[...]{seq[-nc_bases-1:]}</div>', 
                    unsafe_allow_html=True)
        st.markdown(f'<div style="text-align: center;">'
                    f'{prep.Seq(seq[:c_bases+1]).complement()}'
                    f'[...]'f'{prep.Seq(seq[-nc_bases - 1]).complement()}'
                    f'<span style="color: #D00000">'
                    f'{nc_primer[::-1]}</span></div>', 
                    unsafe_allow_html=True)
        st.divider()
        col1, col2 = st.columns(2, gap='large')
        with col1:
            self_dimer1 = check_dimer(c_primer, 
                                      c_primer, 
                                      basepair={'A': 'T', 
                                                'T': 'A', 
                                                'C': 'G', 
                                                'G': 'C'})
            write_format_text('Self-Dimer forward primer\n' + self_dimer1)

        with col2:
            self_dimer2 = check_dimer(nc_primer, 
                                      nc_primer, 
                                      basepair={'A': 'T', 
                                                'T': 'A', 
                                                'C': 'G', 
                                                'G': 'C'})
            write_format_text('Self-Dimer reverse primer\n' + self_dimer2)

        self_dimer12 = check_dimer(c_primer, 
                                   nc_primer, 
                                   basepair={'A': 'T', 'T': 'A', 'C': 'G', 'G': 'C'})
        write_format_text('Dimer between primers\n' + self_dimer12)

    # add a warning if the melting temperature of the primers is too different
    anneal_color = 'green'
    if abs(mts[0] - mts[1]) >= 5:
        st.warning('The difference of Tm should be below 5°C', icon="⚠️")
        anneal_color = 'red'
                    
    # take into account the two main method to calculate the PCR annealing 
    # temperature: IDT method and Phusion Buffer method
    t_anneal = annealing_temp(mts, 
                              seq, 
                              tm_kwargs, 
                              method=st.session_state.annealing_model.split()[0])
    if t_anneal < 50:
        st.warning('Is suggested a temperature of annealing higher than 50°C.', 
                   icon="⚠️")
        anneal_color = 'red'

    # write the annealing temperature in big
    st.markdown(f"### Anneal at: :{anneal_color}[{t_anneal}°C]")

def primers_tab(seq, mt_correct):
    """ Calculate the melting temperature of the primers and check 
    the dimer between the primers and the sequence"""

    # show the settings for the two primers: 
    # choose the number of bases and show the gc content and melting temperature
    st.write('\n')
    col1, col2 = st.columns(2, gap='large')
    mts = [0, 0]
    # settings for the coding primer
    with col1:
        with st.columns(5)[2]:
            st.markdown('###### Forward')

        c_bases = st.slider('Primer length:', 
                            min_value=5, max_value=50, 
                            value=st.session_state.prim_fwd_len,
                            key=f"coding_primer_{st.session_state.auto_count}")
        c_primer = seq[:c_bases]

        mts[0] = round(mt_correct(c_primer),2)

    # settings for the non-coding primer
    with col2:
        with st.columns(5)[2]:
            st.markdown('###### Reverse')

        nc_bases = st.slider('Primer length:', 
                             min_value=5, max_value=50, 
                             value=st.session_state.prim_rev_len, 
                             key=f"non_coding_primer_{st.session_state.auto_count}")
        
        nc_primer = str(prep.Seq(seq[-nc_bases:]).reverse_complement())

        mts[1] = round(mt_correct(nc_primer),2)

    return mts, c_primer, nc_primer

def tm_parameters():
    # load settings if you want to upload custom settings
    load_settings = st.checkbox("Upload previous energy parameter:",
                                help='If you have saved a json file with the '
                                    'energy parameters, you can upload it and'
                                    ' use the same settings.')
    if load_settings:
        saved_setting = st.file_uploader("Upload previous energy parameter "
                                            "(optional):",
                                            on_change=upload_setting_button,
                                        type='json',
                                        help='If you have saved a json file '
                                            'with the energy parameters, you '
                                            'can upload it and use the same '
                                            'settings.')
        
        # if load the settings, upload the valuse from the json file
        if saved_setting and st.session_state['upload_setting']:
            # To read file as bytes:
            session_old = json.load(saved_setting)
            for key, value in session_old.items():
                st.session_state[key] = value
            st.session_state['upload_setting'] = False
            st.rerun()


    # initialize session state values for the energy model
    for key, val in default_values.items():
        if key not in st.session_state:
            st.session_state[key] = val

    # select the energy model
    st.write('Energy model:')
    col1, col2, col3 = st.columns([2, 2, 1])

    # select the melting temperature method
    with col1: 
        mt_type = st.selectbox("Melting Temperature Method", 
                                list(tm_methods), 
                                help=(prep.mt.__doc__), 
                                key='mt_method', 
                                label_visibility='collapsed')
        method = tm_methods[mt_type]

    # select the melting temperature model
    with col2:
        mt_model = st.selectbox("Energy Correction", 
                                list(tm_models[mt_type]), 
                                key='mt_model', 
                                label_visibility='collapsed')
    
    # add a button to show the help of the selected method
    with col3:
        help_model = st.button("Show info", key='energy_model_info')
    if help_model:
        st.help(method)
        
    # buffer correction parameters
    st.divider()
    if method != list(tm_methods)[2]:

        st.write('Buffer corrections (mM):')

        cols = st.columns(8)
        with cols[0]: 
            na = st.number_input('Na', key="Na")
        with cols[1]: 
            k = st.number_input('K', key='K')
        with cols[2]: 
            tris = st.number_input('Tris', key='Tris')
        with cols[3]: 
            mg = st.number_input('Mg', 
                                value=st.session_state["Mg"], 
                                key='Mg')
        with cols[4]: 
            dntps = st.number_input('dNTPs', 
                                        value=st.session_state["dNTPs"],
                                        key='dNTPs')
        with cols[5]: 
            dmso = st.number_input('DMSO (%)', 
                                    value=st.session_state["DMSO (%)"],
                                    key='DMSO (%)')
        with cols[6]: 
            correction_met = st.selectbox('Method', 
                                            [0, 1, 2, 3, 4, 6, 7],
                                            key='Method')
        with cols[7]: 
            help_correction = st.button("Show info", 
                                        key='energy_correction_info')

        tm_kwargs = {'Na': na, 
                        'K': k, 
                        'Tris': tris, 
                        'Mg': mg, 
                        'saltcorr': correction_met,
                        'dNTPs': dntps, 
                        'DMSO': dmso}

    if help_correction:
        # add a button to show the help of the selected method
        st.help(prep.mt.salt_correction)
        st.help(prep.mt.chem_correction)

    primer_conc = st.number_input('Primer conc. (nM)', key="Primer")

    # create the function to calculate the TM
    calculate_mt = make_tm_calculator(method_name=mt_type,
                                        model_name=mt_model,
                                        primer_conc=primer_conc,
                                        tm_kwargs=tm_kwargs)

    # save settings and allow download
    with open('energy_parameters.json', 'w') as settings:
        # cannot save session state as it is, 
        # I have to convert it to a dictionary
        session_dict = {key: st.session_state[key] for key in default_values}
        json.dump(session_dict, settings)
    with open("energy_parameters.json", "rb") as file:
        btn = st.download_button(
            label="Download energy parameters",
            data=file,
            file_name="energy_parameters.json",
            mime="application/json",
            help='Save the current settings (e.g. ions concentration, '
                    'Tm model), so you can easily reload them if you refresh '
                    'the page!'
            )
        
    return calculate_mt, tm_kwargs

def primers_setup():
    if "dna_template" not in st.session_state:
        st.session_state["dna_template"] = ''

    # take the input sequence and sanitize it
    seq = sanitize_input(st.text_input("Input sequence:", 
                                       value = st.session_state["dna_template"]))

    # check the symbols in the sequence
    if set(seq) - symbols:
        st.warning('The sequence contains symbols not included in the '
                    '[IUPAC alphabet](https://www.bioinformatics.org/sms/iupac.html).',
                    icon="⚠️")
    if 'U' in seq:
        st.error('The DNA template contains U', icon=":material/personal_injury:")
    
    if not seq:
        st.stop()
    elif len(seq) < 40:
        st.error('The sequence is too short', icon=":material/personal_injury:")
        st.stop()
    
    ###
    # Various for PCR settings
    ###

    mcol1, mcol2, mcol3, mcol4 = st.columns(4, 
                                            vertical_alignment='center',
                                            gap='large') 
    with mcol1:
        with st.popover("Melting temperature parameters",
                        use_container_width=True):
            
            calculate_mt, tm_kwargs = tm_parameters()
    
    with mcol2:
        with st.popover('Add restriction sites',
                        use_container_width=True):
            restric_site_1 = st.text_input("Restriction site sequence (5'->3')"
                                           " before the promoter:")
            restric_site_2 = st.text_input("Restriction site sequence (5'->3')"
                                           " after the fragment:")
            
    with mcol3:
        with st.popover('Annealing calculation',
                        use_container_width=True):
            st.selectbox("Calculate annealing:", 
                        ['IDT method [2]', 'Phusion method [3]'],
                        key='annealing_model',
                        label_visibility='collapsed',
                        help='Select the method to calculate the annealing '
                             'temperature. Check the references for more information.'
                             )
    with mcol4:
        # default primers length
        with st.popover('**Auto design primers**', 
                        icon=":material/precision_manufacturing:",
                        use_container_width=True):
            
            target_temp = st.number_input('Target melting temperature (°C)', 
                                value=65, min_value=0, max_value=100, step=1)
            if st.button("Auto design primers", 
                      key='auto_design_primers',
                      help='Automatically design the primers for the sequence. '
                           'The primers are designed to have a melting temperature '
                           'and follow best practices (see documentation).'):

                status = st.empty()
                status.info("Designing primers...", 
                            icon=":material/precision_manufacturing:")

                primers, final_mts = auto_design_primers(seq, 
                                                         target_temp, 
                                                         calculate_mt)

                st.session_state.auto_count += 1

                # set the new primers length in the session state
                if (st.session_state.prim_fwd_len != len(primers[0]) or 
                        st.session_state.prim_rev_len != len(primers[1])):
                    st.session_state.prim_fwd_len = len(primers[0])
                    st.session_state.prim_rev_len = len(primers[1])

                status.success(f"Primer designed!", 
                                icon=":material/precision_manufacturing:")
                

    seq = restric_site_1 + seq + str(prep.Seq(restric_site_2).reverse_complement())
    
    mts, c_primer, nc_primer = primers_tab(seq, 
                                           calculate_mt)

    calculate_annealing(seq, mts, c_primer, nc_primer, tm_kwargs)

    # add bibliography
    reference(True)


def md_setup():

    # check if the origami blueprint is good
    if (not st.session_state.get('origami') or
        not all(n in 'AUCG&' for n in st.session_state.origami.sequence)):

        st.error("The origami blueprint is empty or it doesn't contain a sequence",
                 icon=":material/personal_injury:")
        
        col1, col2 = st.columns(2)
        with col1:
            st.page_link("pages/1_Design.py", 
                         label="**:orange[Have a look to the blueprint]**", 
                         icon=":material/draw:")
        
        with col2:
            st.page_link("pages/2_Generate.py", 
                         label="**:orange[Generate the sequence]**", 
                         icon=":material/network_node:")
        st.stop()
    
    # show the origami blueprint
    with st.expander('**Origami preview:**', expanded=True):
        origami_build_view('Origami split view')

    st.divider()
    col1, col2 = st.columns(2)
    with col1:
        oxdna_dir = st.text_input('OxDNA directory '
                                  '(e.g. "~/Documents/software/oxDNA")')
    with col2:
        temp = st.number_input('Temperature (°C)', 
                               value=37, min_value=0, max_value=100, step=1)

    help_mc_rel = ("Use Monte Carlo relaxation to relax the origami before the MD "
                  "simulation, using external forces to keep the basepairing.")
    help_md_rel = ("Use Molecular Dynamics to relax the structure, using forces to"
                   "maintain the basepairing")
    help_md_equ = ("An extra Molecular Dynamics simulation to make sure the Origami"
                   "is fully relaxed and not bias the starting configuration of the " 
                   "production. **External forces for pseudoknots only** are used "
                   "in this simulation, to make sure the pseudoknots are paired "
                   "at the start of the production")
    help_md_run = ("The final Molecular Dynamics simulation, to simulate the "
                   "structure without any external force.")
    
    num_input = st.toggle("Use numerical input", key='num_input')

    col1, col2, col3, col4 = st.columns(4)
    with col1:
        if num_input:
            step_mc_rel = st.number_input('MC relaxation steps', 
                                          min_value=0,
                                          value=int(5e3),
                                          help=help_mc_rel,
                                          )
        else:
            step_mc_rel = st.slider('MC relaxation steps', 
                                    min_value=0,
                                    value=int(5e3),
                                    max_value=int(10e3),
                                    step=int(1e3),
                                    help=help_mc_rel,
                                    format="%0.0e"
                                    )
    with col2:
        if num_input:
            step_md_rel = st.number_input('MD relaxation steps', 
                                          min_value=0,
                                          value=int(1e7),
                                          help=help_md_rel,
                                          )
        else:
            step_md_rel = st.slider('MD relaxation steps', 
                                    min_value=0,
                                    value=int(1e7),
                                    max_value=int(1e8),
                                    step=int(1e6),
                                    help=help_md_rel,
                                    format="%0.0e"
                                    )
    with col3:
        if num_input:
            step_md_equ = st.number_input('MD equilibration steps', 
                                          min_value=0,
                                          value=int(1e8),
                                          help=help_md_equ,
                                          )
        else:
            step_md_equ = st.slider('MD equilibration steps', 
                                    min_value=0,
                                    value=int(1e8),
                                    max_value=int(1e9),
                                    step=int(1e7),
                                    help=help_md_equ,
                                    format="%0.0e"
                                    )
    with col4:
        if num_input:
            step_md_run = st.number_input('MD production steps', 
                                          min_value=0,
                                          value=int(1e9),
                                          help=help_md_run,
                                          )
        else:
            step_md_run = st.slider('MD production steps', 
                                    min_value=0,
                                    value=int(1e9),
                                    max_value=int(1e10),
                                    step=int(1e8),
                                    help=help_md_run,
                                    format="%0.0e"
                                    )
            
    if not oxdna_dir or not temp:
        st.stop()

    zip_path = prep.oxdna_simulations(origami=st.session_state.origami,
                                      oxdna_directory=oxdna_dir,
                                      temperature=temp,
                                      mc_relax_steps=step_mc_rel,
                                      md_relax_steps=step_md_rel,
                                      md_equil_steps=step_md_equ,
                                      md_prod_steps=step_md_run
                                      )
    # add a button to download the zip file
    st.divider()

    col1, col2 = st.columns(2, vertical_alignment='bottom')
    with col1:
        name = st.text_input('Name of the simulation:', value='origami_simulation')
    with col2:
        with open(zip_path, "rb") as file:
            st.download_button(
                label="Download MD simulation files",
                data=file,
                file_name=f"{name}.zip",
                mime="application/zip",
                help='Download the MD simulation files. '
                     'The zip file contains the input files for the MD simulation.'
                )
        
    # delete the zip file
    if os.path.exists(zip_path):
        os.remove(zip_path)
    return

if __name__ == "__main__":
    ### set the logo of the app
    load_logo()
    warnings.filterwarnings("ignore") # ignore warnings

    if "prepare_ind" not in st.session_state:
        st.session_state.prepare_ind = 0
    if 'prim_fwd_len' not in st.session_state:
        st.session_state.prim_fwd_len = 21
    if 'prim_rev_len' not in st.session_state:
        st.session_state.prim_rev_len = 21
    if 'auto_count' not in st.session_state:
        st.session_state.auto_count = 0

    # create the tabs with the functions
    st.header('Prepare', help='Design primers for your DNA template '
                               'or prepare the Origami for OxDNA simulation.')
    option_data = {'Primers': "bi bi-arrow-left-right",
                   'MD simulations': "bi bi-cpu"}

    selected_operation = option_menu(None, 
                                    list(option_data.keys()),
                                    icons=list(option_data.values()),
                                    menu_icon="cast", 
                                    orientation="horizontal",
                                    default_index=st.session_state.prepare_ind,
                                    styles=main_menu_style)
    
    if selected_operation == 'MD simulations':
        if st.session_state.prepare_ind == 0:
            st.session_state.prepare_ind = 1
            st.rerun()
        md_setup()
    elif selected_operation == 'Primers':
        if st.session_state.prepare_ind == 1:
            st.session_state.prepare_ind = 0
            st.rerun()
        primers_setup()
    