{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "290964cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "# show plot in the notebook\n",
    "# allow reload the notebook whenever changes are made\n",
    "#%reload_ext autoreload\n",
    "#%autoreload 2\n",
    "import matplotlib.lines as lines\n",
    "from matplotlib import collections  as mc\n",
    "from collections import *\n",
    "from matplotlib.pyplot import figure\n",
    "from matplotlib.patches import Rectangle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b34cf258",
   "metadata": {},
   "outputs": [],
   "source": [
    "def getChrSize(file):\n",
    "    '''\n",
    "    chromosome size file\n",
    "    return a dict {'LG1': XXXXX}\n",
    "    '''\n",
    "    chrSizeDict = {}\n",
    "    with open(file) as inf:\n",
    "        for line in inf:\n",
    "            line = line.rstrip()\n",
    "            chrom, length = line.split()\n",
    "            chrSizeDict[chrom] = length\n",
    "    return chrSizeDict\n",
    "\n",
    "def getchrPos(LG1=[21.6, 22.4]): # default starting point\n",
    "    #chr1=[-0.4,0.4]\n",
    "    '''\n",
    "    {'chr23': [-0.4, 0.4], 'chr22': [0.6, 1.4], 'chr21': [1.6, 2.4], 'chr20': [2.6, 3.4], \n",
    "    'chr19': [3.6, 4.4], 'chr18': [4.6, 5.4], 'chr17': [5.6, 6.4], 'chr16': [6.6, 7.4], \n",
    "    'chr15': [7.6, 8.4], 'chr14': [8.6, 9.4], 'chr13': [9.6, 10.4], 'chr12': [10.6, 11.4], \n",
    "    'chr11': [11.6, 12.4], 'chr10': [12.6, 13.4], 'chr9': [13.6, 14.4], 'chr8': [14.6, 15.4], \n",
    "    'chr7': [15.6, 16.4], 'chr6': [16.6, 17.4], 'chr5': [17.6, 18.4], 'chr4': [18.6, 19.4], \n",
    "    'chr3': [19.6, 20.4], 'chr2': [20.6, 21.4], 'chr1': [21.6, 22.4]}\n",
    "    '''\n",
    "    chrPosList = {}\n",
    "    for i in range(1,24):\n",
    "        chrPosList['LG{}'.format(i)] = [x - 1 * (i-1) for x in LG1]\n",
    "    return chrPosList\n",
    "\n",
    "def getDNMPos(DNM_pos_file):\n",
    "    '''\n",
    "    return [chr, pos, times] \n",
    "    '''\n",
    "    DNMPosList = []\n",
    "    DNMPosCountDict = {}\n",
    "    with open(DNM_pos_file) as inf:\n",
    "        for line in inf:\n",
    "            line = line.rstrip()\n",
    "            DNMPosList.append(tuple(line.split()[:2]))\n",
    "            DNMPosCountDict = Counter(DNMPosList)\n",
    "            \n",
    "    return DNMPosCountDict\n",
    "\n",
    "def addLines(DNMbed):\n",
    "    \"bg line\"\n",
    "    chrPosList = getchrPos()\n",
    "    DNMPos = getDNMPos(DNMbed)\n",
    "    lines = []\n",
    "    for k,v in DNMPos.items():\n",
    "        chrom, pos = k\n",
    "        indCount = int(v)\n",
    "        lines.append([(pos, chrPosList[chrom][0]),(pos, chrPosList[chrom][1])])\n",
    "    return lines\n",
    "\n",
    "\n",
    "def addPropLines(DNMbed):\n",
    "    \"proportion line\"\n",
    "    chrPosList = getchrPos()\n",
    "    DNMPos = getDNMPos(DNMbed)\n",
    "    lines = []\n",
    "    for k,v in DNMPos.items():\n",
    "        chrom, pos = k\n",
    "        indCount = int(v)\n",
    "        lines.append([(pos, chrPosList[chrom][0]),(pos,chrPosList[chrom][0]+ 0.08*indCount)])\n",
    "    return lines\n",
    "\n",
    "\n",
    "def cdsReg(file):\n",
    "    #file = \"/Users/rainlam/Project/snpeff_analysis/FR26.dnm.annt.stats\"\n",
    "    type_snp_list = []\n",
    "    xyset = []\n",
    "    with open(file, \"r\") as inf:\n",
    "        for line in inf:\n",
    "            tmp = line.split(\"\\t\")\n",
    "            type_tmp = tmp[-1].split(\"|\")[:2]\n",
    "            #print(tmp[:4] + [type_tmp[-1]])\n",
    "            if \"missense\" in type_tmp[-1] or \"synonymous\" in type_tmp[-1]:\n",
    "                type_snp_list.append(tuple(tmp[:4] + [type_tmp[-1]]))\n",
    "    snp_list = list(set(type_snp_list))\n",
    "    for i in type_snp_list:\n",
    "        x = round(21.6 - int(i[0][2:]) + 1 - 0.02, 3)\n",
    "        y = int(i[1]) - 160000\n",
    "        yield x,y\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "7eb71455",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 585048 characters]",
      "text/plain": [
       "<Figure size 10000x8000 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Example data\n",
    "fig = plt.figure(figsize=(10,8), dpi = 1000)\n",
    "#ax = fig.add_axes([1, 1, 1, 1])\n",
    "ax = fig.add_subplot()\n",
    "chrSizeDict = getChrSize('guppy.chrom.sizes')\n",
    "#print(chrSizeDict)\n",
    "\n",
    "chrNames = chrSizeDict.keys()\n",
    "chroms = list(chrNames)[::-1]\n",
    "chr_len = [int(i) for i in chrSizeDict.values()][::-1]\n",
    "\n",
    "\n",
    "hbars = ax.barh(chroms, chr_len, align='center', color = 'grey', alpha = 0.1)\n",
    "ax.set_xticks([0,5000000,10000000,15000000,20000000,25000000,30000000,35000000,40000000,45000000])\n",
    "labels = [0,5,10,15,20,25,30,35,40,45]\n",
    "ax.set_xticklabels(labels)\n",
    "ax.set_xlim(0,45000000)\n",
    "#set_axis_off()\n",
    "ax.spines['top'].set_visible(False)\n",
    "#ax.spines['bottom'].set_visible(False)\n",
    "ax.spines['left'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "# 设置X坐标轴的高低\n",
    "ax.spines[\"bottom\"].set_position(['data',-0.5])\n",
    "\n",
    "#ax.tick_params(axis=u'both', which=u'both',length=0)\n",
    "ax.tick_params(axis='y', which=u'both',length=0)\n",
    "\n",
    "ax.set_xlabel('Genomic Distribution')\n",
    "#ax.text(45450000, -1.35, \"(Mb)\")\n",
    "\n",
    "#line = lines.Line2D([100,300],\n",
    "#                    [-0.2,0.20],\n",
    "#                    lw = 10, color ='green',\n",
    "#                    axes = ax, alpha = 0.3)\n",
    "\n",
    "FR24lines = addLines('./FR24.fam.dnm.bed')\n",
    "FR26lines = addLines('./FR26.fam.dnm.bed')\n",
    "FR31lines = addLines('./FR31.fam.dnm.bed')\n",
    "FR24linesProp = addPropLines('./FR24.fam.dnm.bed')\n",
    "FR26linesProp = addPropLines('./FR26.fam.dnm.bed')\n",
    "FR31linesProp = addPropLines('./FR31.fam.dnm.bed')\n",
    "\n",
    "#### frame CDS DNM\n",
    "\n",
    "#'5076629', 21.6\n",
    "\n",
    "\n",
    "#[(x0,y0),(x1,y1)]\n",
    "lc1 = mc.LineCollection(FR24linesProp, colors='orange', linewidths=3, alpha=1.0)\n",
    "lc2 = mc.LineCollection(FR26linesProp, colors='blue', linewidths=3, alpha=0.8)\n",
    "lc3 = mc.LineCollection(FR31linesProp, colors='green', linewidths=3, alpha=0.8)\n",
    "\n",
    "lc4 = mc.LineCollection(FR24lines, colors='orange', linewidths=3, alpha=0.2)\n",
    "lc5 = mc.LineCollection(FR26lines, colors='blue', linewidths=3, alpha=0.2)\n",
    "lc6 = mc.LineCollection(FR31lines, colors='green', linewidths=3, alpha=0.2)\n",
    "\n",
    "\n",
    "ax.add_collection(lc1)\n",
    "ax.add_collection(lc2)\n",
    "ax.add_collection(lc3)\n",
    "ax.add_collection(lc4)\n",
    "ax.add_collection(lc5)\n",
    "ax.add_collection(lc6)\n",
    "ax.legend(['Fam1','Fam2','Fam3'], title = 'Family',loc=5, ncol=1, frameon=False)\n",
    "\n",
    "currentAxis = plt.gca()\n",
    "# add frame\n",
    "\n",
    "for x,y in cdsReg(\"./FR24.dnm.annt.stats\"):\n",
    "    currentAxis.add_patch(Rectangle((y, x), 320000, 0.85, fill=None, color=\"orange\"))\n",
    "for x,y in cdsReg(\"./FR26.dnm.annt.stats\"):\n",
    "    currentAxis.add_patch(Rectangle((y, x), 320000, 0.85, fill=None, color=\"blue\"))\n",
    "for x,y in cdsReg(\"./FR31.dnm.annt.stats\"):\n",
    "    currentAxis.add_patch(Rectangle((y, x), 320000, 0.85, fill=None, color=\"green\"))\n",
    "\n",
    "#29586980\n",
    "#plt.scatter(22208413,17.6, marker='*', color=\"orange\") \n",
    "plt.scatter(249606, 18.6, marker='*', color=\"blue\")\n",
    "plt.scatter(28244719,19.6, marker='*', color=\"blue\")\n",
    "#plt.scatter(24652180,3.0, marker=\"v\", color=\"black\")\n",
    "#currentAxis.add_patch(Rectangle((5076629 - 160000, 21.58), 320000, 0.85, fill=None, alpha=0.5, color=\"green\"))\n",
    "#'8357440', 20.6\n",
    "# X, Y, width, hight\n",
    "#currentAxis.add_patch(Rectangle((8357440 - 160000, 20.58), 320000, 0.85, fill=None, alpha=0.5, color=\"green\"))\n",
    "\n",
    "\n",
    "#currentAxis = plt.gca()\n",
    "#currentAxis.add_patch(Rectangle((someX -.2, someY), 0.5, 0.3, fill=None, alpha=1, color=\"red\"))\n",
    "\n",
    "\n",
    "\n",
    "#ax.add_line(line)\n",
    "\n",
    "\n",
    "# add GC content <- may remove it later \n",
    "\"\"\"\n",
    "df = pd.read_csv(\"QuH-F25.1k\", sep='\\t') # read window-based GC content\n",
    "# LG23 [21.6, 22.4]\n",
    "for i in range(1,24):\n",
    "    newdf = df[df['CHROM'] == \"LG\" + str(i)]\n",
    "    x = newdf[\"START\"]\n",
    "    y = newdf[\"GC_wo_N\"] + 22.6 - i\n",
    "    ax.plot(x, y,\n",
    "        color='lightgrey', alpha = 0.8)\n",
    "\"\"\"\n",
    "\n",
    "    \n",
    "plt.savefig('/Users/rainlam/Desktop/DNMPlot.jpg', dpi = 1000, pad_inches=0 )\n",
    "plt.savefig('/Users/rainlam/Desktop/DNMPlot.pdf', dpi = 1000, pad_inches=0)\n",
    "\n",
    "plt.show()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
