#!/usr/bin/env python3
__author__ = ["Lavanya Rishishar", "Aroon Chande"]
__copyright__ = "Copyright 2019, Applied Bioinformatics Lab"
__license__ = "GPLv3"

import os
import pickle
import re
import shutil
import subprocess
import sys
from argparse import ArgumentParser, HelpFormatter


VERSION = 1.0

coordinates = {
	"1": {"cx": 128.6, "cy": 1.5, "ht": 1654.5, "wd": 118.6},
	"2": {"cx": 301.4, "cy": 43.6, "ht": 1612.4, "wd": 118.6},
	"3": {"cx": 477.6, "cy": 341.4, "ht": 1314.7, "wd": 118.6},
	"4": {"cx": 655.6, "cy": 517.9, "ht": 1138.1, "wd": 118.6},
	"5": {"cx": 835.4, "cy": 461, "ht": 1195.1, "wd": 118.6},
	"6": {"cx": 1012.4, "cy": 524.2, "ht": 1131.8, "wd": 118.6},
	"7": {"cx": 1198.2, "cy": 608.5, "ht": 1047.5, "wd": 118.6},
	"8": {"cx": 1372.9, "cy": 692.8, "ht": 963.2, "wd": 118.6},
	"9": {"cx": 1554.5, "cy": 724.4, "ht": 931.6, "wd": 118.6},
	"10": {"cx": 1733.8, "cy": 766.6, "ht": 889.4, "wd": 118.6},
	"11": {"cx": 1911.5, "cy": 766.6, "ht": 889.4, "wd": 118.6},
	"12": {"cx": 2095.6, "cy": 769.7, "ht": 886.3, "wd": 118.6},
	"13": {"cx": 129.3, "cy": 2068.8, "ht": 766.1, "wd": 118.6},
	"14": {"cx": 301.6, "cy": 2121.5, "ht": 713.4, "wd": 118.6},
	"15": {"cx": 477.5, "cy": 2153.1, "ht": 681.8, "wd": 118.6},
	"16": {"cx": 656.7, "cy": 2232.2, "ht": 602.8, "wd": 118.6},
	"17": {"cx": 841.2, "cy": 2290.7, "ht": 544.3, "wd": 118.6},
	"18": {"cx": 1015.7, "cy": 2313.9, "ht": 521.1, "wd": 118.6},
	"19": {"cx": 1199.5, "cy": 2437.2, "ht": 397.8, "wd": 118.6},
	"20": {"cx": 1374.4, "cy": 2416.1, "ht": 418.9, "wd": 118.6},
	"21": {"cx": 1553, "cy": 2510.9, "ht": 324.1, "wd": 118.6},
	"22": {"cx": 1736.9, "cy": 2489.8, "ht": 345.1, "wd": 118.6},
	"X": {"cx": 1911, "cy": 1799.6, "ht": 1035.4, "wd": 59},
}

chromSizes = {
	"1": 249250621, "2": 243199373, "3": 198022430, "4": 191154276,
	"5": 180915260, "6": 171115067, "7": 159138663, "8": 146364022,
	"9": 141213431, "10": 135534747, "11": 135006516, "12": 133851895,
	"13": 115169878, "14": 107349540, "15": 102531392, "16": 90354753,
	"17": 81195210, "18": 78077248, "19": 59128983, "20": 63025520,
	"21": 48129895, "22": 51304566, "X": 155270560, "Y": 59373566
}


