import streamlit as st
from rdkit import Chem
from rdkit.Chem import AllChem, Draw

def visualize_molecule(mol_smiles):

    mol = Chem.MolFromSmiles(mol_smiles)

    mol = Chem.AddHs(mol)

    AllChem.Compute2DCoords(mol)

    return mol

# Streamlit App
st.title('Molecular Structure Visualization')
st.write("By Jishnu Setia")

# Input molecule SMILES string
mol_smiles = st.text_input('Enter molecule SMILES:', 'CCl.O.Br.Cl.C=C.C#C.[Na]Br.[Na]Cl.[Na+].[Br-].c1ccccc1.c1c(O)c(I)c(F)c(Br)c1(Cl)')

if mol_smiles:
    mol = visualize_molecule(mol_smiles)
    img = Draw.MolToImage(mol, size=(700, 400))
    st.image(img)
    st.write(" ")
    
    # Instructions for SMILES 
    st.markdown("<h2>SMILES (Simplified Molecular Input Line Entry System)</h2>", unsafe_allow_html=True)
    st.write("SMILES is a string format for representing molecular structures.")
    st.write("You can represent various features using specific characters in SMILES notation:")
    st.write("- Single bonds: Use '-' between atoms (e.g., 'CC' for ethane)")
    st.write("- Double bonds: Use '=' between atoms (e.g., 'C=C' for ethene)")
    st.write("- Triple bonds: Use '#' between atoms (e.g., 'C#C' for ethyne)")
    st.write("- Benzene: Use 'c1ccccc1' (or 'C1=CC=CC=C1') for a benzene ring")
    st.write("- Compounds: Combine individual SMILES strings with '.' (e.g., 'CCO.Cl' for ethanol and chlorine)")
    st.write("- Ionic compounds: Use square brackets (e.g., '[Na]Br' for sodium bromide)")
    st.write("- Ions only: Use the SMILES notation of the ion (e.g., '[Na+]' for sodium ion)")