"""DPP-4 pIC50 Inference Web App — Gradio UI, models loaded from Hugging Face.
Deploy: Hugging Face Space (Gradio) AND GitHub Pages link to the Space.
This app is prepared now; it is TESTED only after training finishes,
once the HF repos below actually exist.

Models (HF repos, loaded on demand by model name):
  - extra-trees      Delower/dpp4-pic50-extra-trees
  - random-forest    Delower/dpp4-pic50-random-forest
  - xgboost          Delower/dpp4-pic50-xgboost
  - dnn              Delower/dpp4-pic50-dnn
  - ltn              Delower/dpp4-pic50-ltn        (NeSy / LTN)
  - gat              Delower/dpp4-pic50-gat        (graph)
  - transformer      Delower/dpp4-pic50-transformer (SMILES Transformer)
  - chemberta        Delower/dpp4-pic50-chemberta  (LLM / ChemBERTa)

Features: single SMILES input, or batch CSV upload (smiles column,
optional compound_id). Preview table, predicted pIC50 / IC50 / class,
download CSV, and two plots (pIC50 distribution, predicted vs wet-lab
for the FDA-12) downloadable as PNG.
"""
import io, os, tempfile
import numpy as np
import pandas as pd
import gradio as gr
import spaces

@spaces.GPU
def _zerogpu_dummy():
    # Dummy GPU function so the Space starts on free ZeroGPU hardware.
    # Never called; prediction code runs on CPU as before.
    return None

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

MODEL_REPOS = {
    "ExtraTrees (best)": "Delower/dpp4-pic50-extra-trees",
    "RandomForest": "Delower/dpp4-pic50-random-forest",
    "XGBoost (tuned)": "Delower/dpp4-pic50-xgboost",
    "DNN": "Delower/dpp4-pic50-dnn",
    "LTN (NeSy)": "Delower/dpp4-pic50-ltn",
    "GAT (graph)": "Delower/dpp4-pic50-gat",
    "Transformer": "Delower/dpp4-pic50-transformer",
    "ChemBERTa (LLM)": "Delower/dpp4-pic50-chemberta",
}
_cache = {}

def load_model(name):
    """Load a model from its HF repo on first use, cache after."""
    if name in _cache:
        return _cache[name]
    from huggingface_hub import hf_hub_download
    import pickle
    repo = MODEL_REPOS[name]
    # model file name is <model>_<features>.pkl, recorded in config.json
    from huggingface_hub import list_repo_files
    files = list_repo_files(repo, repo_type="model")
    pkl = [f for f in files if f.endswith(".pkl")][0]
    path = hf_hub_download(repo_id=repo, filename=pkl, repo_type="model")
    with open(path, "rb") as f:
        model = pickle.load(f)
    _cache[name] = model
    return model

def _ltn_weights():
    """Tuned LTN weights (npz) from Delower/dpp4-pic50-ltn, cached."""
    if "ltn_weights" in _cache:
        return _cache["ltn_weights"]
    from huggingface_hub import hf_hub_download
    path = hf_hub_download(repo_id="Delower/dpp4-pic50-ltn",
                           filename="ltn_model.npz", repo_type="model")
    z = np.load(path)
    keys = sorted((k for k in z.files if k.startswith("w")),
                  key=lambda k: int(k[1:]))
    w = [z[k].astype(np.float64) for k in keys]
    _cache["ltn_weights"] = w
    return w

def featurise_4608(smiles_list):
    """Morgan1024 + RDKitFP1024 + Morgan512 + Morgan2048 (training order)."""
    from rdkit import Chem, DataStructs
    from rdkit.Chem import AllChem
    def bits(fp, n):
        a = np.zeros(n, dtype=np.float32)
        DataStructs.ConvertToNumpyArray(fp, a)
        return a
    rows = []
    for s in smiles_list:
        mol = Chem.MolFromSmiles(s)
        if mol is None:
            rows.append(np.zeros(4608, dtype=np.float32))
            continue
        rows.append(np.concatenate([
            bits(AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=1024), 1024),
            bits(Chem.RDKFingerprint(mol, fpSize=1024), 1024),
            bits(AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=512), 512),
            bits(AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=2048), 2048)]))
    return np.vstack(rows).astype(np.float64)

def predict_ltn(df):
    """Tuned LTN (768-384-192, ELU) forward pass in pure NumPy — the
    Space needs no TensorFlow. Verified vs the training run (<1e-5)."""
    w = _ltn_weights()
    h = featurise_4608(list(df["smiles"]))
    for i in range(0, len(w) - 2, 2):
        h = h @ w[i] + w[i + 1]
        h = np.where(h > 0, h, np.exp(h) - 1.0)  # ELU
    pred = (h @ w[-2] + w[-1]).ravel()
    df = df.copy()
    df["predicted_pIC50"] = pred
    df["predicted_IC50_nM"] = 10 ** (9 - df["predicted_pIC50"])
    df["predicted_class"] = np.where(df["predicted_pIC50"] >= 7, "potent",
        np.where(df["predicted_pIC50"] >= 6, "moderate", "weak"))
    return df

