{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "import ipyparallel as ipp\n",
    "import os, time\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "import numbers\n",
    "from sklearn import preprocessing\n",
    "from subprocess import Popen, PIPE\n",
    "import seaborn as sns\n",
    "from IPython.display import FileLink\n",
    "import urllib.request as urllib2\n",
    "import dill\n",
    "import traceback\n",
    "from pandas import Series, DataFrame\n",
    "import gzip\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore',category=pd.io.pytables.PerformanceWarning)\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "from Bio import SeqIO\n",
    "import pysam\n",
    "from collections import OrderedDict, namedtuple\n",
    "import operator\n",
    "import multiprocessing as mp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vcfutils = \"perl /home/menonm2/g/src/bcftools-1.3/vcfutils.pl\"\n",
    "vcftools = \"/home/menonm2/g/bin/vcftools\"\n",
    "bcftools = \"/home/menonm2/g/src/bcftools-1.3.1/bcftools\"\n",
    "tabix = \"/home/menonm2/g/src/dDocent_run/freebayes/SeqLib/htslib/tabix\"\n",
    "bgzip = \"/home/menonm2/g/src/dDocent_run/freebayes/SeqLib/htslib/bgzip\"\n",
    "java  = \"/home/menonm2/jdk1.8.0_91/bin/java\"\n",
    "plink = \"/home/menonm2/g/src/plink-1.07-x86_64/plink --noweb\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load in the output file generated from dDocent run. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis_dir = '/home/menonm2/eckertlab/Mitra/chapter2/dDocent3489_SNPs_JCB'\n",
    "\n",
    "vcf_file = os.path.join(analysis_dir, \"TotalRawSNPs.vcf\")\n",
    "assert os.path.exists(vcf_file)\n",
    "vcf_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!$vcftools --remove-indels \\\n",
    "--max-missing 0.5 \\\n",
    "--min-alleles 2 \\\n",
    "--max-alleles 2 \\\n",
    "--remove-filtered-all \\\n",
    "--recode \\\n",
    "--recode-INFO-all \\\n",
    "--gzvcf \\\n",
    "$vcf_file \\\n",
    "--out $vcf_file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just renaming files "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vcf_filtered = \"%s.recode.vcf\" % vcf_file\n",
    "vcf_filtered_gz = \"%s.gz\" % vcf_filtered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!$bgzip -c $vcf_filtered > {vcf_filtered_gz}\n",
    "!$tabix {vcf_filtered_gz}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Functions to obtain several summary stats to perform further custom filtering on the vcf file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_vcf_stats(vcf_gz):\n",
    "    \n",
    "    stats = ['depth',\n",
    "            'site-depth',\n",
    "            'site-mean-depth',\n",
    "            'site-quality',\n",
    "            'missing-indv',\n",
    "            'missing-site',\n",
    "            'freq',\n",
    "            'counts',\n",
    "            'hardy',\n",
    "            'het']\n",
    "    \n",
    "    for stat in stats:\n",
    "        !$vcftools --gzvcf $vcf_gz \\\n",
    "        --out $vcf_gz \\\n",
    "        {\"--%s\" % stat} "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.set_option('display.max_columns', 100)\n",
    "\n",
    "def get_MAF(row):\n",
    "    try:\n",
    "        return np.min([row.A1_freq, row.A2_freq])\n",
    "    except:\n",
    "        print(row)\n",
    "        \n",
    "def get_correction(n):\n",
    "    #for finite sample size\n",
    "    return (2*n)/(2*n-1)\n",
    "\n",
    "def calculate_Fis(vals):\n",
    "    try:\n",
    "        data = [float(x) for x in vals.split(\"/\")]\n",
    "        assert len(data) == 3\n",
    "        num_individuals = np.sum(data)\n",
    "        total_alleles = 2*num_individuals\n",
    "        a1_count = 2*data[0]\n",
    "        a2_count = 2*data[2]\n",
    "        het_count = data[1]\n",
    "        a1_count += het_count\n",
    "        a2_count += het_count\n",
    "        a1_freq = a1_count/total_alleles\n",
    "        a2_freq = a2_count/total_alleles\n",
    "        assert a1_freq + a2_freq == 1.0\n",
    "        He = 2 * a1_freq * a2_freq * get_correction(num_individuals)\n",
    "        Ho = het_count/num_individuals\n",
    "        Fis = 1 - (Ho/He)\n",
    "        return Fis\n",
    "    except:\n",
    "        return -9\n",
    "\n",
    "def combine_vcf_stats(filedir, prefix):\n",
    "    hardy_files = !ls {filedir}/{prefix}*.hwe\n",
    "    hardy = pd.read_csv(hardy_files[0], sep=\"\\t\")\n",
    "\n",
    "    hardy.columns = ['CHROM', 'POS', 'OBS(HOM1/HET/HOM2)', 'E(HOM1/HET/HOM2)', 'ChiSq_HWE',\n",
    "       'P_HWE', 'P_HET_DEFICIT', 'P_HET_EXCESS']\n",
    "    hardy.index = hardy.apply(lambda x: \"%s-%d\" % (x.CHROM, x.POS), axis=1)\n",
    "    \n",
    "    loci_files = !ls {filedir}/{prefix}*.l* | grep -v log\n",
    "    loci_df = pd.concat([pd.read_csv(x, sep=\"\\t\", skiprows=0) for x in loci_files], axis=1)\n",
    "    chrom_pos = loci_df.ix[:,0:2]\n",
    "    \n",
    "    frq_files = !ls {filedir}/{prefix}*.frq* | grep -v count\n",
    "    frq_data = []\n",
    "    h = open(frq_files[0])\n",
    "    header = h.readline().strip().split()\n",
    "    for line in h:\n",
    "        frq_data.append(line.strip().split('\\t'))\n",
    "\n",
    "    header = ['CHROM', 'POS', 'N_ALLELES', 'N_CHR', 'A1_FREQ', \"A2_FREQ\"]\n",
    "    frq_df = pd.DataFrame(frq_data)\n",
    "    print(frq_df.columns)\n",
    "    #frq_df = frq_df.drop([6,7],axis=1)\n",
    "    frq_df.columns = header\n",
    "    frq_df.index = frq_df.apply(lambda x: \"%s-%s\" % (x.CHROM, x.POS), axis=1)\n",
    "    \n",
    "    loci_df = loci_df.drop(['CHROM','CHR','POS'], axis=1)\n",
    "    loci_df = pd.concat([chrom_pos, loci_df], axis=1)\n",
    "    loci_df.index = loci_df.apply(lambda x: \"%s-%d\" % (x.CHROM, x.POS), axis=1)\n",
    "    \n",
    "    loci_df = pd.concat([loci_df, frq_df, hardy], axis=1)\n",
    "    loci_df[\"A1_allele\"] = loci_df.apply(lambda row: row.A1_FREQ.split(\":\")[0], axis=1)\n",
    "    loci_df[\"A2_allele\"] = loci_df.apply(lambda row: row.A2_FREQ.split(\":\")[0], axis=1)\n",
    "    \n",
    "    loci_df[\"A1_freq\"] = loci_df.apply(lambda row: float(row.A1_FREQ.split(\":\")[1]), axis=1)\n",
    "    loci_df[\"A2_freq\"] = loci_df.apply(lambda row: float(row.A2_FREQ.split(\":\")[1]), axis=1)\n",
    "    \n",
    "    loci_df['MAF'] = loci_df.apply(get_MAF, axis=1)\n",
    "    loci_df = loci_df.drop(['CHROM', 'POS'], axis=1)\n",
    "    \n",
    "    loci_df['Fis'] = loci_df['OBS(HOM1/HET/HOM2)'].apply(calculate_Fis)\n",
    "    \n",
    "    return loci_df, frq_df, hardy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "get_vcf_stats(vcf_filtered_gz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loci_df, frq_df, hardy = combine_vcf_stats(analysis_dir, \"TotalRawSNPs.vcf.recode.vcf\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loci_df.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loci_df.to_csv(os.path.join(analysis_dir, \"loci_stats_raw.txt\"),\n",
    "              sep=\"\\t\",\n",
    "              index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I need to generate a map of all the position information stored in the vcf file to be able use plink. I am using plink to get a 012 file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chroms = sorted(set([x.split(\"-\")[0] for x in loci_df.index]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(analysis_dir, \"chrom_map_raw.txt\"), \"w\") as o:\n",
    "    for i, c in enumerate(chroms):\n",
    "        o.write(\"%s\\t%d\\n\" % (c, i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_plink_files(vcf_gz):\n",
    "    !$vcftools --gzvcf {vcf_gz} \\\n",
    "    --out {vcf_gz} \\\n",
    "    --plink \\\n",
    "    --chrom-map {os.path.join(analysis_dir, \"chrom_map_raw.txt\")}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "write_plink_files(vcf_filtered_gz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_plink_recode(vcf_gz):\n",
    "    !$plink --recodeA --tab --file {vcf_gz} --out {vcf_gz}_recodeA\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "write_plink_recode(vcf_filtered_gz)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we will use the outputs from the custom functions above to filter SNPs. \n",
    "First look at depth distribution across SNPs to filter based on depth. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loci_df.SUM_DEPTH.describe() "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look at distribution of Fis values, phred score and MAF across SNPs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(loci_df[loci_df.Fis == -9])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(loci_df[loci_df.QUAL >= 10]) - len(loci_df[loci_df.QUAL >= 20])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(loci_df[loci_df.QUAL < 20]), len(loci_df[loci_df.QUAL < 10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(loci_df[loci_df.Fis >= 0.5]), len(loci_df[loci_df.Fis <= -0.5]), len(loci_df[loci_df.MAF < 0.002])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform the filtering based on these stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##MAF should be set lower than the pop with smallest no of indvs\n",
    "#changed depth to match the values in loci_df.SUM_DEPTH min and 75th percentile or 50th percentile\n",
    "def filter_snps(df):\n",
    "        return df[(df.SUM_DEPTH >= 600) & \n",
    "                  (df.SUM_DEPTH < 3931) & \n",
    "                  (df.QUAL >= 20) & \n",
    "                  (df.MAF >= 0.001) & \n",
    "                  (df.Fis < 0.5) & \n",
    "                  (df.Fis > -0.5)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loci_stage1 = filter_snps(loci_df)\n",
    "loci_stage1.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(analysis_dir, \"75perc_positions.txt\"), \"w\") as o:\n",
    "    for elem in loci_stage.index:\n",
    "        o.write(\"%s\\n\" % \"\\t\".join(elem.split(\"-\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for d, vcf_gz in zip([analysis_dir], [vcf_filtered_gz]):\n",
    "    !$vcftools --gzvcf $vcf_gz \\\n",
    "    --remove-indels  \\\n",
    "    --remove-filtered-all \\\n",
    "    --recode \\\n",
    "    --recode-INFO-all \\\n",
    "    --positions {os.path.join(d, \"Raw75perc_positions.txt\")} \\\n",
    "    --out {os.path.join(d, \"Rawgood_snps75perc\")}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for d in [analysis_dir]:\n",
    "    snps = os.path.join(d, \"good_snps75perc.recode.vcf\")\n",
    "    snps_gz = snps + \".gz\"\n",
    "    !$bgzip -c {snps} > {snps_gz}\n",
    "    !$tabix {snps_gz}\n",
    "    \n",
    "    srted = snps_gz + \"_sorted.vcf\"\n",
    "    srted_gz = srted + \".gz\"\n",
    "    !vcf-sort {snps_gz} > {srted}\n",
    "    !$bgzip -c {srted} > {srted_gz}\n",
    "    !$tabix {srted_gz}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for d in [analysis_dir]:\n",
    "    f = os.path.join(d, \"good_snps75perc.recode.vcf.gz_sorted.vcf.gz\")\n",
    "    assert os.path.exists(f)\n",
    "    write_plink_recode(f) #will return the final file in 012 format "
   ]
  }
 ],
 "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}