{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Search PLINDER with your structures and sequences\n",
    "\n",
    "`search()` finds similar proteins, ligand pockets and protein interfaces in PLINDER. Supply an input table for ligand-containing structures, or a FASTA file or structure folder for ligand-free queries."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Input requirements\n",
    "\n",
    "Search and evaluation share a table with `input_id`, `structure_path`, and optional `ligand_path` columns. Evaluation adds `reference_id`. See the {ref}`input table reference <arrange-your-predictions>` for details.\n",
    "\n",
    "Use an assembled PDB/mmCIF receptor with one or more ligand SDFs, or a complete mmCIF containing the ligand coordinates and chemistry. The SDFs must use the receptor's coordinate frame.\n",
    "\n",
    "For ligand-free queries, pass a FASTA file, a structure file, or a folder of PDB/mmCIF structures directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "from plinder.core import PlinderSystem, PlinderRelease\n",
    "from plinder.core.scores import (\n",
    "    CustomProteinSearchConfig,\n",
    "    search,\n",
    ")\n",
    "\n",
    "release_path = os.environ.get(\"PLINDER_DATA_DIR\")\n",
    "release_path = Path(release_path) if release_path else None\n",
    "input_cif = (\n",
    "    Path(os.environ[\"PLINDER_CUSTOM_CIF\"])\n",
    "    if \"PLINDER_CUSTOM_CIF\" in os.environ\n",
    "    else PlinderSystem(\n",
    "        system_id=\"4agi__1__1.C__1.W\",\n",
    "        release=PlinderRelease(data_dir=release_path),\n",
    "    ).receptor_cif\n",
    ")\n",
    "work_dir = Path(\n",
    "    os.environ.get(\"PLINDER_CUSTOM_WORK_DIR\", \"custom_scoring_example\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run a small receptor-only example\n",
    "\n",
    "The example rebuilds a PLINDER receptor mmCIF when `PLINDER_CUSTOM_CIF` is not set. `plinder_entry_ids` limits the comparison to receptor and interface proteins from the selected entries; omit it to search all PLINDER proteins.\n",
    "\n",
    "`mode=\"pockets\"` searches protein and pocket features. Use `mode=\"interfaces\"` for protein interfaces, `mode=\"ligands\"` for ligand comparisons, or `mode=\"both\"` for both. The default `mode=\"auto\"` selects applicable features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "result = search(\n",
    "    input_cif,\n",
    "    output_dir=work_dir,\n",
    "    release=PlinderRelease(data_dir=release_path),\n",
    "    mode=\"pockets\",\n",
    "    backends=(\"mmseqs\",),\n",
    "    plinder_entry_ids=[\"4agi\"],\n",
    "    search_config=CustomProteinSearchConfig(max_seqs=25),\n",
    "    threads=2,\n",
    "    store_aligned_pocket_residues=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "protein_scores = pd.read_parquet(result.protein_scores)\n",
    "residue_pairs = pd.read_parquet(result.aligned_pocket_residues)\n",
    "{\n",
    "    \"protein_score_rows\": len(protein_scores),\n",
    "    \"aligned_pocket_residue_rows\": len(residue_pairs),\n",
    "    \"ligand_scores_written\": result.ligand_scores is not None,\n",
    "    \"interface_scores_written\": result.interface_scores is not None,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Protein scoring is oriented from each PLINDER receptor and ligand pocket to a custom protein chain. `pocket_fident` is therefore the percentage of PLINDER pocket residues that align to identical residues in the custom chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pocket_scores = protein_scores.loc[\n",
    "    protein_scores[\"metric\"].astype(str).eq(\"pocket_fident\"),\n",
    "    [\n",
    "        \"query_system\",\n",
    "        \"query_ligand_id\",\n",
    "        \"target_system\",\n",
    "        \"protein_mapping\",\n",
    "        \"source\",\n",
    "        \"similarity\",\n",
    "    ],\n",
    "].head()\n",
    "pocket_scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When aligned pocket residues are requested, each row records a PLINDER pocket residue and the custom residue selected by the chain mapping and search source used for the reported score."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "residue_pairs[[\n",
    "    \"plinder_system_id\",\n",
    "    \"plinder_chain_instance\",\n",
    "    \"plinder_residue_number\",\n",
    "    \"custom_structure_id\",\n",
    "    \"custom_chain_asym_id\",\n",
    "    \"custom_residue_number\",\n",
    "    \"residue_identical\",\n",
    "]].head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Search ligand poses\n",
    "\n",
    "Each SDF supplies its own ligand chemistry and coordinates. Put the receptor path and ligand path in the same input row, and repeat the input ID for additional ligands."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_ligand_model(receptor_path, ligand_path):\n",
    "    inputs = pd.DataFrame([{\n",
    "        \"input_id\": \"model_1\",\n",
    "        \"structure_path\": str(receptor_path),\n",
    "        \"ligand_path\": str(ligand_path),\n",
    "    }])\n",
    "    return search(\n",
    "        inputs, output_dir=\"ligand_search\", mode=\"ligands\",\n",
    "    )\n",
    "\n",
    "\n",
    "# Examples:\n",
    "# search_ligand_model(\"receptor.pdb\", \"pose.sdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare protein sequences\n",
    "\n",
    "Protein sequences can be searched without coordinates. This MMseqs route writes protein scores, ligand-pocket links, the best link per input sequence, and optionally the aligned pocket residues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_sequences(fasta_path):\n",
    "    return search(\n",
    "        Path(fasta_path),\n",
    "        output_dir=Path(fasta_path).with_suffix(\"\") / \"plinder_search\",\n",
    "        threads=8,\n",
    "        store_aligned_pocket_residues=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For several ligand-free coordinate files, use `search(\"models/\", output_dir=\"search_results\")`. Search uses the first coordinate model as the supplied assembly."
   ]
  }
 ],
 "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.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}