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")

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

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;
}
'''

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

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

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

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

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)
st.markdown('<style>' + open('./style.css').read() + '</style>', unsafe_allow_html=True)

with st.sidebar:
    # st.header('+')
    st.markdown("<h5 style='text-align: center; color:grey;'>Explore</h5>",unsafe_allow_html=True)
    tabs = on_hover_tabs(tabName=['Optimize a molecule', 'Optimize batch', 'About'], 
                         iconName=['science', 'batch_prediction', 'info'], default_choice=0)

if tabs == '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)
elif tabs == 'Optimize a molecule':
    st.markdown("<h2 style='text-align: center;'>Optimize a molecule</h2>",unsafe_allow_html=True)
    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')

    mol = Chem.MolFromSmiles(st.session_state.smiles)
    if mol is None:
        st.markdown("<p style='text-align: center; color: red;'>SMILES is invalid. Please enter a valid SMILES.</p>",unsafe_allow_html=True)
    else:
        score = penalized_logp_standard(mol)
        # with st.columns(3)[1]:
        # st.markdown("<style>{text-align: center;}</style>",unsafe_allow_html=True)
        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)))
        st.metric('Penalized logP score', '%.5f' % (score))

    if mol is not None:
        # col1, col2, col3 = st.columns(3)
        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')
            
            # with st.columns(3)[1]:
                with st.spinner("Operation in progress. Please wait."):

                    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.markdown("<p style='text-align: center; color: red;'>Cannot optimize! Please choose another setting.</p>",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_score = penalized_logp_standard(new_mol)
                    # st.write('New penalized logP score: %.5f' % (new_score))
                        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: cyan;'><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: lightgreen;'>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)
elif tabs == 'Optimize batch':
    st.write('Incoming...')