from pathlib import Path
import torch
# from st_on_hover_tabs import on_hover_tabs
import streamlit as st
st.set_page_config(layout="wide")

model_path = './model.iter-685000'

import sys, os
import rdkit
import rdkit.Chem as Chem
from rdkit.Chem.Draw import MolToImage
from rdkit.Chem import Descriptors
import sascorer
import networkx as nx
from stqdm import stqdm
import base64, io
import pandas as pd
import streamlit_ext as ste

os.environ['KMP_DUPLICATE_LIB_OK']='True'

sys.path.append('%s/fast_jtnn/' % os.path.dirname(os.path.realpath(__file__)))
from mol_tree import Vocab, MolTree
from jtprop_vae import JTPropVAE
from molbloom import buy

css='''
[data-testid="metric-container"] {
    width: fit-content;
    margin: auto;
}

[data-testid="metric-container"] > div {
    width: fit-content;
    margin: auto;
}

[data-testid="metric-container"] label {
    width: fit-content;
    margin: auto;

[data-testid="stDataFrameResizable"] {
    width: fit-content;
    margin: auto;
}
}
'''

st.markdown(f'<style>{css}</style>',unsafe_allow_html=True)

s_buff = io.StringIO()
def img_to_bytes(img_path):
    img_bytes = Path(img_path).read_bytes()
    encoded = base64.b64encode(img_bytes).decode()
    return encoded
def img_to_html(img_path):
    img_html = "<img src='data:image/png;base64,{}' class='img-fluid' style='max-width: 500px;'>".format(
      img_to_bytes(img_path)
    )
    return img_html

_mcf = pd.read_csv('./mcf.csv')
_pains = pd.read_csv('./wehi_pains.csv',
                     names=['smarts', 'names'])
_mcf_filters = [Chem.MolFromSmarts(x) for x in
            _mcf['smarts'].values]
_pains_filters = [Chem.MolFromSmarts(x) for x in
            _pains['smarts'].values]

def mol_passes_filters_custom(mol,
                       allowed=None,
                       isomericSmiles=False):
    """
    Checks if mol
    * passes MCF and PAINS filters,
    * has only allowed atoms
    * is not charged
    """
    allowed = allowed or {'C', 'N', 'S', 'O', 'F', 'Cl', 'Br', 'H'}
    if mol is None:
        return 'NoMol'
    ring_info = mol.GetRingInfo()
    if ring_info.NumRings() != 0 and any(
            len(x) >= 8 for x in ring_info.AtomRings()
    ):
        return 'ManyRings'
    h_mol = Chem.AddHs(mol)
    if any(atom.GetFormalCharge() != 0 for atom in mol.GetAtoms()):
        return 'Charged'
    if any(atom.GetSymbol() not in allowed for atom in mol.GetAtoms()):
        return 'AtomNotAllowed'
    if any(h_mol.HasSubstructMatch(smarts) for smarts in _mcf_filters):
        return 'MCF'
    if any(h_mol.HasSubstructMatch(smarts) for smarts in _pains_filters):
        return 'PAINS'
    smiles = Chem.MolToSmiles(mol, isomericSmiles=isomericSmiles)
    if smiles is None or len(smiles) == 0:
        return 'Isomeric'
    if Chem.MolFromSmiles(smiles) is None:
        return 'Isomeric'
    return 'YES'

def penalized_logp_standard(mol):

    logP_mean = 2.4399606244103639873799239
    logP_std = 0.9293197802518905481505840
    SA_mean = -2.4485512208785431553792478
    SA_std = 0.4603110476923852334429910
    cycle_mean = -0.0307270378623088931402396
    cycle_std = 0.2163675785228087178335699

    log_p = Descriptors.MolLogP(mol)
    SA = -sascorer.calculateScore(mol)

    # cycle score
    cycle_list = nx.cycle_basis(nx.Graph(Chem.rdmolops.GetAdjacencyMatrix(mol)))
    if len(cycle_list) == 0:
        cycle_length = 0
    else:
        cycle_length = max([len(j) for j in cycle_list])
    if cycle_length <= 6:
        cycle_length = 0
    else:
        cycle_length = cycle_length - 6
    cycle_score = -cycle_length
    # print(logP_mean)

    standardized_log_p = (log_p - logP_mean) / logP_std
    standardized_SA = (SA - SA_mean) / SA_std
    standardized_cycle = (cycle_score - cycle_mean) / cycle_std
    return log_p,SA,cycle_score,standardized_log_p + standardized_SA + standardized_cycle

def df_to_file(df):
    s_buff.seek(0)
    df.to_csv(s_buff)
    return s_buff.getvalue().encode()