def draw(opts):
	try:
		input = open(opts.input, 'r')
	except (IOError, EOFError) as e:
		print("Error opening input file!")
		raise(e)
	svg_fn = f"{opts.prefix}.svg"
	try:
		svg_fh = open(svg_fn, 'w')
		svg_fh.write(__head__)
	except (IOError, EOFError) as e:
		print("Error opening output file!")
		raise(e)
	lineNum = 1
	for entry in input:
		if entry.startswith("#"): continue
		entry = entry.rstrip().split("\t")
		if len(entry) != 7:
			print(f"Line number {lineNum} does not have 7 columns")
			os.exit()
		chrm, start, stop, feature, size, col, chrCopy = entry
		chrm = chrm.replace('chr', '')
		start = int(start); stop = int(stop); size = float(size); feature = int(feature); chrCopy = int(chrCopy)
		if 0 > size > 1:
			print(f"Feature size, {size},on line {lineNum} unclear. Please bound the size between 0 (0%) to 1 (100%). Defaulting to 1.")
			size = 1
		if not re.match("^#.{6}", col):
			print(f"Feature color, {col}, on line {lineNum} unclear.  Please define the color in hex starting with #.  Defaulting to #000000.")
			col = "#000000"
		if chrCopy not in [1,2]:
			print(f"Feature chromosome copy, {chrCopy}, on line {lineNum} unclear.  Skipping...")
			lineNum = lineNum + 1
			continue
		lineNum = lineNum + 1
		if feature == 0: # Rectangle
			ftStart = start*coordinates[chrm]["ht"]/chromSizes[chrm]
			ftEnd   = stop*coordinates[chrm]["ht"]/chromSizes[chrm]
			wd = coordinates[chrm]["wd"]*size/2
			if chrCopy == 1:
				x  = coordinates[chrm]["cx"] - wd
			else:
				x  = coordinates[chrm]["cx"]
			y  = coordinates[chrm]["cy"] + ftStart
			ht = ftEnd-ftStart
			svg_fh.write(f"<rect x=\"{x}\" y=\"{y}\" fill=\"{col}\" width=\"{wd}\" height=\"{ht}\"/>" + "\n")
		elif feature == 1: # Circle
			ftStart = start*coordinates[chrm]["ht"]/chromSizes[chrm]
			ftEnd   = stop*coordinates[chrm]["ht"]/chromSizes[chrm]
			r  = coordinates[chrm]["wd"]*size/4
			if chrCopy == 1:
				x  = coordinates[chrm]["cx"] - coordinates[chrm]["wd"]/4
			else:
				x  = coordinates[chrm]["cx"] + coordinates[chrm]["wd"]/4
			y  = coordinates[chrm]["cy"]+(ftStart+ftEnd)/2
			svg_fh.write(f"<circle fill=\"{col}\" cx=\"{x}\" cy=\"{y}\" r=\"{r}\"/>" + "\n")
		elif feature == 2: # Triangle
			ftStart = start*coordinates[chrm]["ht"]/chromSizes[chrm]
			ftEnd   = stop*coordinates[chrm]["ht"]/chromSizes[chrm]
			if chrCopy == 1:
				x  = coordinates[chrm]["cx"] - coordinates[chrm]["wd"]/2
				sx = 38.2*size
			else:
				x  = coordinates[chrm]["cx"] + coordinates[chrm]["wd"]/2
				sx = -38.2*size
			y  = coordinates[chrm]["cy"]+(ftStart+ftEnd)/2
			sy = 21.5*size
			svg_fh.write(f"<polygon fill=\"{col}\" points=\"{x-sx},{y-sy} {x},{y} {x-sx},{y+sy}\"/>" + "\n")
		elif feature == 3: # Line
			y1 = start*coordinates[chrm]["ht"]/chromSizes[chrm]
			y2 = stop*coordinates[chrm]["ht"]/chromSizes[chrm]
			y  = (y1+y2)/2
			y += coordinates[chrm]["cy"]
			miter = size*50
			if chrCopy == 1:
				x1 = coordinates[chrm]["cx"] - coordinates[chrm]["wd"]/2
				x2 = coordinates[chrm]["cx"]
				svg_fh.write(f"<line fill=\"none\" stroke=\"{col}\" stroke-miterlimit=\"10\" x1=\"{x1}\" y1=\"{y}\" x2=\"{x2}\" y2=\"{y}\"/>" + "\n")
			else:
				x1 = coordinates[chrm]["cx"]
				x2 = coordinates[chrm]["cx"] + coordinates[chrm]["wd"]/2
				svg_fh.write(f"<line fill=\"none\" stroke=\"{col}\" stroke-miterlimit=\"10\" x1=\"{x1}\" y1=\"{y}\" x2=\"{x2}\" y2=\"{y}\"/>" + "\n")
		else:
			print(f"Feature type, {feature}, unclear.  Please use either 0, 1, 2 or 3.  Skipping...")
			continue
	svg_fh.write(__tail__)
	if opts.verbose: print(f"\033[92mSuccessfully created SVG\033[0m")


if __name__ == "__main__":
	parser = ArgumentParser(prog="tagore",
							add_help=True,
							description='''
							tagore: a utility for illustrating human chromosomes
							https://github.com/jordanlab/tagore
							''',
							formatter_class=lambda prog: HelpFormatter(prog, width=120, max_help_position=120))

	parser.add_argument('--version', action='version',
						help='Print the software version.',
						version='tagore (version {})'.format(VERSION))

	# Input arguments
	parser.add_argument('-i', '--input', required=True, default=None, metavar='<input.bed>',
							help='Input BED-like file')
	parser.add_argument('-p', '--prefix', required=False, default="out", metavar='[output file prefix]',
							help='Output prefix [Default: "out"]')
	parser.add_argument('-v', '--verbose', required=False, default=False,
							help="Display verbose output",
                    		action="store_true")
	parser.add_argument('-f', '--force', required=False, default=False,
							help="Overwrite output files if they exist already",
                    		action="store_true")
	opts, unknown_args = parser.parse_known_args()
	if unknown_args:
		print("\033[93mOne or more unknown arguments were supplied:\033[0m {}\n".format(' '.join(unknown_args)))
		parser.print_help()
		sys.exit()
	if shutil.which("rsvg", mode=os.X_OK) is None:
		print(f"\033[91mCould not find `rsvg` in PATH.\033[0m")
		sys.exit()
	base_path = os.path.join(sys.prefix, 'lib', 'tagore-data', 'base.svg.p')
	try:
		base = open(base_path, 'rb')
	except (IOError, EOFError) as e:
		print("\033[91mCould not open base.svg.p, this file must be in the same folder as tagore!\033[0m")
		raise(e)
	__head__, __tail__ = pickle.load(base)
	if opts.verbose: print(f"\033[94mDrawing chromosome ideogram using {opts.input}\033[0m")
	if os.path.exists(f"{opts.prefix}.svg") and opts.force is False:
		print(f"\033[93m'{opts.prefix}.svg' already exists.\033[0m")
		ow = input(f"Overwrite {opts.prefix}.svg? [Y/n]:  ") or "y"
		if ow.lower() != "y":
			print(f"\033[93m'tagore will now exit...\033[0m")
			sys.exit()
		else:
			print(f"\033[94mOverwriting existing file and saving to: {opts.prefix}.svg\033[0m")
	else:
		if opts.verbose: print(f"\033[94mSaving to: {opts.prefix}.svg\033[0m")
	draw(opts)
	if opts.verbose: print(f"\033[94mConverting {opts.prefix}.svg -> {opts.prefix}.png\033[0m")
	try:
		subprocess.check_output(f"rsvg {opts.prefix}.svg {opts.prefix}.png", shell=True)
	except subprocess.CalledProcessError as e:
		if opts.verbose: print(f"\033[91mFailed SVG to PNG conversion...\033[0m")
		raise e
	finally:
		if opts.verbose: print(f"\033[92mSuccessfully converted SVG to PNG\033[0m")


