#!/usr/bin/env python
import os
import sys
import matplotlib.cm
import matplotlib.colors
import numpy as np
import pandas as pd
import gzip

# USAGE
USAGE_STATEMENT = "USAGE: python bed2Ideogram.py [in_bed] [out_html] [TYPE={" \
				  "histogram|overlay|heatmap|heatmap_inverted|heatmap_normalized|heatmap_normalized_inverted" \
				  "|heatmap_thresholded|heatmap_zscore|heatmap_chrZScore|heatmap_percentile|heatmap_chrPercentile" \
				  "|heatmap_chrPDF}] [track_description] [ref_genome] [Color or Colormap] [thresholded lower lim] [" \
				  "thresholded upper lim]"

# default colors
DEFAULT_COLOR = '#77F'
DEFAULT_COLORMAP = 'Oranges'
POSS_PLOTTYPES = ['histogram', 'overlay', 'heatmap', 'heatmap_inverted', 'heatmap_normalized',
				  'heatmap_normalized_inverted', 'heatmap_thresholded', 'heatmap_zscore', 'heatmap_chrZScore',
				  'heatmap_percentile', 'heatmap_chrPercentile', 'heatmap_chrPDF']

# template for annotation objects
ANNO_TEMPLATE = (
	'\t{\n'
	'\t\tname: \'{NAME}\',\n'
	'\t\tchr: \'{CHR}\',\n'
	'\t\tstart: {START},\n'
	'\t\tstop: {END},\n'
	'\t\tcolor: \'{COLOR}\''
	'\t},\n'
)

# overall html template
HTML_TEMPLATE = (
	'<head>\n'
	'\t<script src="https://unpkg.com/ideogram@1.0.0/dist/js/ideogram.min.js"></script>\n'
	'</head>\n'
	'<body>\n'
	'<h1>{TRACK_DESC}</h1>\n'
	'<script>\n'
	'\tvar ideogram = new Ideogram({\n'
	'\t\torganism: \'human\',\n'
	'\tannotations: [\n'
	'{ANNOTATIONS}'
	'],\n'
	'\tannotationsLayout: \'{PLOT_TYPE}\',\n'
	'\torientation: \'horizontal\',\n'
	'\tchrHeight: 900,\n'
	'\tresolution: 550,\n'
	'\tassembly: \'{REF_GENOME}\',\n'
	'\t});\n'
	'</script>\n'
	'</body>'
)


# converts bed file to annotation object
def bed_line_to_anno_object(bed_line, cnt, chromosome_sizes_tbl, color, val):
	# extract line
	line_parts = [a.strip() for a in line.split('\t')]
	chrs = line_parts[0]
	if 'chr' in chrs:
		chrs = chrs[3:]
	start = int(line_parts[1])
	end = int(line_parts[2])

	# make sure in bounds
	if ('chr'+chrs) in chromosome_sizes_tbl.Chr.values:
		chrom_size = chromosome_sizes_tbl[chromosome_sizes_tbl.Chr == ('chr'+chrs)].Size.values[0]
		if start > chrom_size:
			start = chrom_size
		if end > chrom_size:
			end = chrom_size
		if start < 1:
			start = 1
		if end < 1:
			end = 1

	# fill annotation template and return
	rtn = ANNO_TEMPLATE
	name = 'region' + str(cnt)
	if str(val) != 'nan':
		name = name + '(val=' + str(val) + ')'
	rtn = rtn.replace('{NAME}', name)
	rtn = rtn.replace('{CHR}', chrs)
	rtn = rtn.replace('{START}', str(start))
	rtn = rtn.replace('{END}', str(end))
	rtn = rtn.replace("{COLOR}", color)
	return rtn