def predict_df(df, model_name):
    if model_name == "LTN (NeSy)":
        return predict_ltn(df)
    model = load_model(model_name)
    # featurise with the model's own feature code (Morgan+physchem
    # for tree models; the model pipeline carries its preprocessor)
    from rdkit import Chem
    from rdkit.Chem import AllChem, Descriptors
    X = []
    for s in df["smiles"]:
        mol = Chem.MolFromSmiles(s)
        fp = list(AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=1024))
        phys = [Descriptors.MolWt(mol), Descriptors.MolLogP(mol),
                Descriptors.TPSA(mol), Descriptors.NumHDonors(mol),
                Descriptors.NumHAcceptors(mol),
                Descriptors.NumRotatableBonds(mol)]
        X.append(fp + phys)
    df = df.copy()
    Xdf = pd.DataFrame(X, columns=[f"feature_{i}" for i in range(1030)])
    df["predicted_pIC50"] = model.predict(Xdf)
    df["predicted_IC50_nM"] = 10 ** (9 - df["predicted_pIC50"])
    df["predicted_class"] = np.where(df["predicted_pIC50"] >= 7, "potent",
        np.where(df["predicted_pIC50"] >= 6, "moderate", "weak"))
    return df

def plot_dist(df):
    fig, ax = plt.subplots()
    ax.hist(df["predicted_pIC50"], bins=20)
    ax.set_xlabel("Predicted pIC50"); ax.set_ylabel("Count")
    ax.set_title("Predicted pIC50 distribution")
    return fig

def plot_fda(df):
    # predicted vs wet-lab, only for rows that have a wet-lab value
    if "wetlab_pIC50" not in df.columns or df["wetlab_pIC50"].dropna().empty:
        return None
    sub = df.dropna(subset=["wetlab_pIC50"])
    fig, ax = plt.subplots()
    ax.scatter(sub["wetlab_pIC50"], sub["predicted_pIC50"])
    ax.plot([4, 10], [4, 10], "--")
    ax.set_xlabel("Wet-lab pIC50"); ax.set_ylabel("Predicted pIC50")
    ax.set_title("Predicted vs wet-lab pIC50 (FDA-12)")
    return fig

def smiles_to_2d(smiles):
    """Return a 2D depiction image (PIL) of the SMILES, or None if invalid."""
    from rdkit import Chem
    from rdkit.Chem import Draw, AllChem
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None
    AllChem.Compute2DCoords(mol)
    return Draw.MolToImage(mol, size=(400, 300))

def single(smiles, model_name):
    df = pd.DataFrame({"compound_id": ["input"], "smiles": [smiles]})
    out = predict_df(df, model_name)
    return out, smiles_to_2d(smiles), plot_dist(out), plot_fda(out)

def batch(file, model_name):
    df = pd.read_csv(file.name)
    out = predict_df(df, model_name)
    fd, path = tempfile.mkstemp(prefix="dpp4_predictions_", suffix=".csv")
    with os.fdopen(fd, "w") as fh:
        out.to_csv(fh, index=False)
    return out, plot_dist(out), plot_fda(out), path

with gr.Blocks(title="Computational Chemistry and Drug Discovery — DPP-4 pIC50 Predictor") as demo:
    gr.Markdown("# Computational Chemistry and Drug Discovery\n## DPP-4 pIC50 Predictor\nSMILES in → predicted pIC50 / IC50. Models are loaded from Hugging Face on demand. Not for clinical use.")
    gr.Markdown("---")
    gr.Markdown("© 2026 Delower Hossain — Computational Chemistry and Drug Discovery. All rights reserved. | Models are in-silico predictions only, not for clinical use.")
    model = gr.Dropdown(list(MODEL_REPOS), value="ExtraTrees (best)", label="Model")
    with gr.Tab("Single SMILES"):
        smi = gr.Textbox(label="SMILES", placeholder="CC(=O)Oc1ccccc1C(=O)O")
        btn = gr.Button("Predict")
        out1 = gr.Dataframe(label="Prediction preview")
        img1 = gr.Image(label="2D structure")
        p1 = gr.Plot(label="pIC50 distribution")
        p2 = gr.Plot(label="Predicted vs wet-lab (FDA-12, if applicable)")
        btn.click(single, [smi, model], [out1, img1, p1, p2])
    with gr.Tab("Batch CSV"):
        f = gr.File(label="CSV with a 'smiles' column (optional 'compound_id')")
        btn2 = gr.Button("Predict batch")
        out2 = gr.Dataframe(label="Prediction preview")
        q1 = gr.Plot(label="pIC50 distribution")
        q2 = gr.Plot(label="Predicted vs wet-lab (FDA-12, if applicable)")
        dl = gr.DownloadButton(label="Download predictions (CSV)")
        btn2.click(batch, [f, model], [out2, q1, q2, dl])

    with gr.Tab("Computational Experiment Results"):
        gr.Markdown("FDA-12 inference with ExtraTrees (original LazyPredict model, RMSE 0.4449). Wet-lab values are comparison-only. RMSE 0.9637, MAE 0.6794 over the 11 drugs with wet-lab values (Gosogliptin excluded).")
        fda_btn = gr.Button("Show FDA-12 results as dataframe")
        fda_out = gr.Dataframe(label="FDA-12 results")
        def show_fda():
            import pandas as pd
            return pd.read_csv("fda.csv") if __import__("os").path.exists("fda.csv") else pd.DataFrame({"note": ["fda.csv is in the Hugging Face repo Delower/dpp4-pic50-extra-trees"]})
        fda_btn.click(show_fda, [], [fda_out])

    gr.Markdown("---\n© 2026 Delower Hossain — Computational Chemistry and Drug Discovery. All rights reserved.")

if __name__ == "__main__":
    demo.launch()