#!/usr/bin/env python

'''
Generate barplot with population distribution from the PopulationDistribution.txt file. Create a barplot for each total value, dividing
into 0 bulge, 1 bulge, 2 bulge etc
Example of PopulationDistribution file (with 1 bulge max):
-Summary_CCATCGGTGGCCGTTTGCCCNNN
EAS     0,0     0,0     0,0     0,0     0,4     0,22    0,0     0,0     0,0     0,0
EUR     0,0     0,0     0,0     0,0     2,1     0,28    0,0     0,0     0,0     0,0
AFR     0,0     1,0     0,1     0,0     1,6     0,42    0,0     0,0     0,0     0,0
AMR     0,0     0,0     0,0     0,0     2,4     0,30    0,0     0,0     0,0     0,0
SAS     0,0     0,0     0,0     0,0     1,1     0,34    0,0     0,0     0,0     0,0
-Summary_GAGTCCGAGCAGAAGAAGAANNN
EAS     0,0     0,0     0,0     1,1     1,18    0,147   0,0     0,0     0,0     0,0
EUR     0,0     0,0     0,0     1,1     3,20    0,128   0,0     0,0     0,0     0,0
AFR     0,0     0,0     0,0     1,1     3,31    0,253   0,0     0,0     0,0     0,0
AMR     0,0     0,0     0,0     1,1     3,26    0,146   0,0     0,0     0,0     0,0
SAS     0,0     0,0     0,0     2,1     3,21    0,153   0,0     0,0     0,0     0,0
'''

# argv 1 is jobid.PopulationDistribution.txt
# argv 2 is total value
# argv 3 is Guide
from matplotlib.legend_handler import HandlerTuple
import matplotlib.colors as mc
import colorsys
from operator import itemgetter, truediv
from os.path import isfile, join
from os import listdir
import os
import warnings
import glob
from itertools import islice
import sys

from pandas.core.tools import numeric
import numpy as np
import scipy.spatial.distance as sp
from math import pi
import pandas as pd
from matplotlib import patches as mpatches
from matplotlib import pyplot as plt
import math
import matplotlib
from matplotlib.ticker import ScalarFormatter
# matplotlib.use("TkAgg")
matplotlib.use('Agg')

warnings.filterwarnings("ignore")


def adjust_lightness(color, amount=0.5):
    try:
        c = mc.cnames[color]
    except:
        c = color
    c = colorsys.rgb_to_hls(*mc.to_rgb(c))
    return colorsys.hls_to_rgb(c[0], max(0, min(1, amount * c[1])), c[2])


plt.rcParams["figure.dpi"] = 400
# matplotlib.rcParams["figure.figsize"] = 10, 3
plt.rcParams['pdf.fonttype'] = 42
plt.rcParams['ps.fonttype'] = 42
plt.style.use('seaborn-poster')

# matplotlib.rcParams["figure.figsize"] = 7, 5
ax = plt.figure()

# ax.set_figheight(5)
# ax.set_figwidth(10)

titlesize = 18
fontsize = 17


barplot_values = dict()  # barplot_values -> EAS -> [1,2]
total = int(sys.argv[2])
guide = sys.argv[3]
max_value = 0
number_bars = 0
previous_bar = []
with open(sys.argv[1]) as summary:
    for line in summary:
        if guide in line:
            while True:
                try:
                    line = next(summary).strip()
                except:
                    break
                if '-Summary_' in line:
                    break
                line = line.split('\t')
                number_bars = len(line[total + 1].split(','))
                # line = EAS 0,7 1,2 5,3 10,11
                barplot_values[line[0]] = dict()
                for count in range(0, total+1):
                    barplot_values[line[0]][count] = [
                        int(x) for x in line[count+1].split(',')]
                    value = sum([int(x) for x in line[count+1].split(',')])
                # previous_bar.append(0)
                if value > max_value:
                    max_value = value

totaldict = barplot_values
barplot_values_tmp = dict()
lower_barplot_values = dict()
for pop in barplot_values:
    barplot_values_tmp[pop] = barplot_values[pop][total]
    sum_up_list = [0]*number_bars
    for count in range(total):
        for elem in range(number_bars):
            sum_up_list[elem] += barplot_values[pop][count][elem]
    summs = 0
    for elem in sum_up_list:
        summs += elem
    lower_barplot_values[pop] = summs

barplot_values = barplot_values_tmp

#print('total', totaldict)
#print('original barplot', barplot_values)
#print('all lower sum up', lower_barplot_values)
# number_bars = number_bars+1  # add bar with all the lower levels of data
# [0 1 2 3 4 5] #NOTE 0 is REFERENCE
ind = np.arange(0, len(barplot_values.keys()), 1)
no_result = False
try:
    max_value += max(lower_barplot_values.values())
    y_range = np.arange(
        0, max_value + math.ceil(max_value/10), math.ceil(max_value/5))
