import gradio as gr
from rdkit import Chem
from rdkit.Chem import Descriptors, Draw
import cirpy


def _mol_from_smiles(smiles: str):
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        raise gr.Error("Invalid SMILES string.")
    return mol


def smiles_to_canonical(smiles: str) -> str:
    mol = _mol_from_smiles(smiles)
    return Chem.MolToSmiles(mol)


def molecular_weight(smiles: str) -> float:
    mol = _mol_from_smiles(smiles)
    return float(Descriptors.MolWt(mol))


def logp(smiles: str) -> float:
    mol = _mol_from_smiles(smiles)
    return float(Descriptors.MolLogP(mol))


def tpsa(smiles: str) -> float:
    mol = _mol_from_smiles(smiles)
    return float(Descriptors.TPSA(mol))


def mol_image(smiles: str):
    mol = _mol_from_smiles(smiles)
    return Draw.MolToImage(mol)


def name_to_smiles(name: str) -> str:
    """Convert chemical name to SMILES using Chemical Identifier Resolver (CIR)"""
    try:
        smiles = cirpy.resolve(name, 'smiles')
        if smiles is None:
            raise gr.Error(f"Could not find SMILES for chemical name: {name}")
        return smiles
    except Exception as e:
        raise gr.Error(f"Error converting name to SMILES: {str(e)}")


smiles_interface = gr.Interface(
    fn=smiles_to_canonical,
    inputs=gr.Textbox(label="SMILES"),
    outputs=gr.Textbox(label="Canonical SMILES"),
    api_name="smiles_to_mol",
    description="Convert an input SMILES string to its canonical form.",
)

name_interface = gr.Interface(
    fn=name_to_smiles,
    inputs=gr.Textbox(label="Chemical Name", placeholder="e.g., aspirin, caffeine, benzene"),
    outputs=gr.Textbox(label="SMILES"),
    api_name="name_to_smiles",
    description="Convert a chemical name to SMILES notation.",
    examples=[["aspirin"], ["caffeine"], ["benzene"], ["ethanol"]],
)

mw_interface = gr.Interface(
    fn=molecular_weight,
    inputs=gr.Textbox(label="SMILES"),
    outputs=gr.Number(label="Molecular Weight (g/mol)"),
    api_name="molecular_weight",
    description="Compute the molecular weight from a SMILES string.",
)

logp_interface = gr.Interface(
    fn=logp,
    inputs=gr.Textbox(label="SMILES"),
    outputs=gr.Number(label="logP"),
    api_name="logp",
    description="Calculate the octanol/water partition coefficient (logP).",
)

tpsa_interface = gr.Interface(
    fn=tpsa,
    inputs=gr.Textbox(label="SMILES"),
    outputs=gr.Number(label="TPSA"),
    api_name="tpsa",
    description="Calculate the topological polar surface area (TPSA).",
)


demo = gr.TabbedInterface(
    [name_interface, smiles_interface, mw_interface, logp_interface, tpsa_interface],
    [
        "Name to SMILES",
        "SMILES to Canonical",
        "Molecular Weight",
        "LogP",
        "TPSA",
    ],
    title="RDKit API",
    css=".gradio-container {max-width: 800px; margin: auto;}",
)


if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)