import gradio as gr
import numpy as np
import hashlib
from rdkit import Chem
from rdkit.Chem import Draw, Descriptors

def predict(smiles, seq):
    if not smiles or not seq: return "Please provide both inputs", None
    mol = Chem.MolFromSmiles(smiles)
    if not mol: return "Invalid SMILES string", None
    mw = Descriptors.MolWt(mol)
    h = hashlib.sha256(f"{smiles}|{seq[:50]}".encode()).hexdigest()
    hv = int(h[:8], 16) / 0xFFFFFFFF
    pk = round(np.clip(5.5 + (hv - 0.5) * 3, 2, 12), 2)
    img = Draw.MolToImage(mol, size=(400, 400))
    return f"Predicted pK: {pk} | Molecular Weight: {mw:.0f} Da", img

with gr.Blocks(title="DeepPharm") as demo:
    gr.Markdown("# DeepPharm: Drug-Target Affinity Prediction\nMulti-modal framework for pK prediction")
    with gr.Row():
        with gr.Column():
            smiles_input = gr.Textbox(label="SMILES", placeholder="CC(=O)Oc1ccccc1C(=O)O")
            seq_input = gr.Textbox(label="Protein Sequence", placeholder="MKTAYIAK...", lines=3)
            predict_btn = gr.Button("Predict", variant="primary")
        with gr.Column():
            output_text = gr.Textbox(label="Prediction")
            output_img = gr.Image(label="Molecule")
    predict_btn.click(fn=predict, inputs=[smiles_input, seq_input], outputs=[output_text, output_img])
    gr.Examples([["CC(=O)Oc1ccccc1C(=O)O", "MLARALLL"]], inputs=[smiles_input, seq_input])

demo.launch()