def download_df(df,id):
    with st.expander(':arrow_down: Download this dataframe'):
        st.markdown("<h4 style='color:tomato;'>Select column(s) to save:</h4>",unsafe_allow_html=True)
        for col in df.columns:
            st.checkbox(col,key=str(id)+'_col_'+str(col))
        st.text_input('File name (.csv):','dataframe',key=str(id)+'_file_name')
        
        ste.download_button('Download',df_to_file(df[[col for col in df.columns if st.session_state[str(id)+'_col_'+str(col)]]]),st.session_state[str(id)+'_file_name']+'.csv')

lg = rdkit.RDLogger.logger() 
lg.setLevel(rdkit.RDLogger.CRITICAL)

    
if 'current_view' not in st.session_state:
    st.session_state['current_view'] = 0
    
if 'current_step' not in st.session_state:
    st.session_state['current_step'] = 0
    
def set_page_view(id):
    st.session_state['current_view'] = id
    
def get_page_view():
    return st.session_state['current_view']
    
def set_step(id):
    st.session_state['current_step'] = id
    
def get_step():
    return st.session_state['current_step']

@st.cache_resource
def load_model():
    vocab = [x.strip("\r\n ") for x in open('./vocab.txt')] 
    vocab = Vocab(vocab)
    model = JTPropVAE(vocab, 450, 56, 20, 3)
    model.load_state_dict(torch.load(model_path,map_location=torch.device('cpu')))
    return model

from streamlit_lottie import st_lottie
import requests
    
def render_animation():
    animation_response = requests.get('https://assets1.lottiefiles.com/packages/lf20_vykpwt8b.json')
    animation_json = dict()
    
    if animation_response.status_code == 200:
        animation_json = animation_response.json()
    else:
        print("Error in the URL")
    return st_lottie(animation_json,height=200,width=300)

def oam_sidebar(step):
    st.title('**Optimize a molecule**')
    prog_bar = st.progress(0)
    # cur_step = get_step()
    if step == 0: prog_bar.progress(0)
    if step == 1: prog_bar.progress(33)
    if step == 2: prog_bar.progress(67)
    if step == 3: prog_bar.progress(100)
    st.markdown('\n')
    
    # st.markdown(get_step())
    color_ls = colorize_step(4,step)
    
    st.markdown("<h4 style='color: "+color_ls[0]+"'>Choose a molecule</h4>",unsafe_allow_html=True)
    st.markdown('|')
    st.markdown("<h4 style='color: "+color_ls[1]+"'>Choose settings</h4>",unsafe_allow_html=True)
    st.markdown('|')
    st.markdown("<h4 style='color: "+color_ls[2]+"'>Optimizing a molecule</h4>",unsafe_allow_html=True)
    st.markdown('|')
    st.markdown("<h4 style='color: "+color_ls[3]+"'>Finished</h4>",unsafe_allow_html=True)

# @st.cache_data(experimental_allow_widgets=True)

# if 'sidebar_con' not in st.session_state:
#     sidebar_con = st.empty()
# def render_sidebar(page,step):
#     sidebar_con.empty()
#     with sidebar_con.container():
#         if page == 0:
#             with st.sidebar():
#                 oam_sidebar(step)

def colorize_step(n_step,cur_step):
    color_list = ['grey']*n_step
    for i in range(cur_step):
        color_list[i] = 'mediumseagreen'
    color_list[cur_step] = 'tomato'
    if cur_step == n_step-1:
        color_list[cur_step] = 'mediumseagreen'
    return color_list

def form_header():
    st.markdown("<h1 style='text-align: center;'>Junction Tree Variational Autoencoder for Molecular Graph Generation (JTVAE)</h1>",unsafe_allow_html=True)
    st.markdown("<h3 style='text-align: center;'>Wengong Jin, Regina Barzilay, Tommi Jaakkola</h3>",unsafe_allow_html=True)

    # determines button color which should be red when user is on that given step
    oam_type = 'primary' if st.session_state['current_view'] == 0 else 'secondary'
    oab_type = 'primary' if st.session_state['current_view'] == 1 else 'secondary'
    ab_type = 'primary' if st.session_state['current_view'] == 2 else 'secondary'

    step_cols = st.columns([.85,.5,.5,.5,.85])    
    step_cols[1].button('Optimize a molecule',on_click=set_page_view,args=[0],type=oam_type,use_container_width=True)
    step_cols[2].button('Optimize a batch',on_click=set_page_view,args=[1],type=oab_type,use_container_width=True)        
    step_cols[3].button('About',on_click=set_page_view,args=[2],type=ab_type,use_container_width=True) 
    st.empty()
    
