import torch
import streamlit as st
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

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


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

st.header('Junction Tree Variational Autoencoder for Molecular Graph Generation (JTVAE)')
st.subheader('Wengong Jin, Regina Barzilay, Tommi Jaakkola')
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)'''

with st.expander('About'):
    st.markdown(descrip)

st.text_input('Enter a SMILES string:','CNC(=O)C1=NC=CC(=C1)OC2=CC=C(C=C2)NC(=O)NC3=CC(=C(C=C3)Cl)C(F)(F)F',key='smiles')

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 standardized_log_p + standardized_SA + standardized_cycle

mol = Chem.MolFromSmiles(st.session_state.smiles)
if mol is None:
    st.write('SMILES is invalid. Please enter a valid SMILES.')
else:
    st.write('Molecule:')
    st.image(MolToImage(mol,size=(300,300)))
    score = penalized_logp_standard(mol)
    st.write('Penalized logP score: %.5f' % (score))

if mol is not None:
    st.slider('Choose learning rate: ',0.0,10.0,0.4,key='lr')
    st.slider('Choose similarity cutoff: ',0.0,3.0,0.4,key='sim_cutoff')
    st.slider('Choose number of iterations: ',1,100,80,key='n_iter')
    vocab = [x.strip("\r\n ") for x in open('./vocab.txt')] 
    vocab = Vocab(vocab)
    if st.button('Optimize'):
        st.write('Testing')

        model = JTPropVAE(vocab, 450, 56, 20, 3)

        model.load_state_dict(torch.load('./model.iter-685000',map_location=torch.device('cpu')))
        
        new_smiles,sim = model.optimize(st.session_state.smiles, sim_cutoff=st.session_state.sim_cutoff, lr=st.session_state.lr, num_iter=st.session_state.n_iter)
        
        del model
        if new_smiles is None:
            st.write('Cannot optimize.')
        else:
            st.write('New SMILES:')
            st.code(new_smiles)
            new_mol = Chem.MolFromSmiles(new_smiles)
            if new_mol is None:
                st.write('New SMILES is invalid.')
            else:
                st.write('New SMILES molecule:')
                st.image(MolToImage(new_mol,size=(300,300)))
                new_score = penalized_logp_standard(new_mol)
                st.write('New penalized logP score: %.5f' % (new_score))
                st.write('Caching ZINC20 if necessary...')
                if buy(new_smiles, catalog='zinc20',canonicalize=True):
                    st.write('This molecule exists.')
                    st.caption('Checked by molbloom.')
                else:
                    st.write('THIS MOLECULE DOES NOT EXIST!')
                    st.caption('Checked by molbloom.')