{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gzip as gz\n",
    "import os\n",
    "import pandas as pd\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# newline separated accessions, swissprot/trembl based\n",
    "f_acc = \"/mnt/prophane-storage/analyses/20201013-iprg_phase_2/data/accessions_found_by_mpa.csv\"\n",
    "\n",
    "# prophane db files, holding a\n",
    "f_acc2tax = \"/mnt/prophane-storage/dbs/uniprot_tax/2019-09-30/uniprot_tax.map.gz\"\n",
    "f_tax2annot = \"/mnt/prophane-storage/dbs/uniprot_tax/2019-09-30/uniprot_tax.txt\"\n",
    "\n",
    "# to create\n",
    "job_taxmap = \"/mnt/prophane-storage/analyses/20201013-iprg_phase_2/data/taxmap.tsv\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "parsed 6578 accessions from /mnt/prophane-storage/analyses/20201013-iprg_phase_2/data/accessions_found_by_mpa.csv\n"
     ]
    }
   ],
   "source": [
    "accs = set()\n",
    "with open(f_acc, \"r\") as h:\n",
    "    for l in h:\n",
    "        accs.add(l.strip())\n",
    "    print(f\"parsed {len(accs)} accessions from {f_acc}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'accessions_found_by_mpa.csv'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.path.basename(f_acc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/miniconda3/envs/jupyter/lib/python3.7/site-packages/ipykernel_launcher.py:6: DeprecationWarning: The default dtype for empty Series will be 'object' instead of 'float64' in a future version. Specify a dtype explicitly to silence this warning.\n",
      "  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[encoded data, 8200 characters]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[encoded data, 6174 characters]\n",
      "Caching accession-TaxID pairs in /mnt/prophane-storage/analyses/20201013-iprg_phase_2/data/acc2tax_found_by_mpa.csv\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "P30871     83333\n",
       "P80859    224308\n",
       "P00350     83333\n",
       "P14062     99287\n",
       "O34499    224308\n",
       "dtype: object"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_acc2tax_job = os.path.join(os.path.dirname(f_acc), \"acc2tax_found_by_mpa.csv\")\n",
    "if os.path.exists(f_acc2tax_job):\n",
    "    print(f\"Reading cached accession-TaxID pairs from {f_acc2tax_job}\")\n",
    "    s_acc2tax = pd.read_csv(f_acc2tax_job, squeeze=True, index_col=0, header=None)\n",
    "else:\n",
    "    s_acc2tax = pd.Series()\n",
    "    with gz.open(f_acc2tax, \"rt\") as h:\n",
    "        next(h)\n",
    "        for i, l in enumerate(h):\n",
    "            a, t = l.strip().split(\"\\t\")\n",
    "            if i % 100000 == 0:\n",
    "                print(i, end=\"\")\n",
    "            if a and (a in accs):\n",
    "                s_acc2tax[a] = t\n",
    "    print(f\"\\nCaching accession-TaxID pairs in {f_acc2tax_job}\")\n",
    "    s_acc2tax.to_csv(f_acc2tax_job)\n",
    "s_acc2tax.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Taxon\n",
      "Mnemonic\n",
      "Scientific name\n",
      "Common name\n",
      "Synonym\n",
      "Other Names\n",
      "Reviewed\n",
      "Rank\n",
      "Lineage\n",
      "Parent\n",
      "Virus hosts\n"
     ]
    }
   ],
   "source": [
    "with open(f_tax2annot, \"r\") as h:\n",
    "    hl = next(h)\n",
    "cols = hl.strip().split(\"\\t\")\n",
    "\n",
    "print(\"\\n\".join(cols))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtype_dict = {\n",
    "    \"Taxon\": int,\n",
    "    \"Mnemonic\": str,\n",
    "    \"Scientific name\": str,\n",
    "    \"Common name\": str,\n",
    "    \"Synonym\": str,\n",
    "    \"Other Names\": str,\n",
    "    \"Reviewed\": str,\n",
    "    \"Rank\": str,\n",
    "    \"Lineage\": str,\n",
    "    \"Parent\": pd.Int64Dtype(),\n",
    "    \"Virus hosts\": str\n",
    "}\n",
    "df_tax = pd.read_csv(f_tax2annot, sep=\"\\t\", usecols=[0,2,7,8,9], dtype=dtype_dict, na_values=[\"\",\"-\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "taxlevels = [\n",
    "    \"superkingdom\",\n",
    "    \"phylum\",\n",
    "    \"class\",\n",
    "    \"order\",\n",
    "    \"family\",\n",
    "    \"genus\",\n",
    "    \"species\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Taxon</th>\n",
       "      <th>Rank</th>\n",
       "      <th>Lineage</th>\n",
       "      <th>Parent</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1980042</td>\n",
       "      <td>Species</td>\n",
       "      <td>unclassified sequences</td>\n",
       "      <td>1566732</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>85621</td>\n",
       "      <td>Species group</td>\n",
       "      <td>Bacteria; Tenericutes; Mollicutes; Acholeplasm...</td>\n",
       "      <td>33926</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>85623</td>\n",
       "      <td>Species group</td>\n",
       "      <td>Bacteria; Tenericutes; Mollicutes; Acholeplasm...</td>\n",
       "      <td>33926</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>85624</td>\n",
       "      <td>Species group</td>\n",
       "      <td>Bacteria; Tenericutes; Mollicutes; Acholeplasm...</td>\n",
       "      <td>33926</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>85629</td>\n",
       "      <td>Species group</td>\n",
       "      <td>Bacteria; Tenericutes; Mollicutes; Acholeplasm...</td>\n",
       "      <td>33926</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     Taxon           Rank                                            Lineage  \\\n",
       "0  1980042        Species                             unclassified sequences   \n",
       "1    85621  Species group  Bacteria; Tenericutes; Mollicutes; Acholeplasm...   \n",
       "2    85623  Species group  Bacteria; Tenericutes; Mollicutes; Acholeplasm...   \n",
       "3    85624  Species group  Bacteria; Tenericutes; Mollicutes; Acholeplasm...   \n",
       "4    85629  Species group  Bacteria; Tenericutes; Mollicutes; Acholeplasm...   \n",
       "\n",
       "    Parent  \n",
       "0  1566732  \n",
       "1    33926  \n",
       "2    33926  \n",
       "3    33926  \n",
       "4    33926  "
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_tax.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_ranks_parents_names_dicts_for_taxids(db_tax_map):\n",
    "    ranks = {}\n",
    "    parents = {}\n",
    "    names = {}\n",
    "    with open(db_tax_map, 'r') as handle:\n",
    "        handle.readline()\n",
    "        for line in handle:\n",
    "            fields = line.strip().split(\"\\t\") + [''] * 9\n",
    "            taxid = fields[0]\n",
    "            name = fields[2]\n",
    "            rank = fields[7].lower()\n",
    "            parent = fields[9]\n",
    "\n",
    "            names[taxid] = name\n",
    "            parents[taxid] = parent\n",
    "            ranks[taxid] = rank\n",
    "    return ranks, parents, names\n",
    "\n",
    "\n",
    "def get_uniprot_taxinfo_dict(db_tax_map, taxlevels):\n",
    "    \"\"\"\n",
    "    read db_tax_map into dict\n",
    "    :param db_tax_map: uniprot_tax.txt\n",
    "    :return: dictionary:\n",
    "        key: taxID\n",
    "        value: string, containing tab-separated entries (zero based index):\n",
    "            taxid, taxids_lineage, taxlevel[0], taxlevel[1], ...\n",
    "    \"\"\"\n",
    "    taxinfo_dict = {}\n",
    "    taxlevel_count = len(taxlevels)\n",
    "    # db_tax_map = uniprot_tax.txt\n",
    "    ranks, parents, names = get_ranks_parents_names_dicts_for_taxids(db_tax_map)\n",
    "    for taxid in names:\n",
    "        taxids = ['-'] * taxlevel_count\n",
    "        lineage = ['-'] * taxlevel_count\n",
    "        orig_taxid = taxid\n",
    "        while taxid != \"\":\n",
    "            rank = ranks[taxid]\n",
    "            if rank in taxlevels:\n",
    "                taxlevel_index = taxlevels.index(rank)\n",
    "                taxids[taxlevel_index] = taxid\n",
    "                lineage[taxlevel_index] = names[taxid]\n",
    "            taxid = parents[taxid]\n",
    "        taxinfo_dict[orig_taxid] = orig_taxid + \"\\t\" + \";\".join(taxids) + \"\\t\" + \"\\t\".join(lineage)\n",
    "    del names\n",
    "    del parents\n",
    "    del ranks\n",
    "    return taxinfo_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "taxinfo_dict = get_uniprot_taxinfo_dict(f_tax2annot, taxlevels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "headline = \"#acc\\tsource\\thit\\ttaxid\\ttaxids_lineage\\t\" + \"\\t\".join(taxlevels) + \"\\n\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(job_taxmap, 'w') as outhandle:\n",
    "    outhandle.write(headline)\n",
    "    with gz.open(f_acc2tax, 'rt') as db_map_handle:\n",
    "        for line in db_map_handle:\n",
    "            fields = line.strip().split(\"\\t\")\n",
    "            target_db_accession = fields[0]\n",
    "            # accession number from map is among best hits\n",
    "            if target_db_accession in accs:\n",
    "                \n",
    "                taxid = fields[1]\n",
    "                lst_annotations_raw = taxinfo_dict[taxid].split(\"\\t\")\n",
    "               \n",
    "                lst_annotations = ['unclassified' if x == \"-\" else x for x in lst_annotations_raw]\n",
    "                \n",
    "                outhandle.write(\n",
    "                    target_db_accession + \"\\ttaxmap_annot [task XX]\\t\" + \"\\t\".join(lst_annotations) + \"\\n\"\n",
    "                )\n"
   ]
  }
 ],
 "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.7.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