def form_body():
    body_container = st.empty()
    ###### Optimize a molecule ######
    if st.session_state['current_view'] == 0:
        body_container.empty()
        with body_container.container():
            Optimize_a_molecule()
    ###### Optimize a batch ######
    if st.session_state['current_view'] == 1: 
        body_container.empty()
        with body_container.container():
            Optimize_a_batch()
    ###### About ######
    if st.session_state['current_view'] == 2:
        body_container.empty()
        with body_container.container():
            About()
        
def About():
    descrip = '''
We seek to automate the design of molecules based on specific chemical properties. In computational terms, this task involves continuous embedding and generation of molecular graphs. Our primary contribution is the direct realization of molecular graphs, a task previously approached by generating linear SMILES strings instead of graphs. Our junction tree variational autoencoder generates molecular graphs in two phases, by first generating a tree-structured scaffold over chemical substructures, and then combining them into a molecule with a graph message passing network. This approach allows us to incrementally expand molecules while maintaining chemical validity at every step. We evaluate our model on multiple tasks ranging from molecular generation to optimization. Across these tasks, our model outperforms previous state-of-the-art baselines by a significant margin.

[https://arxiv.org/abs/1802.04364](https://arxiv.org/abs/1802.04364)'''
    st.markdown(descrip)
    st.markdown("<p style='text-align: center;'>"+
                    img_to_html('about.png')+
                    "</p>", unsafe_allow_html=True)
    
