#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
import pickle
import matplotlib.pyplot as plt

def error(string, code=1):
	sys.stderr.write('ERROR: ' + string + '\n')
	exit(code)

def log(string):
	sys.stderr.write('LOG: ' + string + '\n')

#####

data = pickle.load(open(sys.argv[1],'rb'))

# total_length = 27317
total_length = 27277

#gene_names = ['1a', '1ab', 'S', '4a', '4b', 'E', 'M', 'N']
#gene_borders = [(293,12550),(293,20568),(20570,24091),(24091,24492),(24482,24748),(24750,24983),(24995,25672),(25686,26855)]
# gene_colors = ['#ABDDA4FF', '#66C2A5FF', '#3288BDFF', '#D53E4FFF', '#F46D43FF', '#FDAE61FF', '#FEE08BFF', '#E6F598FF']
# !!! 1ab->1b, no 4b, fix position of E to 24588

gene_names = ['1a', '1b', 'S', '4', 'E', 'M', 'N']
# gene_borders = [(293,12550),(293,20568),(20570,24091),(24091,24587),(24588,24983),(24995,25672),(25686,26855)]
gene_borders = [(293,12550),(12550,20568),(20570,24091),(24091,24587),(24588,24983),(24995,25672),(25686,26855)]
gene_colors = ['#ABDDA4FF', '#66C2A5FF', '#3288BDFF', '#D53E4FFF', '#F46D43FF', '#FEE08BFF', '#E6F598FF']

gene_starts = [g[0] for g in gene_borders]
gene_ends = [g[1] for g in gene_borders]


#####

outn = 'coverage_1_2.pdf'

fig, axes = plt.subplots(2, sharex=True, figsize=(12,6), gridspec_kw = {'height_ratios':[10, 1]})
plt.subplots_adjust(hspace=0)
ax = axes[0]
an = axes[1]

ax.set_yscale('log')

# grid
ax.grid(alpha=0.5, linewidth=0.5)

# remove axis lines
#plt.axis('off')
ax.spines['top'].set_visible(False)
#ax.spines['left'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['bottom'].set_visible(False)
#ax.yaxis.set_visible(False)

for sample in data:
	ax.plot(data[sample][:total_length+1], label=sample, linewidth=0.3)

ax.legend(loc='best', framealpha=1)
ax.set_title('Reference genome coverage of both samples')
ax.set_ylabel('coverage')

# annot
an.hlines([-1.5,-2.5,-1.5,-2.5,-1.5,-2.5,-1.5], gene_starts, gene_ends, colors=gene_colors, linewidths=14)
an.hlines(-2, 1, total_length, colors='k', linewidths=1)
an.vlines([1,total_length], -1, -3, linewidths=1)

an.yaxis.set_visible(False)
an.spines['top'].set_visible(False)
an.spines['left'].set_visible(False)
an.spines['right'].set_visible(False)
an.spines['bottom'].set_visible(False)

an.set_xlabel('genome position')

# ticks
plt.xticks([1,293]+gene_ends+[total_length], [1]+[""]*(len(gene_ends)+1)+[total_length])


plt.savefig(outn, bbox_inches='tight')