except:
    y_range = np.arange(0, 1, 1)
    no_result = True
width = 0.5

population_color = ['#a6cee3', '#1f78b4',
                    '#b2df8a', '#33a02c', '#fb9a99', '#e31a1c']
# hatches = ['/', '\\', '|', '-', '+', 'x', 'o', 'O', '.', '*']
# if 'REF' in barplot_values.keys():
#     population_color.insert(0, 'grey')
# If i have more than 5 superpopulations + ref
if len(barplot_values.keys()) > len(population_color):
    for i in mc.TABLEAU_COLORS:
        population_color.append(i)
        if len(barplot_values.keys()) == len(population_color):
            break
# If i have less than 5 superpopulations + ref
elif len(barplot_values.keys()) < len(population_color):
    population_color = population_color[:len(barplot_values.keys())]

all_bar = []
lower_value_check = False
# add data for lower count data
for value in lower_barplot_values.values():
    previous_bar.append(value)

if max(previous_bar) > 0:
    lower_value_check = True

lower_bar = []
if lower_value_check:
    lower_bar.append(plt.bar(ind, previous_bar, width, color='yellow',
                             align='edge', edgecolor='black'))

for i in range(number_bars):  # For 0 bulge, 1 bulge, 2 bulge ...
    # insert current total bars
    current_bar = [x[i] for x in barplot_values.values()]
    # all_bar.append(plt.bar(ind, current_bar, width, color=[adjust_lightness(x, 1 + i*0.3) if current_bar[pos] != 0 else 'white' for pos, x in enumerate(population_color)], align='edge', bottom=previous_bar, edgecolor='black'))
    all_bar.append(plt.bar(ind, current_bar, width,
                           color=population_color[i], align='edge', bottom=previous_bar, edgecolor='black'))
    previous_bar = [x + previous_bar[k] for k, x in enumerate(current_bar)]

# p1 = plt.bar(ind, barplot_values.values(), width, color=population_color, align='edge')       #color = '#67a9cf'

# p2 = plt.bar(ind, barplot_values.values() , width, color=[adjust_lightness(x, 1.3) for x in population_color], align='edge', bottom = list(barplot_values.values()))

# #print('number bars', number_bars)
legend_labels = []
handles_color = []
if lower_value_check:
    legend_labels.append('MM+B < '+str(total))
    handles_color.append(lower_bar[0][0])
# for x in range(min(number_bars, total+1)):
for x in range(number_bars):
    if x == 0:
        if barplot_values['REF'][0] != 0:
            # got Total mms and 0 bulges
            legend_labels.append(str(total) + ' MM')
            # handles_color.append((all_bar[0][:]))
            handles_color.append(all_bar[0][0])
    else:
        if (total - x) < 0:  # To avoid negative numbers on MM values
            if barplot_values['REF'][x] != 0:
                legend_labels.append('0 MM + ' + str(x) + ' B')
            # handles_color.append((all_bar[x][:]))
                handles_color.append(all_bar[x][0])
        else:
            if barplot_values['REF'][x] != 0:
                legend_labels.append(str(total - x) + ' MM + ' + str(x) + ' B')
                handles_color.append(all_bar[x][0])

legend_labels.reverse()
# handles_color = [(x[:]) for x in all_bar]
handles_color.reverse()
plt.legend(handles_color, legend_labels, fontsize=fontsize, handlelength=5,
           handler_map={tuple: HandlerTuple(ndivide=None)}, title="MM mismatches, B bulges", title_fontsize=fontsize)
# first param is for the colored rectangles of legens, second parameter for labels, handlelength is size of rectangles, handlermap is for grouping different colors in single label
# [(first bar color, second bar color, ...), (first bar light color, second bar light color,...)]
plt.title('Targets with up to ' +
          str(total)+' mismatches and/or bulges by superpopulation', size=titlesize)

if no_result:
    plt.annotate('No targets found with ' + str(total) + ' mismatches + bulges',
                 (2.5, 0), size=fontsize, ha='center', va='center')  # 2.5 is x position
    sys.exit()


plt.xticks(ind+0.25, barplot_values.keys(), size=fontsize)

size_y_ticks = fontsize
digits = int(math.log10(max_value))+1
if digits > 5:  # Reduce dimension of y label because it can exceed plot size
    size_y_ticks = max(16, size_y_ticks - (2*(digits-5)))

plt.yticks(y_range, size=size_y_ticks)  # , size=size_y_ticks)

plt.tight_layout()
# plt.subplots_adjust(top=0.95, bottom=0.06, left=0.1, right=0.99)
plt.savefig("populations_distribution_" + guide + '_' +
            str(total) + "total" + ".png", format='png')