def Optimize_a_molecule():
    st.markdown("<h2 style='text-align: center;'>Optimize a molecule</h2>",unsafe_allow_html=True)
    with st.expander(':snowman: :blue[Instruction]'):
        guide = """<h4 style='color:tomato;'>Steps to optimize a molucule</h4>
        1. Select from examples, or manually enter a valid SMILES string of a molecule.</br>
        2. Configure the settings to generate a new molecule. The new molecule should have a higher penalized LogP value.</br>
            - Learning rate: How 'far' from the molecule that you want to search.</br>
            - Similarity cutoff: How 'similar' to the molecule that you want to search.</br>
            - Number of iterations: Number of generation trials.</br>
        <h4 style='color:darkturquoise;'>Annotation</h4>
        <b>SMILES</b> - Simplified molecular-input line-entry system</br>
        <b>LogP</b> - The log of the partition coefficient of a solute between octanol and water, at near infinite dilution</br>
        <b>SA score</b> - Synthetic Accessibility Score (lower is better)</br>
        <b>Cycle score</b> - A number of carbon rings of size larger than 6 (lower is better)</br>
        <b>Penalized LogP</b> - Standardized score of <i>LogP - SA score - Cycle score</i></br>
        <b>Similarity</b> - Molecular similarity is calculated via Morgan fingerprint of radius 2 with Tanimoto similarity</br>
        """
        st.markdown(guide,unsafe_allow_html=True)
        
    with st.sidebar:
        sidebar_con = st.empty()
    # sidebar_con.empty()
    with sidebar_con.container():
            set_step(0)
            oam_sidebar(0)
    oab_sel_container = st.empty()
    if 'checked_single' not in st.session_state:
        st.session_state.checked_single = 'NO'
    with oab_sel_container.container():
            oam_sel_col = st.columns([3,7])
            oam_sel_col[0].selectbox("Select an example",options=['-','Sorafenib','Pazopanib','Sunitinib'],key='mode',on_change=reset_oam_state)
            if st.session_state.mode == '-':
                st.session_state.smiles = oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):',max_chars=200)
            elif st.session_state.mode == 'Sorafenib':
                st.session_state.smiles = 'CNC(=O)C1=NC=CC(=C1)OC2=CC=C(C=C2)NC(=O)NC3=CC(=C(C=C3)Cl)C(F)(F)F'
                oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):','CNC(=O)C1=NC=CC(=C1)OC2=CC=C(C=C2)NC(=O)NC3=CC(=C(C=C3)Cl)C(F)(F)F',max_chars=200,disabled=True)
            elif st.session_state.mode == 'Pazopanib':
                st.session_state.smiles = 'CC1=C(C=C(C=C1)NC2=NC=CC(=N2)N(C)C3=CC4=NN(C(=C4C=C3)C)C)S(=O)(=O)N'
                oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):','CC1=C(C=C(C=C1)NC2=NC=CC(=N2)N(C)C3=CC4=NN(C(=C4C=C3)C)C)S(=O)(=O)N',max_chars=200,disabled=True)
            elif st.session_state.mode == 'Sunitinib':
                st.session_state.smiles = 'CCN(CC)CCNC(=O)C1=C(NC(=C1C)C=C2C3=C(C=CC(=C3)F)NC2=O)C'
                oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):','CCN(CC)CCNC(=O)C1=C(NC(=C1C)C=C2C3=C(C=CC(=C3)F)NC2=O)C',max_chars=200,disabled=True)
            st.columns([4,1,4])[1].button('Check SMILES',on_click=check_single,args=[st.session_state.smiles])
    
    check_single_con = st.empty()
    if 'smiles_selected' in st.session_state:
        if st.session_state.smiles_selected:
            with check_single_con.container():
                if 'checked_single' in st.session_state:
                    if st.session_state.checked_single == 'EnterError':
                        st.markdown("<p style='text-align: center; color: red;'><b>Please enter a SMILES string.</b></p>",unsafe_allow_html=True)
                        sidebar_con.empty()
                        with sidebar_con.container():
                                set_step(0)
                                oam_sidebar(0)
                    elif st.session_state.checked_single == 'MolError':
                        st.markdown("<p style='text-align: center; color: red;'><b>SMILES is invalid. Please enter a valid SMILES.</b></p>",unsafe_allow_html=True)
                        sidebar_con.empty()
                        with sidebar_con.container():
                                set_step(0)
                                oam_sidebar(0)
                    elif st.session_state.checked_single == 'YES':
                        st.markdown("<b>Canonicalized SMILES</b>",unsafe_allow_html=True)
                        st.code(st.session_state.smiles)
                        st.markdown("<p style='text-align: center; color: mediumseagreen'>MOSES filters passed successfully.</p>",unsafe_allow_html=True)
                        mol = Chem.MolFromSmiles(st.session_state.smiles)
                        imgByteArr = io.BytesIO()
                        MolToImage(mol,size=(400,400)).save(imgByteArr,format='PNG')
                        st.markdown("<p style='text-align: center;'>"+
                                f"<img src='data:image/png;base64,{base64.b64encode(imgByteArr.getvalue()).decode()}' class='img-fluid'>"+
                                "</p>", unsafe_allow_html=True)
                        # st.image(MolToImage(mol,size=(300,300)))
                        col1, col2, col3, col4 = st.columns(4)
                        col1.metric('LogP', '%.5f' % (st.session_state.logp))
                        col2.metric('SA score', '%.5f' % (-st.session_state.sa))
                        col3.metric('Cycle score', '%d' % (-st.session_state.cycle))
                        col4.metric('Penalized LogP', '%.5f' % (st.session_state.pen_p))
                        
                        st.session_state.smiles_checked = True
                        # render_sidebar()
                        # col1, col2, col3 = st.columns(3)
                        sidebar_con.empty()
                        with sidebar_con.container():
                                set_step(1)
                                oam_sidebar(1)
                        with st.form(":gear: Settings"):
                                st.slider('Choose learning rate: ',0.0,5.0,0.4,key='lr_s')
                                st.slider('Choose similarity cutoff: ',0.0,1.0,0.4,key='sim_cutoff_s')
                                st.slider('Choose number of iterations: ',1,100,80,key='n_iter_s')
                                st.session_state.optim_single_butt = st.form_submit_button("Optimize")
                    else:
                        st.markdown("<b>Canonicalized SMILES</b>",unsafe_allow_html=True)
                        st.code(st.session_state.smiles)
                        st.markdown("<p style='text-align: center; color: red;'><b>MOSES filters passed failed. Please use another molecule.</b></p>",unsafe_allow_html=True)
                        sidebar_con.empty()
                        with sidebar_con.container():
                                set_step(0)
                                oam_sidebar(0)
        else: check_single_con.empty()
    
    optim_single_con = st.empty()
    if 'optim_single_butt' in st.session_state:
        if st.session_state.optim_single_butt and st.session_state.smiles_checked:
            sidebar_con.empty()
            with sidebar_con.container():
                    set_step(2)
                    oam_sidebar(2)
            with optim_single_con.container():
                ani_con = st.empty()
                with ani_con.container():
                    st.markdown('Operation in progress. Please wait...')
                    render_animation()
                    new_smiles,sim = optim_single(st.session_state.smiles,st.session_state.lr_s,st.session_state.sim_cutoff_s,st.session_state.n_iter_s)
                ani_con.empty()
                sidebar_con.empty()
                with sidebar_con.container():
                        set_step(3)
                        oam_sidebar(3)
                if new_smiles is None:
                    st.markdown("<h4 style='text-align: center; color: red;'>Cannot optimize! Please choose another setting.</h4>",unsafe_allow_html=True)
                else:
                    st.markdown("<b style='text-align: center;'>New SMILES</b>",unsafe_allow_html=True)
                    st.code(new_smiles)
                    new_mol = Chem.MolFromSmiles(new_smiles)
                    if new_mol is None:
                        st.markdown("<p style='text-align: center; color: red;'>New SMILES is invalid! Please choose another setting.</p>",unsafe_allow_html=True)
                        # st.write('New SMILES is invalid.')
                    else:
                        # st.write('New SMILES molecule:')
                        imgByteArr = io.BytesIO()
                        MolToImage(new_mol,size=(400,400)).save(imgByteArr,format='PNG')
                        st.markdown("<p style='text-align: center;'>"+
                                    f"<img src='data:image/png;base64,{base64.b64encode(imgByteArr.getvalue()).decode()}' class='img-fluid'>"+
                                    "</p>", unsafe_allow_html=True)
                        
                        new_moses_passed = mol_passes_filters_custom(new_mol)
                        if new_moses_passed=='YES':
                            st.markdown("<p style='text-align: center; color: mediumseagreen'>MOSES filters passed successfully.</p>",unsafe_allow_html=True)
                        else:
                            st.markdown("<p style='text-align: center; color: red;'><b>MOSES filters passed failed.</b></p>",unsafe_allow_html=True)
                        st.session_state.new_logp,st.session_state.new_sa,st.session_state.new_cycle,st.session_state.new_pen_p = penalized_logp_standard(new_mol)
                    # st.write('New penalized logP score: %.5f' % (new_score))
                        col12, col22, col32, col42 = st.columns(4)
                        col12.metric('LogP', '%.5f' % (st.session_state.new_logp),'%.5f'%(st.session_state.new_logp-st.session_state.logp))
                        col22.metric('SA score', '%.5f' % (-st.session_state.new_sa),'%.5f'%(-st.session_state.new_sa+st.session_state.sa),delta_color='inverse')
                        col32.metric('Cycle score', '%d' % (-st.session_state.new_cycle),'%d'%(-st.session_state.new_cycle+st.session_state.cycle),delta_color='inverse')
                        col42.metric('Penalized LogP', '%.5f' % (st.session_state.new_pen_p),'%.5f'%(st.session_state.new_pen_p-st.session_state.pen_p))
                        # st.metric('New penalized logP score','%.5f' % (new_score), '%.5f'%(new_score-score))
                        st.metric('Similarity','%.5f' % (sim))
                        # st.write('Caching ZINC20 if necessary...')
                        with st.spinner("Caching ZINC20 if necessary..."):
                            if buy(new_smiles, catalog='zinc20',canonicalize=True):
                                st.write('This molecule exists.')
                                st.markdown("<h3 style='text-align: center; color: darkturquoise;'><b>This molecule exists.</h3>",unsafe_allow_html=True)
                            else:
                                # st.write('THIS MOLECULE DOES NOT EXIST!')
                                st.markdown("<h3 style='text-align: center; color: mediumseagreen;'>THIS MOLECULE DOES NOT EXIST!</h3>",unsafe_allow_html=True)
                        st.markdown("<p style='text-align: center; color: grey;'>Checked using molbloom</p>",unsafe_allow_html=True)
        else: optim_single_con.empty()

