import gradio as gr
import openai
import requests
from transformers import pipeline
from langchain.llms import OpenAI as LangOpenAI
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from rdkit import Chem
from rdkit.Chem import AllChem, Draw
from rdkit.Chem.Draw import rdMolDraw2D
import base64
from io import BytesIO
import py3Dmol
import re

# === Advanced LLM Configs ===
openai.api_key = "your-openai-api-key"

# BioGPT and OpenAI Coordinated Agents
def get_literature_insights(disease, symptoms):
    bio_gpt = pipeline("text-generation", model="microsoft/BioGPT-Large")
    prompt = f"Recent drug research for {disease} with symptoms: {symptoms}."
    return bio_gpt(prompt, max_length=200)[0]['generated_text']

def get_openai_smiles(disease, symptoms):
    prompt = f"Suggest 3 valid, drug-like SMILES strings that can potentially treat {disease} (symptoms: {symptoms}). Return only SMILES strings separated by space."
    response = openai.Completion.create(
        engine="text-davinci-003",
        prompt=prompt,
        max_tokens=100
    )
    return response.choices[0].text.strip()

# === Multi-agent Molecular Pipeline ===
def drug_discovery(disease, symptoms):
    # Agent 1: Literature from BioGPT
    literature = get_literature_insights(disease, symptoms)

    # Agent 2: Molecule SMILES from OpenAI
    smiles_result = get_openai_smiles(disease, symptoms)
    smiles_matches = re.findall(r"(?<![A-Za-z0-9])[A-Za-z0-9@+\-\[\]\(\)=#$]{5,}(?![A-Za-z0-9])", smiles_result)

    smiles = None
    for match in smiles_matches:
        if Chem.MolFromSmiles(match):
            smiles = match
            break
    if not smiles:
        smiles = "C1=CC=CC=C1"  # fallback

    # RDKit 2D Drawing
    mol = Chem.MolFromSmiles(smiles)
    AllChem.Compute2DCoords(mol)
    drawer = rdMolDraw2D.MolDraw2DCairo(300, 300)
    drawer.DrawMolecule(mol)
    drawer.FinishDrawing()
    img_data = drawer.GetDrawingText()
    img_base64 = base64.b64encode(img_data).decode("utf-8")
    img_html = f'''<div style="text-align:center; margin-top: 10px; animation: fadeIn 2s ease-in-out;">
        <img src="data:image/png;base64,{img_base64}" alt="2D Molecule" 
             style="border-radius: 16px; box-shadow: 0 6px 20px rgba(0,255,255,0.3);">
        <div style='color: #eee; margin-top: 8px;'>💊 2D Drug Structure</div></div>'''

    # RDKit 3D Molecule + py3Dmol
    mol3d = Chem.AddHs(mol)
    AllChem.EmbedMolecule(mol3d)
    AllChem.UFFOptimizeMolecule(mol3d)
    molblock = Chem.MolToMolBlock(mol3d)
    viewer = py3Dmol.view(width=420, height=420)
    viewer.addModel(molblock, "mol")
    viewer.setStyle({"stick": {"colorscheme": "cyanCarbon"}})
    viewer.setBackgroundColor("black")
    viewer.zoomTo()
    viewer.spin(True)
    viewer_html_raw = viewer._make_html()
    viewer_html = f'''<div style="text-align:center; margin-top: 20px; animation: zoomIn 2s ease-in-out;">
        <iframe srcdoc="{viewer_html_raw.replace('"', '&quot;')}" width="440" height="440" frameborder="0"></iframe>
        <div style='color: #eee; margin-top: 8px;'>🧬 3D Molecule</div></div>'''

    return literature, smiles, img_html, viewer_html

# === Gradio UI ===
disease_input = gr.Textbox(label="🧬 Disease (e.g., glioblastoma)", value="glioblastoma")
symptom_input = gr.Textbox(label="🩸 Symptoms (e.g., seizures, nausea)", value="seizures, nausea")
lit_output = gr.Textbox(label="📚 Literature from BioGPT")
smiles_output = gr.Textbox(label="🧪 SMILES Representation")
img_output = gr.HTML(label="🔬 2D Structure")
viewer_output = gr.HTML(label="🧬 3D Molecule")

custom_css = """
@keyframes fadeIn {
    from {opacity: 0;}
    to {opacity: 1;}
}
@keyframes zoomIn {
    from {transform: scale(0.5); opacity: 0;}
    to {transform: scale(1); opacity: 1;}
}
body {
    background: linear-gradient(to right, #141e30, #243b55);
    color: #ffffff;
    font-family: 'Segoe UI', sans-serif;
}
.gradio-container {
    animation: fadeIn 2s ease-in-out;
}
"""

iface = gr.Interface(
    fn=drug_discovery,
    inputs=[disease_input, symptom_input],
    outputs=[lit_output, smiles_output, img_output, viewer_output],
    title="💊 AI-Powered Drug Discovery System (CS3235 Project)",
    description="This real-time LLM-based platform suggests drugs for diseases without known treatments, generates 2D/3D molecules, and provides literature justifications using BioGPT + OpenAI + RDKit.",
    css=custom_css
)

iface.launch(share=True)