# extract heat map color
def get_color(chrs, heat_val, cmap, data_stat, hm_flags):
	# extract data_stat
	data_min = data_stat[0]
	data_max = data_stat[1]
	data_mean = data_stat[2]
	data_median = data_stat[3]
	data_std = data_stat[4]
	data_percentiles = data_stat[5]
	chr_percentiles = data_stat[6]
	chr_means = data_stat[7]
	chr_stds = data_stat[8]
	chr_sums = data_stat[9]

	# hm flags
	inverted = hm_flags[0]
	thresholded = hm_flags[1]
	thresh = hm_flags[2]
	z_score = hm_flags[3]
	percentile = hm_flags[4]
	chr_percentile = hm_flags[5]
	chr_zscore = hm_flags[6]
	chr_pdf = hm_flags[7]

	# normalizations
	if thresholded:
		data_min = thresh[0]
		data_max = thresh[1]
		if heat_val < data_min:
			heat_val = data_min
		if heat_val > data_max:
			heat_val = data_max
	elif z_score:
		heat_val = (heat_val - data_mean) / data_std
		data_min = -3.0
		data_max = 3.0
		if heat_val < data_min:
			heat_val = data_min
		if heat_val > data_max:
			heat_val = data_max
	elif chr_zscore:
		data_min = -3.0
		data_max = 3.0
		if chr_stds[chrs] == 0:
			heat_val = 0.0
		else:
			heat_val = (heat_val - chr_means[chrs]) / chr_stds[chrs]
			if heat_val < data_min:
				heat_val = data_min
			if heat_val > data_max:
				heat_val = data_max
	elif percentile:
		data_min = 0.0
		data_max = 100.0
		heat_val = abs(heat_val - data_percentiles).argmin()
	elif chr_percentile:
		data_min = 0.0
		data_max = 100.0
		heat_val = abs(heat_val - chr_percentiles[chrs]).argmin()
	elif chr_pdf:
		data_min = 0.0
		data_max = 1.0
		if chr_sums[chrs] == 0:
			heat_val = 0.0
		else:
			heat_val = heat_val / chr_sums[chrs]
	normalized = (heat_val - data_min) / (data_max - data_min)
	if inverted:
		normalized = 1.0 - normalized
	rgba = cmap(normalized)
	return matplotlib.colors.to_hex(rgba)


def get_header_free_file_pointer(in_bed):
	if '.bgz' in in_bed or 'bed.gz' in in_bed:
		fin = gzip.open(in_bed, 'r')
		gzipped = True
	else:
		fin = open(in_bed, 'r')
		gzipped = False
	line = fin.readline()
	if gzipped:
		line = line.decode('UTF-8')
	possible_header_parts = line.split('\t')
	if possible_header_parts[1].isdigit():
		fin.close()
		if gzipped:
			fin = gzip.open(in_bed, 'r')
		else:
			fin = open(in_bed, 'r')
	return fin, gzipped


def get_empirical_statistics(in_bed):
	(fin, gzipped) = get_header_free_file_pointer(in_bed)
	data_min = np.inf
	data_max = -1 * np.inf
	data_sum = 0
	num_val = 0
	vals = list()
	chr_vals = dict()
	for line in fin:
		if gzipped:
			line = line.decode('UTF-8')
		line_parts = [a.strip() for a in line.split('\t')]
		chrs = line_parts[0]
		val = float(line_parts[3])
		if val < data_min:
			data_min = val
		if val > data_max:
			data_max = val
		vals.append(val)
		if chrs not in chr_vals:
			chr_vals[chrs] = list()
		chr_vals[chrs].append(val)
		data_sum = data_sum + val
		num_val = num_val + 1
	fin.close()
	data_mean = np.nanmean(vals)
	data_median = np.nanmedian(vals)
	data_std = np.nanstd(vals)
	data_percentiles = np.nanpercentile(vals, range(0, 101))
	chr_percentiles = dict()
	chr_means = dict()
	chr_stds = dict()
	chr_sums = dict()
	for chrs in chr_vals.keys():
		chr_percentiles[chrs] = np.nanpercentile(chr_vals[chrs], range(0, 101))
		chr_means[chrs] = np.nanmean(chr_vals[chrs])
		chr_stds[chrs] = np.nanstd(chr_vals[chrs])
		chr_sums[chrs] = np.nansum(chr_vals[chrs])
	return data_min, data_max, data_mean, data_median, data_std, data_percentiles, chr_percentiles, chr_means, chr_stds, chr_sums


# args
if len(sys.argv) < 5:
	print(USAGE_STATEMENT)
	sys.exit()

# basic args
in_bed = sys.argv[1]
out_HTML_file = sys.argv[2]
plot_type = sys.argv[3]
track_desc = sys.argv[4]
if plot_type not in POSS_PLOTTYPES:
	print(USAGE_STATEMENT)
	sys.exit()

# reference genome
ref_genome = 'hg19'
if len(sys.argv) >= 6:
	ref_genome = sys.argv[5]
if ref_genome == 'hg19':
	ref_genome_str = 'GRCh37'
elif ref_genome=='hg38':
	ref_genome_str = 'GRCh38'

# load chromsome sizes (based on ref genome)
if ref_genome=='hg19':
	chromosome_sizes_tbl = pd.read_csv(os.environ['HG19CHROMSIZES'],sep='\t',names=['Chr','Size'])
elif ref_genome=='hg38':
	chromosome_sizes_tbl = pd.read_csv(os.environ['HG38CHROMSIZES'],sep='\t',names=['Chr','Size'])

# color map / color of overlay
colormap = DEFAULT_COLORMAP
color = DEFAULT_COLOR
if len(sys.argv) >= 7:
	cc = sys.argv[5]
	if 'heatmap' in plot_type:
		colormap = cc
	else:
		color = cc

