import gradio as gr
from transformers import pipeline
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem.Draw import rdMolDraw2D
import base64
import re
import py3Dmol
import time

# Load model once
bio_gpt = pipeline("text-generation", model="microsoft/BioGPT-Large")

def drug_discovery(disease, symptoms):
    # Simplified and efficient medical prompt
    prompt = (
        f"You're a biomedical AI. A new disease shows symptoms: '{symptoms}'. "
        f"Suggest 5 generic drug names and 5 SMILES strings that could help treat this. "
        f"List drug names first, then SMILES strings in separate lines like:\n"
        f"Drugs: Aspirin, Ibuprofen, Paracetamol, ...\n"
        f"SMILES: C1=CC=CC=C1 C(C(=O)O)N ..."
    )

    try:
        start = time.time()
        result = bio_gpt(prompt, max_length=150, do_sample=True, temperature=0.6)[0]['generated_text']
    except Exception as e:
        return f"BioGPT error: {e}", "", "", ""

    # Extract drug names and SMILES
    drugs_match = re.search(r"Drugs:\s*(.+)", result)
    smiles_match = re.search(r"SMILES:\s*(.+)", result)

    drug_names = drugs_match.group(1).strip() if drugs_match else "Unknown"
    raw_smiles = smiles_match.group(1).strip() if smiles_match else "C1=CC=CC=C1"

    smiles_list = re.findall(r"(?<![A-Za-z0-9])[A-Za-z0-9@+\-\[\]\(\)=#$]{5,}(?![A-Za-z0-9])", raw_smiles)
    smiles_list = list({sm for sm in smiles_list if Chem.MolFromSmiles(sm)})[:3]

    if not smiles_list:
        smiles_list = ["C1=CC=CC=C1"]

    img_html, viewer_htmls = "", ""
    for smiles in smiles_list:
        mol = Chem.MolFromSmiles(smiles)
        AllChem.Compute2DCoords(mol)
        drawer = rdMolDraw2D.MolDraw2DCairo(250, 250)
        drawer.DrawMolecule(mol)
        drawer.FinishDrawing()
        img_data = drawer.GetDrawingText()
        img_base64 = base64.b64encode(img_data).decode("utf-8")
        img_html += f'''
        <div style="display:inline-block; margin:10px;">
            <img src="data:image/png;base64,{img_base64}" width="120" height="120">
            <p style="color:white; font-size:12px;">{smiles}</p>
        </div>'''

        # 3D View
        mol3d = Chem.AddHs(mol)
        AllChem.EmbedMolecule(mol3d, randomSeed=42)
        AllChem.UFFOptimizeMolecule(mol3d)
        mb = Chem.MolToMolBlock(mol3d)
        viewer = py3Dmol.view(width=240, height=240)
        viewer.addModel(mb, "mol")
        viewer.setStyle({"stick": {"colorscheme": "cyanCarbon"}})
        viewer.setBackgroundColor("black")
        viewer.zoomTo()
        viewer_html_raw = viewer._make_html()
        viewer_htmls += f'''
        <div style="display:inline-block; margin:10px;">
            <iframe srcdoc="{viewer_html_raw.replace('"', '&quot;')}" width="240" height="240" frameborder="0"></iframe>
        </div>'''

    duration = round(time.time() - start, 2)
    literature_summary = f"📋 Drug candidates (auto-generated in {duration}s):\n{drug_names}"

    return literature_summary, ", ".join(smiles_list), img_html, viewer_htmls

# Gradio UI setup

iface = gr.Interface(
    fn=drug_discovery,
    inputs=[
        gr.Textbox(label="🧬 Enter Unknown Disease or Name", value="X-disease"),
        gr.Textbox(label="📝 Symptoms", value="fever, joint pain")
    ],
    outputs=[
        gr.Textbox(label="🔖 AI Literature Summary"),
        gr.Textbox(label="🧪 SMILES List"),
        gr.HTML(label="🖼️ 2D Molecules"),
        gr.HTML(label="🔬 3D Molecules")
    ],
    title="🧪 Drug Discovery for Unknown Diseases",
    description="BioGPT + RDKit-powered system to suggest potential drug molecules for unknown or rare diseases.",
    css="""
        body { background-color: #111; color: #eee; }
        .gradio-container { animation: fadeIn 1.5s ease-in-out; }
    """
)

iface.launch(share=True)