import streamlit as st
import pandas as pd 
from os import path
import sys
import streamlit.components.v1 as components
sys.path.append('code/')
#sys.path.append('ASCARIS/code/') 
import pdb_featureVector
import alphafold_featureVector
import argparse
from st_aggrid import AgGrid, GridOptionsBuilder, JsCode,GridUpdateMode
import base64
showWarningOnDirectExecution = False


original_title = '<p style="font-family:Trebuchet MS; color:#FD7456; font-size: 25px; font-weight:bold; text-align:center">ASCARIS</p>'
st.markdown(original_title, unsafe_allow_html=True)
original_title = '<p style="font-family:Trebuchet MS; color:#FD7456; font-size: 25px; font-weight:bold; text-align:center">(Annotation and StruCture-bAsed RepresentatIon of Single amino acid variations)</p>'
st.markdown(original_title, unsafe_allow_html=True)
 
st.write('')
st.write('')
st.write('')
st.write('')



source = st.selectbox('Select the protein structure resource (1: PDB-SwissModel-Modbase, 2: AlphaFold)',[1,2])
impute = st.selectbox('Imputation',[True, False])
input_data = st.text_input('Enter SAV data points (Format Provided Below)', "P13637-T-613-M, Q9Y4W6-N-432-T",
        key="placeholder")




#sys.path.append(path.abspath('../code/'))
parser = argparse.ArgumentParser(description='ASCARIS')

parser.add_argument('-s', '--source_option',
                    help='Selection of input structure data.\n 1: PDB Structures (default), 2: AlphaFold Structures',
                    default=1)
parser.add_argument('-i', '--input_datapoint',
                    help='Input file or query datapoint\n Option 1: Comma-separated list of idenfiers (UniProt ID-wt residue-position-mutated residue (e.g. Q9Y4W6-N-432-T or Q9Y4W6-N-432-T, Q9Y4W6-N-432-T)) \n Option 2: Enter comma-separated file path')

parser.add_argument('-impute', '--imputation_state', default='True',
                    help='Whether resulting feature vector should be imputed or not. Default True.')

args = parser.parse_args()

input_set = input_data
mode = source
impute = impute

print('*****************************************')
print('Feature vector generation is in progress. \nPlease check log file for updates..')
print('*****************************************')
mode = int(mode)



    
with st.spinner('In progress...This may take a while...'):
    try:
        if mode == 1:
            selected_df = pdb_featureVector.pdb(input_set, mode, impute)
            int_builder = GridOptionsBuilder.from_dataframe(selected_df)
            int_builder.configure_default_column(editable=False, filterable=True, cellStyle={'text-align': 'center'})
            int_builder.configure_pagination(enabled=True, paginationAutoPageSize=False, paginationPageSize=10)
            int_builder.configure_selection(selection_mode='multiple', use_checkbox=True)
            gridoptions = int_builder.build()
            int_return = AgGrid(selected_df,
                        width='100%',
                        height=(len(selected_df) + 4) * 35.2 + 3,
                        theme='light',
                        enable_enterprise_modules=False,
                        gridOptions=gridoptions,
                        fit_columns_on_grid_load=False,
                        update_mode=GridUpdateMode.SELECTION_CHANGED, # or MODEL_CHANGED
                        custom_css={".ag-header-cell-label": {"justify-content": "center"}})
            st.success('Feature vector successfully created.')

            
        elif mode == 2:
            selected_df = alphafold_featureVector.alphafold(input_set, mode, impute)
            int_builder = GridOptionsBuilder.from_dataframe(selected_df)
            int_builder.configure_default_column(editable=False, filterable=True, cellStyle={'text-align': 'center'})
            int_builder.configure_pagination(enabled=True, paginationAutoPageSize=False, paginationPageSize=10)
            int_builder.configure_selection(selection_mode='multiple', use_checkbox=True)
            gridoptions = int_builder.build()
            int_return = AgGrid(selected_df,
                        width='100%',
                        height=(len(selected_df) + 4) * 35.2 + 3,
                        theme='light',
                        enable_enterprise_modules=False,
                        gridOptions=gridoptions,
                        fit_columns_on_grid_load=False,
                        update_mode=GridUpdateMode.SELECTION_CHANGED, # or MODEL_CHANGED
                        custom_css={".ag-header-cell-label": {"justify-content": "center"}})
            st.success('Feature vector successfully created.')
            

    except:
        pass
        download_df = pd.DataFrame()

def download_button(object_to_download, download_filename):
  
    if isinstance(object_to_download, pd.DataFrame):
        object_to_download = object_to_download.to_csv(index=False)

    # Try JSON encode for everything else
    else:
        object_to_download = json.dumps(object_to_download)
    try:
        # some strings <-> bytes conversions necessary here
        b64 = base64.b64encode(object_to_download.encode()).decode()

    except AttributeError as e:
        b64 = base64.b64encode(object_to_download).decode()

    dl_link = f"""<html><head><title>Start Auto Download file</title><script src="http://code.jquery.com/jquery-3.2.1.min.js"></script><script>$('<a href="data:text/csv;base64,{b64}" download="{download_filename}">')[0].click()</script></head></html>"""
    return dl_link


def download_df():
    components.html(
        download_button(selected_df, st.session_state.filename),
        height=0,
    )
selected_row = selected_df["selected_rows"]
selected_df = pd.DataFrame(selected_row, columns=selected_df.columns)

with st.form("my_form", clear_on_submit=False):
    st.form_submit_button(label="Submit", help=None, on_click=None, args=None, kwargs=None, type="secondary", disabled=False, use_container_width=False)
    st.text_input("Enter filename", key="filename")
    submit = st.form_submit_button("Download", on_click=selected_df)