import streamlit as st
import py3Dmol
from rdkit import Chem
from rdkit.Chem import AllChem, Draw
import streamlit.components.v1 as components

# --- Helper Functions ---

def smi_to_3d_mol(smiles):
    """
    Converts a SMILES string to an RDKit Mol object with 3D coordinates.
    """
    try:
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return None
        mol = Chem.AddHs(mol)
        AllChem.EmbedMolecule(mol)
        AllChem.MMFFOptimizeMolecule(mol, maxIters=200)
        return mol
    except Exception as e:
        st.error(f"Error processing SMILES: {e}")
        return None

def render_3d_viewer(mol, style="stick", size=(500, 500), surface=False, opacity=0.5):
    """
    Creates a py3Dmol viewer and returns its HTML representation for Streamlit.
    """
    assert style in ('line', 'stick', 'sphere', 'carton')
    mblock = Chem.MolToMolBlock(mol)
    viewer = py3Dmol.view(width=size[0], height=size[1])
    viewer.addModel(mblock, 'mol')
    viewer.setStyle({style:{}})
    if surface:
        viewer.addSurface(py3Dmol.SAS, {'opacity': opacity})
    viewer.zoomTo()
    # Return the viewer's HTML/JS representation for Streamlit
    return viewer.js()

def render_2d_image(name, ingredients):
    """
    Displays the name and a 2D image of a molecule in a Streamlit column.
    """
    mol = Chem.MolFromSmiles(ingredients)
    if mol:
        st.write(f"**{name}**")
        img = Draw.MolToImage(mol)
        st.image(img, caption=ingredients)
    else:
        st.warning(f"Could not generate molecule for '{name}'.")


# --- Main Streamlit App ---

st.set_page_config(layout="wide")
st.title('⚛️🧬 Molecule Modeler 🧬⚛️')

# --- Interactive 3D Viewer Section ---
st.header("Interactive 3D Molecule Viewer")
# Create two columns for the input and the viewer
c1, c2 = st.columns([1, 2]) # Input column is 1/3 of the width

with c1:
    st.write('Enter a SMILES string below to generate a 3D model.')
    compound_smiles = st.text_input(
        'SMILES Input',
        'CC(=O)Oc1ccccc1C(=O)O' # Default: Aspirin
    )
    style = st.selectbox('Select Viewer Style', ['stick', 'sphere', 'line', 'carton'])
    show_surface = st.checkbox('Show Molecular Surface')
    opacity = 0.5
    if show_surface:
        opacity = st.slider('Surface Opacity', 0.1, 1.0, 0.5)

    st.write('**What is SMILES?**')
    st.info('''
    The **Simplified Molecular-Input Line-Entry System (SMILES)** is a text-based representation of a molecule's structure.
    - [Learn more on Wikipedia](https://en.wikipedia.org/wiki/Simplified_molecular-input_line-entry_system)
    - [Find compounds on PubChem](https://pubchem.ncbi.nlm.nih.gov/).
    ''')


# Process input and display the 3D viewer in the second column
with c2:
    if compound_smiles:
        mol_3d = smi_to_3d_mol(compound_smiles)
        if mol_3d:
            st.write('**3D Chemical Structure**')
            viewer_html = render_3d_viewer(mol_3d, style=style, surface=show_surface, opacity=opacity)
            components.html(viewer_html, height=500, width=500)
        else:
            st.error("Invalid SMILES string. Could not generate the molecule.")
    else:
        st.info("Enter a SMILES string to generate a 3D model.")


st.divider()

# --- Predefined Molecules Section ---
st.header("Molecule Examples (2D)")

col1, col2, col3 = st.columns(3)

with col1:
    st.subheader("Common Molecules")
    render_2d_image("Ibuprofen", "CC(C)CC1=CC=C(C=C1)C(C)C(=O)O")
    render_2d_image("Caffeine", "CN1C=NC2=C1C(=O)N(C(=O)N2C)C")
    render_2d_image("Ethanol", "CCO")
    render_2d_image("Glucose", "OC[C@@H](O1)[C@@H](O)[C@H](O)[C@@H](O)[C@H](O)1")

with col2:
    st.subheader("Vitamins")
    render_2d_image("Thiamine (B1)", "OCCc1c(C)[n+](cs1)Cc2cnc(C)nc2N")
    render_2d_image("Vitamin E", "CC(C)CCC[C@@H](C)CCC[C@@H](C)CCC[C@]1(C)CCc2c(C)c(O)c(C)c(C)c2O1")
    render_2d_image("Vitamin D3", "C[C@@H](CCCC(C)C)[C@H]1CC[C@H]2/C(=C/C=C/3\C(=C)CCC[C@H]3O)/CC[C@]12C")

with col3:
    st.subheader("COVID-19 Antivirals")
    render_2d_image("Favipiravir", "C1=C(N=C(C(=O)N1)C(=O)N)F")
    render_2d_image("Remdesivir", "CCC(CC)COC(=O)[C@H](C)N[P@](=O)(OC[C@@H]1[C@H]([C@H]([C@](O1)(C#N)C2=CC=C3N2N=CN=C3N)O)O)OC4=CC=CC=C4")
    render_2d_image("Ritonavir", "CC(C)C1=NC(=CS1)CN(C)C(=O)N[C@@H](C(C)C)C(=O)N[C@@H](CC2=CC=CC=C2)C[C@@H]([C@H](CC3=CC=CC=C3)NC(=O)OCC4=CN=CS4)O")

st.divider()

# --- Markdown Section with SMILES Lists ---
st.header("Drug Information & SMILES Strings")

markDown = '''
### Opioids
* **Fentanyl:** `CCC(=O)N(C1CCN(CC1)CCC2=CC=CC=C2)C3=CC=CC=C3`
* **Hydrocodone:** `CN1CC[C@]23[C@@H]4[C@H]1CC5=C2C(=C(C=C5)OC)O[C@H]3C(=O)CC4`
* **Oxycodone:** `CN1CC[C@]23[C@@H]4C(=O)CC[C@]2([C@H]1CC5=C3C(=C(C=C5)OC)O4)O`

### CNS Depressants
* **Diazepam (Valium):** `CN1C(=O)CN=C(C2=C1C=CC(=C2)Cl)C3=CC=CC=C3`
* **Alprazolam (Xanax):** `CC1=NN=C2N1C3=C(C=C(C=C3)Cl)C(=NC2)C4=CC=CC=C4`
* **Zolpidem (Ambien):** `CC1=CC=C(C=C1)C2=C(N3C=C(C=CC3=N2)C)CC(=O)N(C)C`

### Stimulants
* **Dextroamphetamine:** `C[C@@H](CC1=CC=CC=C1)N`
* **Methylphenidate (Ritalin):** `COC(=O)C(C1CCCCN1)C2=CC=CC=C2`
* **Amphetamine (base):** `CC(N)CC1=CC=CC=C1`
'''
st.markdown(markDown)