import pandas as pd
import streamlit as st
import os, shutil
from tiger import tiger_exhibit, load_transcripts, TARGET_LEN, NUCLEOTIDE_TOKENS


@st.cache_data
def convert_df(df):
    # IMPORTANT: Cache the conversion to prevent computation on every rerun
    return df.to_csv().encode('utf-8')


# title and initialization
st.title('TIGER Cas13 Efficacy Prediction')
st.session_state['manual_seq'] = ''
st.session_state['fasta_seq'] = ''

# run mode selection
with st.form(key='calc_options'):
    c1_names = ['On-target', 'On- and off-target']
    option = st.radio('Select mode:', c1_names, index=0)
    submitButton = st.form_submit_button(label='Choose options')

# text input
manual_entry = st.form('text')
manual_input = manual_entry.text_input(
    label='Enter a target transcript:',
    # value='ATGCAGGACGCGGAGAACGTGGCGGTGCCCGAGGCGGCCGAGGAGCGCGC',
    placeholder='Upper or lower case')
if manual_input:
    if len(manual_input) < TARGET_LEN:
        manual_entry.write('Transcript must be at least {:d} bases.'.format(TARGET_LEN))
    else:
        st.session_state['manual_seq'] = manual_input
manual_mode = manual_entry.form_submit_button(label='calculate')

# fasta input
fasta_form = st.form('fasta')
fasta_input = fasta_form.file_uploader(label='Upload a fasta file:')
if fasta_input:
    if os.path.exists('temp'):
        shutil.rmtree('temp')
    os.makedirs('temp')
    st.write(fasta_input.name)
    fpath = os.path.join('temp', fasta_input.name)
    with open(fpath, 'w') as f:
        f.write(fasta_input.getvalue().decode('utf-8'))
    transcript_tbl = load_transcripts([fpath])
    fasta_form.text('fasta file contents')
    fasta_form.write(transcript_tbl)
    seq = transcript_tbl['seq'][0]
    st.session_state['fasta_seq'] = seq
fasta_mode = fasta_form.form_submit_button(label='Calculate')

# input-specific configuration
if manual_mode:
    src_seq = st.session_state['manual_seq']
    status_text = manual_entry.empty()
    status_bar = manual_entry.progress(0)
elif fasta_mode:
    src_seq = st.session_state['fasta_seq']
    status_text = fasta_form.empty()
    status_bar = fasta_form.progress(0)
else:
    src_seq = status_bar = status_text = None

# valid input
if src_seq and all([True if nt.upper() in NUCLEOTIDE_TOKENS.keys() else False for nt in src_seq]):
    on_target, off_target = tiger_exhibit(pd.DataFrame(dict(id=['ManualEntry'], seq=[src_seq])),
                                          status_bar, status_text, check_off_targets=option == 'On and Off Target')
    on_target.rename(columns={'Guide': '23 nt guide sequence'}, inplace=True)
    if len(on_target) > 0:
        if on_target.iloc[0]['On-target ID'] == 0:
            on_target.drop(['On-target ID'], axis=1, inplace=True)
    st.write('On-target predictions: ', on_target)
    st.download_button(label='Download', data=convert_df(on_target), file_name='on_target.csv', mime='text/csv')
    if option == 'On and Off Target' and len(off_target) > 0:
        off_target.rename(columns={'Guide': '23 nt guide sequence'}, inplace=True)
        st.write('Off-target predictions: ', off_target)
        st.download_button(label='Download', data=convert_df(off_target), file_name='off_target.csv', mime='text/csv')
    elif option == 'On and Off Target' and len(off_target) == 0:
        st.write('We did not find any off-target effects!')

# invalid input
else:
    st.write('Invalid input!')