def check_single(smiles):
    # render_view()
    st.session_state.smiles_selected = True
    check_single_con = st.empty()
    
    # optim = False
    with check_single_con.container():
        if len(smiles) == 0:
            st.session_state.checked_single = 'EnterError'
        else:
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                st.session_state.checked_single = 'MolError'
            else:
                st.session_state.smiles = Chem.MolToSmiles(mol)
                st.session_state.logp,st.session_state.sa,st.session_state.cycle,st.session_state.pen_p = penalized_logp_standard(mol)
                moses_passed = mol_passes_filters_custom(mol)
                st.session_state.checked_single = moses_passed


def optim_single(smiles,lr,sim_cutoff,n_iter):
    # render_view()
    
        
        model = load_model()

        new_smiles,sim = model.optimize(smiles, sim_cutoff=sim_cutoff, lr=lr, num_iter=n_iter)
        
        return new_smiles,sim

        
            
def Optimize_a_batch():
    st.empty()

def reset_oam_state():
    st.session_state.smiles_selected = False
    st.session_state.checked_single = 'NO'
    st.session_state.smiles_checked = False
    set_step(0)
    
def rerun():
    st.experimental_rerun()
def render_view():
    # render_sidebar(st.session_state.current_view,st.session_state.current_step)
    form_header()
    form_body()


render_view()