{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OmniBiMol Variant Priority Pipeline — Walkthrough\n",
    "\n",
    "This notebook demonstrates the full variant-to-therapy workflow using only free-tier Hugging Face resources.\n",
    "\n",
    "**Repository:** [omshrivastava/omnibimol-variant-priority](https://huggingface.co/omshrivastava/omnibimol-variant-priority)\n",
    "\n",
    "**Best model AUROC:** 0.982 (XGBoost on precomputed scores)\n",
    "\n",
    "**Hardware:** CPU-only (free-tier sandbox)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q datasets xgboost scikit-learn pandas matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Load public datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "import pandas as pd\n",
    "\n",
    "# Variant pathogenicity data with precomputed scores\n",
    "ds_scores = load_dataset(\"songlab/clinvar\", split=\"test\")\n",
    "df_scores = ds_scores.to_pandas()\n",
    "print(f\"Variants: {len(df_scores)}\")\n",
    "print(df_scores.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Protein feature data (fallback when precomputed scores unavailable)\n",
    "ds_prot_test = load_dataset(\"Rain021217/clinvar-pathogenicity-prediction-dataset\", split=\"test\")\n",
    "df_prot_test = ds_prot_test.to_pandas()\n",
    "print(f\"Protein variants: {len(df_prot_test)}\")\n",
    "print(df_prot_test.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Gene-disease evidence from OpenTargets\n",
    "ds_ot = load_dataset(\"opentargets/clinical_evidence\", split=\"train\")\n",
    "df_ot = ds_ot.to_pandas()\n",
    "print(f\"Clinical evidence rows: {len(df_ot)}\")\n",
    "print(df_ot.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Train XGBoost model (best AUROC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.impute import SimpleImputer\n",
    "import xgboost as xgb\n",
    "from sklearn.metrics import roc_auc_score, f1_score, average_precision_score, precision_recall_curve\n",
    "\n",
    "score_cols = [c for c in df_scores.columns if c not in [\"chrom\", \"pos\", \"ref\", \"alt\", \"label\"]]\n",
    "X = df_scores[score_cols]\n",
    "y = df_scores[\"label\"].astype(int).values\n",
    "\n",
    "imputer = SimpleImputer(strategy=\"median\")\n",
    "scaler = StandardScaler()\n",
    "X_imp = imputer.fit_transform(X)\n",
    "X_scaled = scaler.fit_transform(X_imp)\n",
    "\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    X_scaled, y, test_size=0.2, random_state=42, stratify=y\n",
    ")\n",
    "\n",
    "clf = xgb.XGBClassifier(\n",
    "    n_estimators=200, max_depth=4, learning_rate=0.05,\n",
    "    subsample=0.8, colsample_bytree=0.8, eval_metric=\"logloss\",\n",
    "    random_state=42, n_jobs=4\n",
    ")\n",
    "clf.fit(X_train, y_train)\n",
    "\n",
    "y_proba = clf.predict_proba(X_test)[:, 1]\n",
    "print(f\"AUROC: {roc_auc_score(y_test, y_proba):.4f}\")\n",
    "print(f\"F1:    {f1_score(y_test, (y_proba >= 0.5).astype(int)):.4f}\")\n",
    "print(f\"AP:    {average_precision_score(y_test, y_proba):.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Plot ROC and PR curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import roc_curve, precision_recall_curve\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
    "\n",
    "fpr, tpr, _ = roc_curve(y_test, y_proba)\n",
    "axes[0].plot(fpr, tpr, lw=2)\n",
    "axes[0].plot([0, 1], [0, 1], \"k--\", alpha=0.5)\n",
    "axes[0].set_xlabel(\"FPR\")\n",
    "axes[0].set_ylabel(\"TPR\")\n",
    "axes[0].set_title(f\"ROC (AUC={roc_auc_score(y_test, y_proba):.3f})\")\n",
    "axes[0].grid(True, alpha=0.3)\n",
    "\n",
    "precision, recall, _ = precision_recall_curve(y_test, y_proba)\n",
    "ap = average_precision_score(y_test, y_proba)\n",
    "axes[1].plot(recall, precision, lw=2)\n",
    "axes[1].set_xlabel(\"Recall\")\n",
    "axes[1].set_ylabel(\"Precision\")\n",
    "axes[1].set_title(f\"PR (AP={ap:.3f})\")\n",
    "axes[1].grid(True, alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Feature importance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fi = pd.DataFrame({\n",
    "    \"feature\": score_cols,\n",
    "    \"importance\": clf.feature_importances_\n",
    "}).sort_values(\"importance\", ascending=True)\n",
    "\n",
    "plt.figure(figsize=(8, 5))\n",
    "plt.barh(fi[\"feature\"], fi[\"importance\"])\n",
    "plt.xlabel(\"Importance\")\n",
    "plt.title(\"XGBoost Feature Importance\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Single-variant inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example variant: input precomputed scores\n",
    "sample = np.array([[1.5, 0.3, 0.7, -1.2, 0.1, -2.5, -0.8, 0.0]])\n",
    "sample_imp = imputer.transform(sample)\n",
    "sample_scaled = scaler.transform(sample_imp)\n",
    "proba = clf.predict_proba(sample_scaled)[0, 1]\n",
    "\n",
    "tier = \"Tier 1\" if proba >= 0.9 else \"Tier 2\" if proba >= 0.7 else \"Tier 3\" if proba >= 0.5 else \"Tier 4\"\n",
    "print(f\"Pathogenicity: {proba:.4f} → {tier}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Gene-disease evidence lookup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build gene-disease association scores from OpenTargets\n",
    "evidence = df_ot.groupby([\"targetId\", \"diseaseId\"]).agg(\n",
    "    max_phase=(\"clinicalPhase\", lambda x: x.fillna(0).max()),\n",
    "    n_trials=(\"nctid\", \"nunique\"),\n",
    "    n_datasources=(\"datasourceId\", \"nunique\")\n",
    ").reset_index()\n",
    "\n",
    "evidence[\"evidence_score\"] = (\n",
    "    0.4 * (evidence[\"max_phase\"] / 4.0) +\n",
    "    0.4 * np.minimum(evidence[\"n_trials\"] / 10.0, 1.0) +\n",
    "    0.2 * np.minimum(evidence[\"n_datasources\"] / 5.0, 1.0)\n",
    ")\n",
    "\n",
    "print(f\"Unique gene-disease pairs: {len(evidence)}\")\n",
    "print(evidence.nlargest(5, \"evidence_score\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Save artifacts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle, json\n",
    "\n",
    "clf.save_model(\"xgb_precomputed.json\")\n",
    "with open(\"imputer.pkl\", \"wb\") as f: pickle.dump(imputer, f)\n",
    "with open(\"scaler.pkl\", \"wb\") as f: pickle.dump(scaler, f)\n",
    "\n",
    "metrics = {\n",
    "    \"AUROC\": float(roc_auc_score(y_test, y_proba)),\n",
    "    \"F1\": float(f1_score(y_test, (y_proba >= 0.5).astype(int))),\n",
    "    \"AP\": float(average_precision_score(y_test, y_proba))\n",
    "}\n",
    "with open(\"metrics.json\", \"w\") as f:\n",
    "    json.dump(metrics, f, indent=2)\n",
    "\n",
    "print(\"Artifacts saved!\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}