# heat map flags
if 'heatmap' in plot_type:

	# empirical statistics
	cmap = matplotlib.cm.get_cmap(colormap)
	emp = get_empirical_statistics(in_bed)

	# use empirical statistics
	data_min = emp[0]
	data_max = emp[1]
	data_mean = emp[2]
	data_median = emp[3]
	data_std = emp[4]
	data_percentiles = emp[5]
	chr_percentiles = emp[6]
	chr_means = emp[7]
	chr_stds = emp[8]
	chr_sums = emp[9]
	if 'inverted' in plot_type:
		inverted = True
	else:
		inverted = False
	if 'thresholded' in plot_type:
		thresholded = True
		thresh = [float(sys.argv[6]), float(sys.argv[7])]
	else:
		thresholded = False
		thresh = [-1, -1]
	if 'zscore' in plot_type:
		z_score = True
	else:
		z_score = False
	if 'chrZScore' in plot_type:
		chr_z_score = True
	else:
		chr_z_score = False
	if 'percentile' in plot_type:
		percentile = True
	else:
		percentile = False
	if 'chrPercentile' in plot_type:
		chr_percentile = True
	else:
		chr_percentile = False
	if 'chrPDF' in plot_type:
		chr_pdf = True
	else:
		chr_pdf = False
	hm_flags = (inverted, thresholded, thresh, z_score, percentile, chr_percentile, chr_z_score, chr_pdf)

# generate annotations
(fin, gzipped) = get_header_free_file_pointer(in_bed)
cnt = 0
temp_anno_file = out_HTML_file + '.anno'
temp_anno = open(temp_anno_file, 'w')
for line in fin:

	# decode line (if needed)
	if gzipped:
		line = line.decode('UTF-8')

	# generate parameters
	if 'heatmap' in plot_type:
		line_parts = [a.strip() for a in line.split('\t')]
		chrs = line_parts[0]
		val = float(line_parts[3])
		pt_color = get_color(chrs, val, cmap, emp, hm_flags)
		if z_score:
			val = (val - data_mean) / data_std
		elif chr_z_score:
			if chr_stds[chrs] == 0:
				val = 0.0
			else:
				val = (val - chr_means[chrs]) / chr_stds[chrs]
		elif percentile:
			val = abs(val - data_percentiles).argmin()
		elif chr_percentile:
			val = abs(val - chr_percentiles[chrs]).argmin()
		elif chr_pdf:
			if chr_sums[chrs] == 0:
				val = 0.0
			else:
				val = val / chr_sums[chrs]
	else:
		pt_color = color
		val = np.nan

	# generate anno object
	anno_text = bed_line_to_anno_object(line, cnt, chromosome_sizes_tbl, pt_color, val)

	#write to file
	temp_anno.write(anno_text)

	#inc
	cnt = cnt + 1
fin.close()
temp_anno.close()

# convert anno text to output html
anno_text = open(temp_anno_file, 'r').readlines()
anno_text = '\n'.join(anno_text)
anno_text = anno_text[0:-2]
out_HTML = HTML_TEMPLATE
out_HTML = out_HTML.replace('{ANNOTATIONS}', anno_text.rstrip())
ideogram_js_plot_type = plot_type
if 'heatmap' in plot_type:
	ideogram_js_plot_type = 'overlay'
	track_desc = track_desc + '<br>raw data statistics(min=' + str(data_min) + ', max=' + str(
		data_max) + ', mean=' + str(
		data_mean) + ', median=' + str(data_median) + ', std=' + str(data_std) + ')'
out_HTML = out_HTML.replace('{PLOT_TYPE}', ideogram_js_plot_type)
if 'zscore' in plot_type:
	track_desc = track_desc + '<br>Z-SCORE Scaled (3-sigma)'
elif 'chrZScore' in plot_type:
	track_desc = track_desc + '<br>CHROM Z-SCORE Scaled (3-sigma, per chromosome)'
elif 'percentile' in plot_type:
	track_desc = track_desc + '<br>PERCENTILE Scaled (0-100)'
elif 'chrPercentile' in plot_type:
	track_desc = track_desc + '<br>CHROM PERCENTILE Scaled (0-100, per chromosome)'
elif 'chrPDF' in plot_type:
	track_desc = track_desc + '<br>CHROM Probability Normalized Scaled<br>'
if 'inverted' in plot_type:
	track_desc = track_desc + '<br>INVERTED COLORING'
out_HTML = out_HTML.replace('{TRACK_DESC}', track_desc)
out_HTML = out_HTML.replace('{REF_GENOME}', ref_genome_str)

# actually write html file
fout = open(out_HTML_file, 'w')
fout.write(out_HTML)
fout.close()

# cleanup
os.system('rm ' + temp_anno_file)