sample_arg_identification_general.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3aba88e6-0c5a-460f-bcd2-7c776df38f4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import subprocess\n",
    "import warnings\n",
    "from collections import Counter\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from Bio import SeqIO\n",
    "from pandas.core.common import SettingWithCopyWarning\n",
    "import pysam\n",
    "\n",
    "warnings.simplefilter(action=\"ignore\", category=SettingWithCopyWarning)\n",
    "warnings.simplefilter(action=\"ignore\", category=FutureWarning)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "527aa6eb-7f1a-4b30-9818-88e38111647f",
   "metadata": {},
   "source": [
    "# Data location"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c88b1dd8-97f2-4071-a4f6-774647c70e07",
   "metadata": {},
   "outputs": [],
   "source": [
    "threads = 30\n",
    "ont_reads = \"ont_reads\"\n",
    "short_reads_1_location = \"short_reads_1_location\"\n",
    "short_reads_2_location = \"short_reads_2_location\"\n",
    "working_directory = \"working_directory\"\n",
    "Path(working_directory).mkdir(parents=True, exist_ok=True)\n",
    "# database locations\n",
    "card_database = \"card_database\"#card_data/nucleotide_fasta_protein_homolog_model.fasta\"\n",
    "card_database_general = \"card_database_general\"#/card/card_database_v3.1.4.fasta\"\n",
    "blast_nt_database = \"blast_nt_database\"#blast/nt/nt\"\n",
    "mge_database = \"mge_database\"#/MGEs90.fasta\"\n",
    "centrifuge_database = \"centrifuge_database\"#centrifuge/abv\"\n",
    "mob_suite_database = \"mob_suite_database\"#mob_suite/databases\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2da7a152-f2ce-4dd9-be42-c2112ef0f71e",
   "metadata": {},
   "outputs": [],
   "source": [
    "ont_card_blast = os.path.join(working_directory, \"ont_card_blast.out\")\n",
    "ont_arg_related_reads_list = os.path.join(working_directory, \"ont_arg_related.list\")\n",
    "ont_arg_related_reads = os.path.join(working_directory, \"ont_arg_related_reads.fasta\")\n",
    "ont_arg_mge_association = os.path.join(working_directory, \"ont_arg_mge_blast.out\")\n",
    "ont_arg_related_reads_plasflow = os.path.join(working_directory, \"ont_arg_related_reads_plasflow\")\n",
    "ont_arg_related_reads_plasflow_unclassified =f\"{ont_arg_related_reads_plasflow}_unclassified.fasta\"\n",
    "ont_arg_related_reads_plasflow_unclassified_blast = os.path.join(working_directory, \"ont_arg_related_reads_plasflow_unclassified_blast.out\")\n",
    "ont_plasmid_reads = os.path.join(working_directory, \"ont_plasmid_reads.fasta\")\n",
    "ont_plasmid_reads_list_file = os.path.join(working_directory, \"ont_plasmid_reads.list\")\n",
    "ont_mob_suite = os.path.join(working_directory, \"ont_plasmid_mob_suite_analysis.out\")\n",
    "ont_abundance_report = os.path.join(working_directory, \"ont_abundance.csv\")\n",
    "ont_centrifuge_classification_out = os.path.join(working_directory, \"ont_centrifuge.out\")\n",
    "ont_centrifuge_classification_report = os.path.join(working_directory, \"ont_centrifuge.report\")\n",
    "ont_arg_reads_best_hit = os.path.join(working_directory, \"ont_arg_reads_best_hit.csv\")\n",
    "ont_chromosome_list_file = os.path.join(working_directory, \"ont_chromosome.list\")\n",
    "\n",
    "bgi_contigs_folder = os.path.join(working_directory, \"bgi.contig.fa\")\n",
    "bgi_contigs = os.path.join(bgi_contigs_folder, \"bgi.contig.fa\")\n",
    "bgi_rgi_folder = os.path.join(working_directory, \"rgi\")\n",
    "Path(bgi_rgi_folder).mkdir(parents=True, exist_ok=True)\n",
    "bgi_rgi_report_prefix = os.path.join(bgi_rgi_folder, \"rgi_output\")\n",
    "bgi_rgi_report = os.path.join(bgi_rgi_folder, \"rgi_output.txt\")\n",
    "bgi_arg_related_reads_list = os.path.join(working_directory, \"bgi_arg_related.list\")\n",
    "bgi_arg_related_contigs = os.path.join(working_directory, \"bgi_arg_related_contigs.fasta\")\n",
    "bgi_arg_mge_association = os.path.join(working_directory, \"bgi_arg_mge_blast.out\")\n",
    "bgi_arg_related_reads_plasflow = os.path.join(working_directory, \"bgi_arg_related_reads_plasflow\")\n",
    "bgi_arg_related_reads_plasflow_unclassified = f\"{bgi_arg_related_reads_plasflow}_unclassified.fasta\"\n",
    "bgi_arg_related_reads_plasflow_unclassified_blast = os.path.join(working_directory, \"bgi_arg_related_reads_plasflow_unclassified_blast.out\")\n",
    "bgi_contig_regions_bed = os.path.join(working_directory, \"bgi_contig_regions.bed\")\n",
    "bgi_bowtie2_index = os.path.join(bgi_contigs_folder, \"bgi.contig\")\n",
    "bgi_bowtie2_output = os.path.join(bgi_contigs_folder, \"bgi_read_contig.sam\")\n",
    "bgi_reads_contig_mapping = os.path.join(bgi_contigs_folder, \"bgi_read_contig_sorted.bam\")\n",
    "bgi_total_contig_cov = os.path.join(working_directory, \"bgi_total_contig.cov\")\n",
    "bgi_plasmid_reads_list_file = os.path.join(working_directory, \"bgi_plasmid_reads.list\")\n",
    "bgi_plasmid_reads = os.path.join(working_directory, \"bgi_plasmid_reads.fasta\")\n",
    "bgi_mob_suite = os.path.join(working_directory, \"bgi_plasmid_mob_suite_analysis.out\")\n",
    "bgi_abundance_report = os.path.join(working_directory, \"bgi_abundance.csv\")\n",
    "bgi_centrifuge_classification_out = os.path.join(working_directory, \"bgi_centrifuge.out\")\n",
    "bgi_centrifuge_classification_report = os.path.join(working_directory, \"bgi_centrifuge.report\")\n",
    "bgi_mge_reads_best_hit = os.path.join(working_directory, \"bgi_arg_mge_best_hit.csv\")\n",
    "bgi_chromosome_list_file = os.path.join(working_directory, \"bgi_chromosome.list\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45644260-2392-4193-9aff-a2c285ae711b",
   "metadata": {},
   "source": [
    "# Functions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc21fd39-2d01-4a55-bfd5-7d465d62e69c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_card_blast(reads_loc, output_loc, blastdb=card_database):\n",
    "    output_fmt = '\"6 qseqid sseqid pident length mismatch gapopen qstart qend sstart send evalue bitscore slen\"'\n",
    "    cmd = f\"blastn  -task megablast -num_threads {threads} -query {reads_loc} -outfmt {output_fmt} -db {blastdb} -out {output_loc}\"\n",
    "    print(f\"running command {cmd}\")\n",
    "    os.system(cmd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10e1b040-6592-4024-b622-3f648c6e6617",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_seqtk_subseq(query_file, candidate_list, output_file):\n",
    "    cmd = f\"seqtk subseq {query_file} {candidate_list} > {output_file}\"\n",
    "    print(f\"running command {cmd}\")\n",
    "    os.system(cmd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cd55b67-d828-4570-ad04-c1b87922b4b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def export_candidate_read_list(candidate_list_file, best_hit_list):\n",
    "    with open(candidate_list_file, \"w\") as arg_related_f:\n",
    "        best_hit_read_list = []\n",
    "        for i in best_hit_list:\n",
    "            best_hit_read_list.append(i[1])  # i[1]: read id\n",
    "        for i in set(best_hit_read_list):\n",
    "            arg_related_f.writelines(f\"{i}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9004e91-1b69-40b2-9731-8301a94bf47d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_candidate_list_to_file(candidate_list_file, candidate_list):\n",
    "    with open(candidate_list_file, \"w\") as arg_contig_list_f:\n",
    "        for i in candidate_list:\n",
    "            arg_contig_list_f.writelines(f\"{i}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfce7984-63bf-43f7-8bff-c2a35f1dd46e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_arg_mge_association(\n",
    "    mge_best_hit, arg_best_hit, max_distance_range=[0, 100, 200, 500, 1000]\n",
    "):\n",
    "    mge_best_hit_df = pd.DataFrame(\n",
    "        mge_best_hit,\n",
    "        columns=[\n",
    "            \"index\",\n",
    "            \"qseqid\",\n",
    "            \"sseqid\",\n",
    "            \"pident\",\n",
    "            \"length\",\n",
    "            \"mismatch\",\n",
    "            \"gapopen\",\n",
    "            \"qstart\",\n",
    "            \"qend\",\n",
    "            \"sstart\",\n",
    "            \"send\",\n",
    "            \"evalue\",\n",
    "            \"bitscore\",\n",
    "            \"slen\",\n",
    "            \"qcov\",\n",
    "        ],\n",
    "    )\n",
    "    result_distance_dict = dict.fromkeys(max_distance_range)\n",
    "    for k in max_distance_range:\n",
    "        arg_mge_association_df = pd.DataFrame(columns=[\"ARG\", \"read_id\", \"mge_id\"])\n",
    "        for read_to_arg in arg_best_hit:\n",
    "            arg_name = read_to_arg[2]\n",
    "            seq_name = read_to_arg[1]\n",
    "            seq_arg_qstart = min(read_to_arg[7], read_to_arg[8])\n",
    "            seq_arg_qend = max(read_to_arg[7], read_to_arg[8])\n",
    "\n",
    "            arg_related_mge_df = mge_best_hit_df[mge_best_hit_df[\"qseqid\"] == seq_name]\n",
    "            for j in arg_related_mge_df.iterrows():\n",
    "                mge_align = j[1]\n",
    "                mge_align_qstart = min(mge_align.qstart, mge_align.qend)\n",
    "                mge_align_qend = max(mge_align.qstart, mge_align.qend)\n",
    "                mge_type = mge_align.sseqid\n",
    "                if (seq_arg_qstart <= mge_align_qend + k) and (\n",
    "                    seq_arg_qend >= mge_align_qstart - k\n",
    "                ):\n",
    "                    arg_related_mge = mge_type\n",
    "                    #                     print(arg_related_mge, arg_name, seq_name, seq_arg_qstart, seq_arg_qend, mge_align_qend, mge_align_qstart)\n",
    "                    arg_mge_association_df = arg_mge_association_df.append(\n",
    "                        {\n",
    "                            \"ARG\": arg_name,\n",
    "                            \"read_id\": seq_name,\n",
    "                            \"mge_id\": arg_related_mge,\n",
    "                        },\n",
    "                        ignore_index=True,\n",
    "                    )\n",
    "        result_distance_dict[k] = arg_mge_association_df\n",
    "    return result_distance_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1c1e2e2-025d-4882-8351-134a0e6ffd46",
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_all_arg_mge_distance(\n",
    "    mge_best_hit, arg_best_hit):\n",
    "    mge_best_hit_df = pd.DataFrame(\n",
    "        mge_best_hit,\n",
    "        columns=[\n",
    "            \"index\",\n",
    "            \"qseqid\",\n",
    "            \"sseqid\",\n",
    "            \"pident\",\n",
    "            \"length\",\n",
    "            \"mismatch\",\n",
    "            \"gapopen\",\n",
    "            \"qstart\",\n",
    "            \"qend\",\n",
    "            \"sstart\",\n",
    "            \"send\",\n",
    "            \"evalue\",\n",
    "            \"bitscore\",\n",
    "            \"slen\",\n",
    "            \"qcov\",\n",
    "        ],\n",
    "    )\n",
    "    arg_mge_association_df = pd.DataFrame(columns=[\"ARG\", \"read_id\", \"mge_id\", \"distance\", \"mge_pident\", \"mge_lencov\"])\n",
    "    for read_to_arg in arg_best_hit:\n",
    "        arg_name = read_to_arg[2]\n",
    "        seq_name = read_to_arg[1]\n",
    "        seq_arg_qstart = min(read_to_arg[7], read_to_arg[8])\n",
    "        seq_arg_qend = max(read_to_arg[7], read_to_arg[8])\n",
    "        seq_arg_middle = (seq_arg_qstart+seq_arg_qend)/2.0\n",
    "        arg_related_mge_df = mge_best_hit_df[mge_best_hit_df[\"qseqid\"] == seq_name]\n",
    "        for j in arg_related_mge_df.iterrows():\n",
    "            mge_align = j[1]\n",
    "            mge_align_qstart = min(mge_align.qstart, mge_align.qend)\n",
    "            mge_align_qend = max(mge_align.qstart, mge_align.qend)\n",
    "            mge_align_middle = (mge_align_qstart+mge_align_qend)/2.0\n",
    "            mge_type = mge_align.sseqid\n",
    "            arg_related_mge = mge_type\n",
    "            arg_mge_distance = None\n",
    "            if seq_arg_qstart > mge_align_qstart:\n",
    "                arg_mge_distance = seq_arg_qstart - mge_align_qend\n",
    "            else:\n",
    "                arg_mge_distance = mge_align_qstart -seq_arg_qend\n",
    "            arg_mge_association_df = arg_mge_association_df.append(\n",
    "                {\n",
    "                    \"ARG\": arg_name,\n",
    "                    \"read_id\": seq_name,\n",
    "                    \"mge_id\": arg_related_mge,\n",
    "                    \"distance\": arg_mge_distance,\n",
    "                    \"mge_pident\": mge_align.pident,\n",
    "                    \"mge_lencov\": mge_align.qcov,\n",
    "                },\n",
    "                ignore_index=True,\n",
    "            )\n",
    "        \n",
    "    return arg_mge_association_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c1c857b-9c7b-4540-b107-70d7f883788c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_centrifuge(reads_loc, classification_output, report_file):\n",
    "    cmd = f\"centrifuge -x {centrifuge_database} -U {reads_loc} -f -S {classification_output} --report-file {report_file} -p {threads}\"\n",
    "    print(f\"running command {cmd}\")\n",
    "    os.system(cmd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4885367-c91c-4817-92bb-77faeb130b12",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_nt_blast(reads_loc, output_loc, blastdb=blast_nt_database):\n",
    "    output_fmt = '\"6 qseqid sseqid pident length mismatch gapopen qstart qend sstart send evalue bitscore slen stitle\"'\n",
    "    cmd = f\"blastn  -task megablast -num_threads {threads} -query {reads_loc} -outfmt {output_fmt} -db {blastdb} -out {output_loc}\"\n",
    "    print(f\"running command {cmd}\")\n",
    "    os.system(cmd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "021747eb-d140-455d-a446-dca71a5a9578",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_mge_blast(reads_loc, output_loc, blastdb=mge_database):\n",
    "    output_fmt = '\"6 qseqid sseqid pident length mismatch gapopen qstart qend sstart send evalue bitscore slen\"'\n",
    "    cmd = f\"blastx -num_threads {threads} -query {reads_loc} -outfmt {output_fmt} -db {blastdb} -out {output_loc}\"\n",
    "    print(f\"running command {cmd}\")\n",
    "    os.system(cmd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c724041c-835a-4c3f-897e-82fe04d45f31",
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_blast_alignment(blastn_card_result, threshold=0.7, overlap_borderline=20):\n",
    "    \"\"\"\n",
    "    Given a sample's reads' blast result against card database, run quality filter and length filter.\n",
    "    Make sure there is no overlap when there are multiple hits for one read, fuzzy the overlapping borderline by 20bps.\n",
    "    Return:\n",
    "        best_hit_list: this variable consists the best arg hit of each read in the database.\n",
    "        Hits do not meet the lencov and pident requirements are discarded.\n",
    "    \"\"\"\n",
    "    blast_output_fmt = [\n",
    "        \"qseqid\",\n",
    "        \"sseqid\",\n",
    "        \"pident\",\n",
    "        \"length\",\n",
    "        \"mismatch\",\n",
    "        \"gapopen\",\n",
    "        \"qstart\",\n",
    "        \"qend\",\n",
    "        \"sstart\",\n",
    "        \"send\",\n",
    "        \"evalue\",\n",
    "        \"bitscore\",\n",
    "        \"slen\",\n",
    "    ]\n",
    "    blastn_card_df = pd.read_csv(\n",
    "        blastn_card_result, sep=\"\\t\", header=None, names=blast_output_fmt\n",
    "    )\n",
    "    blastn_card_df_pident = blastn_card_df[blastn_card_df.pident >= threshold * 100]\n",
    "    blastn_card_df_pident[\"lencov\"] = (\n",
    "        blastn_card_df_pident.length / blastn_card_df_pident.slen\n",
    "    )\n",
    "    blastn_card_df_lencov = blastn_card_df_pident[\n",
    "        blastn_card_df_pident.lencov >= threshold\n",
    "    ]\n",
    "    seqid_list = list(set(list(blastn_card_df_lencov.qseqid)))\n",
    "    best_hit_list = []\n",
    "    for seqid in seqid_list:\n",
    "        seq_df = blastn_card_df_lencov[\n",
    "            blastn_card_df_lencov.qseqid == seqid\n",
    "        ].sort_values(by=[\"pident\", \"lencov\"], ascending=[False, False])\n",
    "        seq_df.reset_index(inplace=True)\n",
    "        seq_df_values = seq_df.values\n",
    "        seq_alignment_block = []\n",
    "        for i in seq_df_values:\n",
    "            qstart = min(i[7], i[8])  # sometimes there are reverse ones\n",
    "            qend = max(i[7], i[8])\n",
    "            append_flag = 1\n",
    "            for j in seq_alignment_block:\n",
    "                if (qend >= j[0] + overlap_borderline) and (\n",
    "                    qstart <= j[1] - overlap_borderline\n",
    "                ):\n",
    "                    j[2].append([i])\n",
    "                    j[0] = min(qstart, j[1])\n",
    "                    j[1] = max(qend, j[0])\n",
    "                    append_flag = 0\n",
    "            if append_flag == 1:\n",
    "                seq_alignment_block.append([qstart, qend, [i]])\n",
    "        for i in seq_alignment_block:\n",
    "            best_hit_list.append(i[2][0])\n",
    "    return best_hit_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44ef861b-5c9a-4508-aa73-bc0f39863daa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_blast_names(megablast_result, blast_cutoff=5):\n",
    "    \"\"\"\n",
    "    Given a blast result against nt database, this function return top 5 hits with stitle of blast result in the database\n",
    "    \"\"\"\n",
    "    reads_list = []\n",
    "    nt_blast_output_fmt = [\n",
    "        \"qseqid\",\n",
    "        \"sseqid\",\n",
    "        \"pident\",\n",
    "        \"length\",\n",
    "        \"mismatch\",\n",
    "        \"gapopen\",\n",
    "        \"qstart\",\n",
    "        \"qend\",\n",
    "        \"sstart\",\n",
    "        \"send\",\n",
    "        \"evalue\",\n",
    "        \"bitscore\",\n",
    "        \"slen\",\n",
    "        \"stitle\",\n",
    "    ]\n",
    "    megablast_df = pd.read_csv(\n",
    "        megablast_result, sep=\"\\t\", header=None, names=nt_blast_output_fmt\n",
    "    )\n",
    "    #     for record in SeqIO.parse(read_file, \"fasta\"):\n",
    "    #         reads_list.append(record.id)\n",
    "    reads_blast_dict = {}\n",
    "    for read_id in list(set(list(megablast_df.qseqid))):\n",
    "        #         if read_id in list(megablast_df.qseqid):\n",
    "        cutoff_df = megablast_df.groupby(\"qseqid\").get_group(read_id)[:blast_cutoff]\n",
    "        reads_blast_dict.update({read_id: cutoff_df.stitle.values})\n",
    "    return reads_blast_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21c57abe-70ce-4b4e-aa39-f8fb484ee20e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_read_location_from_blast(blast_result_dict):\n",
    "    location_dict = {}\n",
    "    for read_id in blast_result_dict.keys():\n",
    "        first_hit = blast_result_dict[read_id][0]\n",
    "        if \"plasmid\" in first_hit.split(\" \"):\n",
    "            location_dict.update({read_id: \"plasmid\"})\n",
    "        elif \"chromosome\" in first_hit.split(\" \"):\n",
    "            location_dict.update({read_id: \"chromosome\"})\n",
    "        else:\n",
    "            location_dict.update({read_id: \"unclassified\"})\n",
    "    return location_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "281153e9-8709-4324-abbe-954e18625574",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_short_read_cov(\n",
    "    bgi_rgi_output,\n",
    "    region_bed,\n",
    "    total_contig_cov,\n",
    "    short_read_mapping,\n",
    "    samtools_quality_filter=20,\n",
    "):\n",
    "    # /home/project_data/wastewater_arg/bgi/inf_reads/mapping/inf_read_contig_sorted.bam\n",
    "    with open(region_bed, \"w\") as rgi_alignment_region:\n",
    "        for i in bgi_rgi_output.iterrows():\n",
    "            content = i[1]\n",
    "            ctg_split = content[\"Contig\"].split(\"_\")\n",
    "            contig_name = f\"{ctg_split[0]}_{ctg_split[1]}\"\n",
    "            #             contig_name = content[\"Contig\"]\n",
    "            start = content[\"Start\"]\n",
    "            stop = content[\"Stop\"]\n",
    "            rgi_alignment_region.writelines(f\"{contig_name}\\t{start}\\t{stop}\\n\")\n",
    "    samtools_cmd = f\"samtools bedcov -j -Q {samtools_quality_filter} {region_bed} {short_read_mapping} > {total_contig_cov}\"\n",
    "    print(f\"running command {samtools_cmd}\")\n",
    "    os.system(samtools_cmd)\n",
    "    ctg_base_count_df = pd.read_csv(\n",
    "        total_contig_cov,\n",
    "        sep=\"\\t\",\n",
    "        header=None,\n",
    "        names=[\"ctg_name\", \"start\", \"stop\", \"total_base_count\"],\n",
    "    )\n",
    "    ctg_base_count_df[\"avg_base_count\"] = ctg_base_count_df[\"total_base_count\"] / (\n",
    "        ctg_base_count_df[\"stop\"] - ctg_base_count_df[\"start\"]\n",
    "    )\n",
    "    ctg_base_count_df[\"ctg_name_id\"] = bgi_rgi_output.reset_index().Contig\n",
    "    return ctg_base_count_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c8b866f-637e-4121-8aea-189508903a75",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rgi_output_filter(file_loc, identity_threshold=0.7, length_threshold=0.7):\n",
    "    \"\"\"\n",
    "    Quality filter for rgi results.\n",
    "    Return: a dataframe consists rgi result with qualified reads\n",
    "    \"\"\"\n",
    "    i_t = identity_threshold * 100\n",
    "    l_t = length_threshold * 100\n",
    "    rgi_df = pd.read_csv(file_loc, sep=\"\\t\")\n",
    "    df_after_i_t = rgi_df[rgi_df.Best_Identities >= i_t]\n",
    "    df_after_l_t = df_after_i_t[\n",
    "        df_after_i_t[\"Percentage Length of Reference Sequence\"] >= l_t\n",
    "    ]\n",
    "    return df_after_l_t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82d29b21-f8c7-4f99-8b03-8e1d472a66d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_rgi_contig_list(rgi_output_df):\n",
    "    \"\"\"\n",
    "    given a filtered rgi output, return a list of contigs that is associated with args\n",
    "    \"\"\"\n",
    "    contig_name_list = set()\n",
    "    for i in rgi_output_df.Contig:\n",
    "        contig_name_l = i.split(\"_\")\n",
    "        contig_name = \"_\".join(contig_name_l[:2])\n",
    "        contig_name_list.add(contig_name)\n",
    "    return list(contig_name_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43064194-7091-423f-b85b-d7f7636a6eb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_location_and_abundance(\n",
    "    arg_related_reads_plasflow,\n",
    "    arg_best_hit,\n",
    "    mge_best_hit,\n",
    "    arg_related_reads_plasflow_unclassified_blast_location_dict,\n",
    "    lr=True,\n",
    "    sr_avg_read_cov=None,\n",
    "):\n",
    "    chromosome_list = []\n",
    "    best_hit_columns = [\n",
    "        \"index\",\n",
    "        \"qseqid\",\n",
    "        \"sseqid\",\n",
    "        \"pident\",\n",
    "        \"length\",\n",
    "        \"mismatch\",\n",
    "        \"gapopen\",\n",
    "        \"qstart\",\n",
    "        \"qend\",\n",
    "        \"sstart\",\n",
    "        \"send\",\n",
    "        \"evalue\",\n",
    "        \"bitscore\",\n",
    "        \"slen\",\n",
    "        \"qcov\",\n",
    "    ]\n",
    "    # PREPROCESS\n",
    "    # read dataframe and csv\n",
    "    arg_related_reads_plasflow_df = pd.read_csv(arg_related_reads_plasflow, sep=\"\\t\")\n",
    "    mge_best_hit_df = pd.DataFrame(mge_best_hit, columns=best_hit_columns)[\n",
    "        [\"qseqid\", \"sseqid\"]\n",
    "    ]\n",
    "    # get arg name to read id dict\n",
    "    # get arg name to arg acc dict\n",
    "    if lr:\n",
    "        arg_best_hit_df = pd.DataFrame(arg_best_hit, columns=best_hit_columns)[\n",
    "            [\"qseqid\", \"sseqid\"]\n",
    "        ]\n",
    "        arg_best_hit_df[\"arg_acc\"] = arg_best_hit_df[\"sseqid\"].apply(extract_arg_acc)\n",
    "        arg_best_hit_df[\"arg_name\"] = arg_best_hit_df[\"arg_acc\"].apply(extract_arg)\n",
    "        arg_reads_dict = dataframe_to_dict(arg_best_hit_df, \"arg_name\", \"qseqid\")\n",
    "        arg_acc_name_dict = dict(zip(arg_best_hit_df.arg_name, arg_best_hit_df.arg_acc))\n",
    "    else:\n",
    "        arg_best_hit_df = arg_best_hit\n",
    "        arg_best_hit_df[\"contig_name\"] = arg_best_hit_df[\"Contig\"].apply(extract_contig)\n",
    "        arg_best_hit_df[\"arg_acc\"] = arg_best_hit_df[\"ARO\"].apply(add_ARO)\n",
    "        arg_reads_dict = dataframe_to_dict(\n",
    "            arg_best_hit_df, \"Best_Hit_ARO\", \"contig_name\"\n",
    "        )\n",
    "        arg_reads_id_dict = dataframe_to_dict(arg_best_hit_df, \"Best_Hit_ARO\", \"Contig\")\n",
    "        arg_acc_name_dict = dict(\n",
    "            zip(arg_best_hit_df.Best_Hit_ARO, arg_best_hit_df.arg_acc)\n",
    "        )\n",
    "    # get read id to mge type dict\n",
    "    mge_best_hit_df[\"mge_type\"] = mge_best_hit_df[\"sseqid\"].apply(extract_mge)\n",
    "    mge_reads_dict = dataframe_to_dict(mge_best_hit_df, \"qseqid\", \"mge_type\")\n",
    "    # get read id to plasflow location dict\n",
    "    plawflow_df_info = arg_related_reads_plasflow_df[[\"contig_name\", \"label\"]]\n",
    "    plawflow_df_info[\"label\"] = plawflow_df_info[\"label\"].apply(extract_plasflow)\n",
    "    read_plasflow_dict = dict(zip(plawflow_df_info.contig_name, plawflow_df_info.label))\n",
    "\n",
    "    # GENERATE RESULTS\n",
    "    # result dataframe\n",
    "    stats_column = [\n",
    "        \"transposase\",\n",
    "        \"IntI1\",\n",
    "        \"integrase\",\n",
    "        \"recombinase\",\n",
    "        \"plasmid\",\n",
    "        \"chromosome\",\n",
    "        \"unclassified\",\n",
    "        \"arg_name\",\n",
    "        \"arg_ref_len\",\n",
    "        \"arg_read_num\",\n",
    "    ]\n",
    "    abundance_df = pd.DataFrame(columns=stats_column)\n",
    "\n",
    "    result_dict = {}\n",
    "    for arg_name in arg_reads_dict.keys():\n",
    "        transposase_num = 0\n",
    "        IntI1_num = 0\n",
    "        integrase_num = 0\n",
    "        recombinase_num = 0\n",
    "        plasmid_num = 0\n",
    "        chromosome_num = 0\n",
    "        unclassified_num = 0\n",
    "        read_list = arg_reads_dict[arg_name]\n",
    "        if lr:\n",
    "            read_num = len(read_list)\n",
    "        else:\n",
    "            read_id_list = arg_reads_id_dict[arg_name]\n",
    "            sr_cov_dict = dict(\n",
    "                zip(sr_avg_read_cov.ctg_name_id, sr_avg_read_cov.avg_base_count)\n",
    "            )\n",
    "            read_num = 0\n",
    "            for contig in read_id_list:\n",
    "                read_num += sr_cov_dict[contig]\n",
    "        arg_len = aro_len_dict[arg_acc_name_dict[arg_name]]\n",
    "        # plasflow location and blast location\n",
    "        for read_id in read_list:\n",
    "            plasflow_location = read_plasflow_dict[read_id]\n",
    "            if plasflow_location == \"plasmid\":\n",
    "                plasmid_num += 1\n",
    "            elif plasflow_location == \"chromosome\":\n",
    "                chromosome_list.append(read_id)\n",
    "                chromosome_num += 1\n",
    "            else:\n",
    "                if (\n",
    "                    read_id\n",
    "                    in arg_related_reads_plasflow_unclassified_blast_location_dict.keys()\n",
    "                ):\n",
    "                    blast_location = (\n",
    "                        arg_related_reads_plasflow_unclassified_blast_location_dict[\n",
    "                            read_id\n",
    "                        ]\n",
    "                    )\n",
    "                    if blast_location == \"plasmid\":\n",
    "                        plasmid_num += 1\n",
    "                    elif blast_location == \"chromosome\":\n",
    "                        chromosome_list.append(read_id)\n",
    "                        chromosome_num += 1\n",
    "                    else:\n",
    "                        unclassified_num += 1\n",
    "                else:\n",
    "                    unclassified_num += 1\n",
    "            # mges\n",
    "            if read_id in mge_reads_dict.keys():\n",
    "                mge_list = mge_reads_dict[read_id]\n",
    "                for mge_type in mge_list:\n",
    "                    if mge_type == \"transposase\":\n",
    "                        transposase_num += 1\n",
    "                    elif mge_type == \"IntI1\":\n",
    "                        IntI1_num += 1\n",
    "                    elif mge_type == \"integrase\":\n",
    "                        integrase_num += 1\n",
    "                    elif mge_type == \"recombinase\":\n",
    "                        recombinase_num += 1\n",
    "\n",
    "        abundance_df = abundance_df.append(\n",
    "            {\n",
    "                \"transposase\": transposase_num,\n",
    "                \"IntI1\": IntI1_num,\n",
    "                \"integrase\": integrase_num,\n",
    "                \"recombinase\": recombinase_num,\n",
    "                \"plasmid\": plasmid_num,\n",
    "                \"chromosome\": chromosome_num,\n",
    "                \"unclassified\": unclassified_num,\n",
    "                \"arg_name\": arg_name,\n",
    "                \"arg_ref_len\": arg_len,\n",
    "                \"arg_read_num\": read_num,\n",
    "            },\n",
    "            ignore_index=True,\n",
    "        )\n",
    "    return abundance_df, chromosome_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "301f4037-8834-4c36-a87e-1f986a2fcb62",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_plasmid_read_id_list(\n",
    "    arg_related_reads_plasflow,\n",
    "    arg_related_reads_plasflow_unclassified_blast_location_dict,\n",
    "):\n",
    "    plasflow_df = pd.read_csv(arg_related_reads_plasflow, sep=\"\\t\")\n",
    "    plawflow_df_info = plasflow_df[[\"contig_name\", \"label\"]]\n",
    "    plawflow_df_info[\"label\"] = plawflow_df_info[\"label\"].apply(extract_plasflow)\n",
    "    read_plasflow_dict = dict(zip(plawflow_df_info.contig_name, plawflow_df_info.label))\n",
    "    plasmid_read_list = []\n",
    "    for i in read_plasflow_dict:\n",
    "        if read_plasflow_dict[i] == \"plasmid\":\n",
    "            plasmid_read_list.append(i)\n",
    "        elif read_plasflow_dict[i] == \"unclassified\":\n",
    "            if i in list(\n",
    "                arg_related_reads_plasflow_unclassified_blast_location_dict.keys()\n",
    "            ):\n",
    "                if (\n",
    "                    arg_related_reads_plasflow_unclassified_blast_location_dict[i]\n",
    "                    == \"plasmid\"\n",
    "                ):\n",
    "                    plasmid_read_list.append(i)\n",
    "    return plasmid_read_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d0baa10-ab54-4691-8deb-c02aa802de06",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_mob_suite(in_file, out_file):\n",
    "    cmd = f\"mob_typer -i {os.path.abspath(in_file)} -o {os.path.abspath(out_file)} -n {threads} -x -f\"\n",
    "    print(f\"{cmd}\")\n",
    "\n",
    "\n",
    "#     os.system(cmd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cc141d6-6cae-4cee-aa3f-752a8f2ae305",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_mge(mge_str):\n",
    "    return mge_str.split(\"|\")[-2]\n",
    "\n",
    "\n",
    "def extract_arg(arg_acc):\n",
    "    aro_name = aro_acc_name_dict[arg_acc].split(\"[\")[0][:-1]\n",
    "    return aro_name\n",
    "\n",
    "\n",
    "def extract_arg_acc(mge_str):\n",
    "    return mge_str.split(\"|\")[-2]\n",
    "\n",
    "\n",
    "def extract_plasflow(mge_str):\n",
    "    return mge_str.split(\".\")[0]\n",
    "\n",
    "\n",
    "def extract_contig(mge_str):\n",
    "    contig_l = mge_str.split(\"_\")\n",
    "    return f\"{contig_l[0]}_{contig_l[1]}\"\n",
    "\n",
    "\n",
    "def add_ARO(aro_num):\n",
    "    return f\"ARO:{aro_num}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42221a40-3c86-4011-987d-a3f42deab9f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def output_best_hit(arg_best_hit, output_csv):\n",
    "    best_hit_columns = [\n",
    "        \"index\",\n",
    "        \"qseqid\",\n",
    "        \"sseqid\",\n",
    "        \"pident\",\n",
    "        \"length\",\n",
    "        \"mismatch\",\n",
    "        \"gapopen\",\n",
    "        \"qstart\",\n",
    "        \"qend\",\n",
    "        \"sstart\",\n",
    "        \"send\",\n",
    "        \"evalue\",\n",
    "        \"bitscore\",\n",
    "        \"slen\",\n",
    "        \"qcov\",\n",
    "    ]\n",
    "    arg_best_hit_df = pd.DataFrame(arg_best_hit, columns=best_hit_columns)\n",
    "    arg_best_hit_df.to_csv(output_csv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e76623d4-5962-45e0-8a7e-9409786c164f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dataframe_to_dict(transform_df, column_1, column_2):\n",
    "    \"\"\"\n",
    "    Transform a dataframe into dictionary. column 1 is the key, column 2 is the value.\n",
    "    In this function, value can be multiple ones, so they are stored in a list.\n",
    "    \"\"\"\n",
    "    ret_dict = {}\n",
    "    for data_series in transform_df.iterrows():\n",
    "        row_content = data_series[1]\n",
    "        if row_content[column_1] in list(ret_dict.keys()):\n",
    "            ret_dict[row_content[column_1]].append(row_content[column_2])\n",
    "        else:\n",
    "            ret_dict.update({row_content[column_1]: [row_content[column_2]]})\n",
    "    return ret_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e79cbca4-9403-4781-a0d4-abb683d4cad2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_list_to_file(writing_list, writing_file):\n",
    "    with open(writing_file, \"w\") as f:\n",
    "        for i in writing_list:\n",
    "            f.writelines(f\"{i}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a1a2326-f842-4208-a165-2dfa02f81082",
   "metadata": {},
   "source": [
    "# ARG identification with CARD database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20173236-bd51-4be2-8c6f-2ab9ec91f209",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Parse CARD database, link ARO acc to its length and its name\n",
    "\"\"\"\n",
    "aro_len_dict = {}\n",
    "aro_acc_name_dict = {}\n",
    "with open(card_database) as card_db:\n",
    "    for record in SeqIO.parse(card_db, \"fasta\"):\n",
    "        record_id_list = record.description.split(\"|\")\n",
    "        # print(record_id_list)\n",
    "        record_acc = record_id_list[4]\n",
    "        record_name = record_id_list[5]\n",
    "        record_len = len(record.seq)\n",
    "        aro_len_dict.update({record_acc: record_len})\n",
    "        aro_acc_name_dict.update({record_acc: record_name})\n",
    "with open(card_database_general) as card_db:\n",
    "    for record in SeqIO.parse(card_db, \"fasta\"):\n",
    "        record_id_list = record.description.split(\"|\")\n",
    "        record_acc = record_id_list[0]\n",
    "        #         record_name = record_id_list[5]\n",
    "        record_len = len(record.seq)\n",
    "        aro_len_dict.update({record_acc: record_len})\n",
    "#         aro_acc_name_dict.update({record_acc:record_name})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d119cf7-b8b6-41f5-874b-d632509f3516",
   "metadata": {},
   "source": [
    "# Short read assembly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4066d8c8-0005-4178-9e54-d67bb58255c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.system(f\"megahit -1 {short_reads_1_location} -2 {short_reads_2_location} -t {threads} -o {bgi_contigs}\")\n",
    "# remove contigs less than 500 bps\n",
    "contig_seq = os.path.join(bgi_contigs_folder, \"final.contigs.fa\")\n",
    "os.system(f\"seqkit seq -m 500 {contig_seq} > {bgi_contigs}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "197c6153-c6d3-444f-b05c-304369633a28",
   "metadata": {},
   "source": [
    "## ONT fasta reads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a789076-c3fc-450f-ac0c-534a30067d97",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run blast against CARD database for inf, eff and fin reads\n",
    "run_card_blast(ont_reads, ont_card_blast)\n",
    "# get the best hit of all qualified reads-arg pairs\n",
    "ont_arg_best_hit = check_blast_alignment(ont_card_blast)\n",
    "output_best_hit(ont_arg_best_hit, ont_arg_reads_best_hit)\n",
    "# export arg-related reads as a list\n",
    "export_candidate_read_list(ont_arg_related_reads_list, ont_arg_best_hit)\n",
    "# extract ONT arg related reads\n",
    "run_seqtk_subseq(ont_reads, ont_arg_related_reads_list, ont_arg_related_reads)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03fda134-25f0-4ea6-8bd1-ca8624b1585b",
   "metadata": {},
   "source": [
    "## short reads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba4e9310-1c7b-4c46-914c-d86646abd694",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.system(f\"rgi main -i {bgi_contigs} -o {bgi_rgi_report_prefix} -n {threads}\")\n",
    "# get filtered rgi results\n",
    "rgi_filtered = rgi_output_filter(bgi_rgi_report)\n",
    "# get arg-related contigs as a list\n",
    "rgi_contig_list = get_rgi_contig_list(rgi_filtered)\n",
    "# write arg-related contigs list to file\n",
    "write_candidate_list_to_file(bgi_arg_related_reads_list, rgi_contig_list)\n",
    "# extract BGI arg related contigs\n",
    "run_seqtk_subseq(bgi_contigs, bgi_arg_related_reads_list, bgi_arg_related_contigs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df23210a-cd02-4112-9512-0a82676198ab",
   "metadata": {},
   "source": [
    "# ARG-associated reads MGE identification"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22f7da1b-c422-4c7e-8f2d-ad0c5ee35073",
   "metadata": {},
   "source": [
    "## ONT reads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "645fc34f-d268-486e-b7dc-bd25eafe5ac2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run blast against mge reads for arg-related inf, eff and fin reads\n",
    "run_mge_blast(ont_arg_related_reads, ont_arg_mge_association)\n",
    "# get each read' best mge hits\n",
    "ont_mge_best_hit = check_blast_alignment(ont_arg_mge_association)\n",
    "# generate a dictionary that have arg and mge within a certain distance\n",
    "ont_arg_mge_distance_range = get_arg_mge_association(ont_mge_best_hit, ont_arg_best_hit)\n",
    "check_all_arg_mge_distance(ont_mge_best_hit, ont_arg_best_hit).to_csv(os.path.join(working_directory,\"ont_arg_mge_distance.csv\"))\n",
    "for range_num in ont_arg_mge_distance_range:\n",
    "    ont_arg_mge_distance_range[range_num].to_csv(os.path.join(working_directory, f\"inf_arg_mge_association_{range_num}.csv\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd289f99-6378-4733-82d1-2df19a902fe6",
   "metadata": {},
   "source": [
    "## Short reads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1802e01b-e30c-465d-9b29-a4280edd3ea3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run blast against mge reads for arg-related inf, eff and fin reads\n",
    "run_mge_blast(bgi_arg_related_contigs, bgi_arg_mge_association)\n",
    "# get each contig' best mge hits\n",
    "bgi_mge_best_hit = check_blast_alignment(bgi_arg_mge_association)\n",
    "output_best_hit(bgi_mge_best_hit, bgi_mge_reads_best_hit)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "113d9b06-7f79-4dca-9302-f715400fa768",
   "metadata": {},
   "source": [
    "#### RGI does not have blast-like output for locations, so cannot determine if a ARG is flanked by MGE"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc2b52e6-396f-405b-902a-b9e56a9227ef",
   "metadata": {},
   "source": [
    "## PlasFlow analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab388e2f-77dc-4f21-8bc2-79b22e8cf3ec",
   "metadata": {},
   "source": [
    "### PlasFlow has to be run at another environment, here are commands for copy-pasting:\n",
    "## ONT reads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "288851d1-e1ff-43ec-8bbf-33cd428ed5a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"conda activate plasflow\")\n",
    "print(f\"PlasFlow.py --input {os.path.abspath(ont_arg_related_reads)} --output {os.path.abspath(ont_arg_related_reads_plasflow)}\")\n",
    "print(\"conda activate base\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "847dc815-f061-4d84-af51-9316a2ffd6fd",
   "metadata": {},
   "source": [
    "## Short reads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1574e040-4f36-46ad-81ef-98810f27ba4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"conda activate plasflow\")\n",
    "print(f\"PlasFlow.py --input {os.path.abspath(bgi_arg_related_contigs)} --output {os.path.abspath(bgi_arg_related_reads_plasflow)}\")\n",
    "print(\"conda activate base\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6caf80e-e4b7-4bf9-be3b-ab5ab270c7b9",
   "metadata": {},
   "source": [
    "## Classify PlasFlow unclassified reads with blastn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08a26970-c59d-4564-80b8-652bd0c583ee",
   "metadata": {},
   "source": [
    "## ONT reads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5056d3b4-42e5-4c6f-b4c7-fdab960d068a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use megablast to check unclassified results from plasflow\n",
    "run_nt_blast(ont_arg_related_reads_plasflow_unclassified, ont_arg_related_reads_plasflow_unclassified_blast,)\n",
    "ont_arg_related_reads_plasflow_unclassified_blast_dict = get_blast_names(ont_arg_related_reads_plasflow_unclassified_blast)\n",
    "ont_arg_related_reads_plasflow_unclassified_blast_location_dict = (get_read_location_from_blast(ont_arg_related_reads_plasflow_unclassified_blast_dict))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75833114-8df1-48bb-9659-9e7ab7a9d2e8",
   "metadata": {},
   "source": [
    "## Short reads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7305a8a1-d496-4674-8a44-435b121fbe98",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use megablast to check unclassified results from plasflow\n",
    "run_nt_blast(bgi_arg_related_reads_plasflow_unclassified, bgi_arg_related_reads_plasflow_unclassified_blast,)\n",
    "bgi_arg_related_reads_plasflow_unclassified_blast_dict = get_blast_names(bgi_arg_related_reads_plasflow_unclassified_blast)\n",
    "bgi_arg_related_reads_plasflow_unclassified_blast_location_dict = (get_read_location_from_blast(bgi_arg_related_reads_plasflow_unclassified_blast_dict))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb70a8d9-a3b0-456c-8c7b-e190d032fd58",
   "metadata": {},
   "source": [
    "# Combine Plasflow, blast and MGE information. Add abundance calculation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be26761b-c54f-4599-bc9e-dda609afea2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# map reads to contigs\n",
    "subprocess.run([\"bowtie2-build\", bgi_contigs, bgi_bowtie2_index], stdout=subprocess.DEVNULL)\n",
    "subprocess.run([\"bowtie2\", \"-x\", bgi_bowtie2_index, \"-1\", short_reads_1_location, \"-2\",  short_reads_2_location, \"-S\",  bgi_bowtie2_output, \"-p\", str(threads)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3d73fba-457c-461c-9df9-7af5952713bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "pysam.sort(\"-@\", str(threads), \"-o\", bgi_reads_contig_mapping, bgi_bowtie2_output)\n",
    "pysam.index(bgi_reads_contig_mapping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a09121a-7ab2-4328-98e6-0cc142022d11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get short read contigs average per base coverage\n",
    "bgi_read_cov = get_short_read_cov(\n",
    "    rgi_filtered,\n",
    "    bgi_contig_regions_bed,\n",
    "    bgi_total_contig_cov,\n",
    "    bgi_reads_contig_mapping,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e113865-e648-4adf-a208-1d8e12be2029",
   "metadata": {},
   "outputs": [],
   "source": [
    "ont_loc_abundance_report, ont_chromosome_list = get_location_and_abundance(\n",
    "    ont_arg_related_reads_plasflow,\n",
    "    ont_arg_best_hit,\n",
    "    ont_mge_best_hit,\n",
    "    ont_arg_related_reads_plasflow_unclassified_blast_location_dict,\n",
    ")\n",
    "write_list_to_file(ont_chromosome_list, ont_chromosome_list_file)  \n",
    "ont_loc_abundance_report.to_csv(ont_abundance_report)   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1b1527c-6a96-4755-af2c-afcd78121850",
   "metadata": {},
   "outputs": [],
   "source": [
    "bgi_loc_abundance_report, bgi_chromosome_list = get_location_and_abundance(\n",
    "    bgi_arg_related_reads_plasflow,\n",
    "    rgi_filtered,\n",
    "    bgi_mge_best_hit,\n",
    "    bgi_arg_related_reads_plasflow_unclassified_blast_location_dict,\n",
    "    lr=False,\n",
    "    sr_avg_read_cov=bgi_read_cov,\n",
    ")\n",
    "write_list_to_file(bgi_chromosome_list, bgi_chromosome_list_file) \n",
    "bgi_loc_abundance_report.to_csv(bgi_abundance_report)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03af6980-eacc-4da4-aad5-9a9d7bcd7568",
   "metadata": {},
   "source": [
    "# MOB-suite for plasmid annotation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ccb16be-1f12-4fe5-9cf7-3fa365a5bc58",
   "metadata": {},
   "outputs": [],
   "source": [
    "ont_plasmid_read_list = extract_plasmid_read_id_list(\n",
    "    ont_arg_related_reads_plasflow,\n",
    "    ont_arg_related_reads_plasflow_unclassified_blast_location_dict,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aabe15c2-ee6e-4800-a8fe-b5d0136973e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "bgi_plasmid_read_list = extract_plasmid_read_id_list(\n",
    "    bgi_arg_related_reads_plasflow,\n",
    "    bgi_arg_related_reads_plasflow_unclassified_blast_location_dict,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "145dbf5c-1d4d-43a0-ba36-57ccd6ab9c17",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_candidate_list_to_file(ont_plasmid_reads_list_file, ont_plasmid_read_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21d9c2c2-f81b-4a25-bf9c-f497142e7f79",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_candidate_list_to_file(bgi_plasmid_reads_list_file, bgi_plasmid_read_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f6c9951-d943-4d89-a7d2-0628da7c1ad7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract ONT plasmid reads\n",
    "run_seqtk_subseq(ont_reads, ont_plasmid_reads_list_file, ont_plasmid_reads)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41bc4477-ee4e-48e1-b844-ccee3a541cdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract short bgi plasmid reads\n",
    "run_seqtk_subseq(bgi_contigs, bgi_plasmid_reads_list_file, bgi_plasmid_reads)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "287c9772-7e87-463c-a51d-a22c66526534",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run mob suite\n",
    "# NEED TO RUN MOB INIT ON COMMAND LINE FIRST\n",
    "os.system(f\"mob_typer -i {os.path.abspath(ont_plasmid_reads)} -o {os.path.abspath(ont_mob_suite)} -n {threads} -x -f\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d18b5d53-62e8-4af7-a755-c4a81515b4f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.system(f\"mob_typer -i {os.path.abspath(bgi_plasmid_reads)} -o {os.path.abspath(bgi_mob_suite)} -n {threads} -x -f\")\n",
    "# run_mob_suite(bgi_plasmid_reads, bgi_mob_suite)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "489d1501-2818-4ade-b27e-bb008ff1a12d",
   "metadata": {},
   "source": [
    "# Centrifuge for classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "654c11cc-2952-4a4a-95d9-944640d61da1",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_centrifuge(\n",
    "    ont_arg_related_reads,\n",
    "    ont_centrifuge_classification_out,\n",
    "    ont_centrifuge_classification_report,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92ceb34d-c864-44a2-abe5-26b6a0c5c4f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_centrifuge(\n",
    "    bgi_arg_related_contigs,\n",
    "    bgi_centrifuge_classification_out,\n",
    "    bgi_centrifuge_classification_report,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d42483ea-8599-429a-9ed6-58b5746b4502",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}