print("Suggested usage: python BTSpred.py -in INFILE")
print("make sure predisi.txt and predsl.txt files are ready")

import argparse
import os
import sys
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.Alphabet import IUPAC

#os.system('source ~/.bash_profile')

parser = argparse.ArgumentParser(description='How to use argparse')
parser.add_argument('-in', '--infile', help='fasta to be processed', default="chromera.fasta")

args = parser.parse_args()

filename = args.infile.split(".")[0]

def query_yes_no(question, default="yes"):
	"""Ask a yes/no question via raw_input() and return their answer.

	"question" is a string that is presented to the user.
	"default" is the presumed answer if the user just hits <Enter>.
		It must be "yes" (the default), "no" or None (meaning
		an answer is required of the user).

	The "answer" return value is True for "yes" or False for "no".
	"""
	valid = {"yes": True, "y": True, "ye": True,
			 "no": False, "n": False}
	if default is None:
		prompt = " [y/n] "
	elif default == "yes":
		prompt = " [Y/n] "
	elif default == "no":
		prompt = " [y/N] "
	else:
		raise ValueError("invalid default answer: '%s'" % default)

	while True:
		sys.stdout.write(question + prompt)
		if sys.platform in ["darwin", "win32"]:
			choice = input().lower()
		elif sys.platform.startswith("linux"):
			choice = raw_input().lower()
		else:
			quit("unrecognized OS, check how to use raw input")

		if default is not None and choice == '':
			return valid[default]
		elif choice in valid:
			return valid[choice]
		else:
			sys.stdout.write("Please respond with 'yes' or 'no' "
							 "(or 'y' or 'n').\n")

monstr_dic = {}
###########################################################################
#		read in seq names to extract nt/aa models of interest:		   #
###########################################################################
inproteins = SeqIO.parse(args.infile, 'fasta')
names = []
for sequence in inproteins:
	annotation = " ".join(sequence.description.split(" ")[1:])
	name = sequence.name
	if "-reversed" in name:
		name = name.replace("-reversed", "")
	if name not in names:
		names.append(name)
	seq = sequence.seq
	monstr_dic[name] = 	{"Annotation": annotation, "Name": name, "Seq": seq}

#print("Now printing control table, number of items per dictionary key...")
#for sequence in monstr_dic.keys():
#	print(sequence, len(monstr_dic[sequence]))

###########################################################################
#		  read in predictions from PrediSi, PredSL, TMHMM:			   #
###########################################################################
predisi = open("predisi.txt").readlines()
predsl = open("predsl.txt").readlines()

print("parsing PrediSi and PredSL predictions and sequences into data dictionary")
for line in predisi:
	line = line.split("\t")
	#FASTA-ID	Cleavage Position	Signal Peptide ?	Score
	#[0]		[1]					[2]					[3]
	if line[0] in names:
		name = line[0]
		if line[2] == "Y" and int(line[1]) < 101: #sometimes we see signal very far in the protein
			monstr_dic[name].update({"PrediSi site": line[1], "PrediSi SPscore": float(line[3])})
		else:
			monstr_dic[name].update({"PrediSi site": "-", "PrediSi SPscore": float(line[3])})

for line in predsl:
	line = line.split()
	#sequence id	mTP score	SP score	prediction	cleavage site
	#[0]			[1]			[2]			[3]			[4]
	if line[0] in names:
		name = line[0]
		#line[3] for nonplant prediction - complex algae
		if line[3] == "secreted":
			monstr_dic[name].update({"PredSL site": line[4], "PredSL SPscore": float(line[2])})
		else:
			monstr_dic[name].update({"PredSL site": "-", "PredSL SPscore": float(line[2])})

###########################################################################
#					  transit peptide detection:						 #
###########################################################################

print("preparing output for MultiLoc2")

with open('cleaved-' + args.infile, 'a') as result:
	for protein in monstr_dic.keys():
		#print(protein, monstr_dic[protein])
		try:
			print(protein)
			SP_si = monstr_dic[protein]["PrediSi site"]
			SP_sl = monstr_dic[protein]["PredSL site"]
			if SP_si == "-":
				if SP_sl != "-":
					SP_sl = int(SP_sl)
					SP_cleaved_sl = '>{}\n{}\n'.format(protein, monstr_dic[protein]["Seq"][SP_sl:])
					#SP_cleaved_sl = '>{}@PredSL\n{}\n'.format(protein, monstr_dic[protein]["Seq"][SP_sl:])
					result.write(SP_cleaved_sl)
			elif SP_sl == "-":
				SP_si = int(SP_si)
				SP_cleaved_si = '>{}\n{}\n'.format(protein, monstr_dic[protein]["Seq"][SP_si:])
				#SP_cleaved_si = '>{}@PrediSi\n{}\n'.format(protein, monstr_dic[protein]["Seq"][SP_si:])
				result.write(SP_cleaved_si)
			elif int(SP_si) == int(SP_sl):
				SP_si = int(SP_si)
				SP_sl = int(SP_sl)
				SP_cleaved = '>{}\n{}\n'.format(protein, monstr_dic[protein]["Seq"][SP_si:])
				result.write(SP_cleaved)
			else:
				SP_si = int(SP_si)
				SP_sl = int(SP_sl)
				SP_cleaved_si = '>{}\n{}\n'.format(protein, monstr_dic[protein]["Seq"][SP_si:])
				SP_cleaved_sl = '>{}@PredSL\n{}\n'.format(protein, monstr_dic[protein]["Seq"][SP_sl:])
				result.write(SP_cleaved_si)
				result.write(SP_cleaved_sl)
		except ValueError as VE:
			with open('result_errors.txt', 'a') as errors:
				errors.write('SP cleavage error: {}\n{}\n\n'.format(monstr_dic[protein][5], str(VE)))

print("cleaved proteins written to {}... now running MultiLoc2 plant option to determine transit peptide probabilities.".format("cleaved-" + args.infile))

if query_yes_no("Do you wish to perform MultiLoc2 on cleaved sequences?") is True:
	os.system('python /home/manager/MultiLoc2/src/multiloc2_prediction.py -fasta={} -predictor=LowRes -origin=plant -result=pred-ML2-cleave.txt -output=simple'.format("cleaved-" + args.infile))
#use the following if your script failed, but you have made the new MultiLoc2 predictions
#if os.path.isfile("pred-ML2-cleave.txt") == False:
#	os.system('python /home/manager/MultiLoc2/src/multiloc2_prediction.py -fasta={} -predictor=LowRes -origin=plant -result=pred-ML2-cleave.txt -output=simple'.format("cleaved-" + args.infile))

print("MultiLoc2 finished, now adding predictions to data dictionary")

with open("pred-ML2-cleave.txt") as infileML2:
	ML2C = infileML2.read().split('\n')

alternativecleavage = open("alternativecleavage-" + filename + ".txt", "w")
alternativecleavage_dict = {}

for line in ML2C:
	if len(line.split('\t')) != 1:
		line = line.split('\t')
		preds = sorted(line[1:])
		name = line[0]
		#['chloroplast: ', 'cytoplasmic: ', 'mitochondrial: ', 'nuclear: ', 'secretory pathway: ']
		if "@PredSL" in name or "@PrediSi" in name:
			alternativecleavage.write("alternative cleavage for: {}\n".format("\t".join(line)))
			name = name.split("@")[0]
			try:
				cTP = preds[0].split()[1]
				mTP = preds[2].split()[1]
				combinedMLpred = float(cTP) + float(mTP)
				if combinedMLpred > 0.3:
					ourguess = "BIPARTITE, plastid"
				else:
					ourguess = "secretory"
				#print(cTP, mTP)
				if name in alternativecleavage_dict:
					if combinedMLpred > alternativecleavage_dict[name]: #here we choose the "better prediction"
						monstr_dic[name].update({"ML2post-cleavage": "pt: {} (mt: {})".format(cTP, mTP), "ourguess": ourguess, "ourscore": combinedMLpred})
				else:
					alternativecleavage_dict[name] = combinedMLpred
					monstr_dic[name].update({"ML2post-cleavage": "pt: {} (mt: {})".format(cTP, mTP), "ourguess": ourguess, "ourscore": combinedMLpred})

			except ValueError as VE:
				with open('result_errors.txt', 'a') as errors:
					errors.write('ML2 processing error: {}\n{}\n\n'.format(line, str(VE)))

		else:
			try:
				cTP = preds[0].split()[1]
				mTP = preds[2].split()[1]
				combinedMLpred = float(cTP) + float(mTP)
				if combinedMLpred > 0.3:
					ourguess = "BIPARTITE, plastid"
				else:
					ourguess = "secretory"
				#print(cTP, mTP)
				monstr_dic[name].update({"ML2post-cleavage": "pt: {} (mt: {})".format(cTP, mTP), "ourguess": ourguess, "ourscore": combinedMLpred})
			except ValueError as VE:
				with open('result_errors.txt', 'a') as errors:
					errors.write('ML2 processing error: {}\n{}\n\n'.format(line, str(VE)))
print("Transit peptide predictions made, check alternativecleavage" + filename + ".txt for ambiguous predictions.")
alternativecleavage.close()

print("Now writing results...")
outtable = 'FINALTABLE-' + filename + '.txt'

#solution without pandas
with open('FINALTABLE-' + filename, 'w') as outtable, open('FINALPREDS-' + filename, 'w') as outpreds:
	for name in monstr_dic:
		item = monstr_dic[name]
		ML2pred = item.get("ML2post-cleavage", "-")
		ourplastidpred = item.get("ourguess", "non-plastid")
		ourplastidscore = "%.3f" % ((max(item["PrediSi SPscore"], item["PredSL SPscore"]) + 3*item.get("ourscore", 0))/4)
		outtable.write("{}\t{}\t{}\t{}\n".format(item["Name"],item["PrediSi site"],item["PredSL site"],ML2pred))
		outpreds.write("{}\t{}\t{}\n".format(item["Name"],ourplastidpred, ourplastidscore))
		print(item["Name"], item["PrediSi SPscore"], item["PredSL SPscore"], item.get("ourscore", 0))

print("Writing done. Hooray.")


from pycamia import info_manager

__info__ = info_manager(
    project = "PyCAMIA",
    package = "micomputing",
    author = "Yuncheng Zhou",
    create = "2021-12",
    fileinfo = "File to visualize features of 3D volume.",
    help = "Use `from micomputing import plot as plt`.",
    requires = "matplotlib"
).check()

__all__ = """
    set_font
    plot
    subplots
    background
    imshow
    imsshow
    volshow
    maskshow
    bordershow
    fieldshow
    gridshow
""".split()

import math
from .trans import DenseDisplacementField, Transformation
from .stdio import IMG
from .funcs import reorient

with __info__:
    import batorch as bt
    from matplotlib import pyplot as plt
    from matplotlib import colors as mc
    from matplotlib.pyplot import *
    from pyoverload import *
    from pycamia import prod, to_tuple, argmax, avouch, touch
    __all__.extend([x for x in dir(plt) if not x.startswith('_') and x not in __all__])
    plt.rcParams['font.sans-serif'] = 'Arial Unicode MS'
    plt.rcParams['axes.unicode_minus'] = False
    plt.rcParams['savefig.dpi'] = 300

canvas = None
colors = ['crimson', 'forestgreen', 'royalblue', 'red', 'green', 'blue', 'gold', 'purple', 'gray', 'pink', 'darkgreen', 'dodgerblue']
assert all(c in mc.CSS4_COLORS for c in colors)
subplot_rows = 1
subplot_cols = 1
subplot_id = 1
default_orient = None

TRS_list = ["trans", "trs", "transversal", "transverse", "ax", "axial",
            "axis", "crs", "cross", "hor", "horizontal", "slice"]
COR_list = ["cor", "coro", "coronal", "coron"]
SAG_list = ["sag", "sagit", "sagittal", "vert", "vertical"]
name_to_orient = {k: 'AR' for k in TRS_list}
name_to_orient.update({k: 'SR' for k in COR_list})
name_to_orient.update({k: 'SA' for k in SAG_list})

def set_font(font: str):
    plt.rcParams['font.sans-serif'] = font

def set_orient(ref):
    if isinstance(ref, IMG): ref = ref.orientation
    global default_orient
    default_orient = ref

def to_image(data: array, nslice: (int, null)=None, dim: int=-1, as_orient=None, to_orient='AR'):
    if as_orient is None: as_orient = default_orient
    if to_orient is None: to_orient = 'AR' # default orient for output
    as_orient = name_to_orient.get(as_orient.lower(), as_orient.upper()) if as_orient else None
    to_orient = name_to_orient.get(to_orient.lower(), to_orient.upper()) if to_orient else None
    axis_upper = {'L':'R', 'I':'S', 'P':'A'}
    to_orilist = [axis_upper.get(o, o) for o in to_orient]
    if as_orient and len(as_orient) == 3 and len(to_orilist) < 3:
        to_orilist += [x for x in "ARS" if x not in to_orilist]
        to_orient += to_orilist[-1]
    if isinstance(data, IMG) and as_orient is None:
        data = data.reorient(to_orient)
        if data.ndim == 3 and len(to_orilist) < 3:
            to_orilist += [x for x in "ARS" if x not in to_orilist]
            to_orient += to_orilist[-1]
        if isinstance(dim, int): # LPI => to_orient
            from_orilist = [axis_upper.get(o, o) for o in data.orientation]
            remapper = {k:to_orilist.index(from_orilist[k]) for k in range(-3, 3)}
            dim = remapper[dim]
    sp = getattr(data, 'spacing', None)
    data = bt.tensor(data).squeeze()
    if sp is not None:
        for d, s in enumerate(sp):
            data = data.amplify(round(s), d)
    if as_orient is not None: data = reorient(data, from_orient=as_orient, to_orient=to_orient)
    if isinstance(dim, str) and len(dim) == 2: dim = argmax([o in dim for o in to_orient])
    if data.ndim <= 1: raise TypeError("Please don't use 'plot.imshow' to demonstrate an array or a scalar. ")
    if data.nspace > 3: raise TypeError(f"'plot.imshow' takes 2 or 3D-data as input (currently {data.shape}), please reduce the dimension manually or specify special dimensions to reduce. ")
    if data.nspace == 3:
        if data.has_batch: data = data.pick(0, [])
        if data.has_channel:
            if dim < 0: dim += data.nspace
            dim += 1 if data.channel_dim <= dim else 0
            if nslice is None: nslice = data.size(dim) // 2
            data = data.pick(nslice, dim) \
                       .sample(number=min(data.channel_size, 3), random=False, dim={}) \
                       .movedim(data.channel_dimension, -1)
        else:
            if data.space[-1] <= 3: pass
            elif data.space[0] <= 3: data = data.movedim(0, 2)
            if nslice is None: nslice = data.size(dim) // 2
            data = data.pick(nslice, dim)
    elif data.nspace == 2:
        if data.has_batch: data = data.pick(0, [])
        if data.has_channel:
            data = data.sample(number=min(data.channel_size, 3), random=False, dim={}) \
                       .movedim(data.channel_dimension, -1)
    elif data.ndim == 3: data = data.pick(0, []).remove_special_()
    elif data.ndim == 2: data = data.remove_special_()
    else: TypeError(f"Cannot extrat image from data of size {data.shape}. ")
    if data.ndim == 3 and data.size(2) < 3: data = bt.cat(data, bt.zeros(data.shape[:-1] + (3 - data.size(2),)), 2)
    if data.ndim == 2 and data.max() == data.min(): data = data.multiply(3, -1)
    return data.float()

def to_RGB(*color):
    if len(color) == 0: return (1.,) * 3
    elif len(color) == 1:
        c = color[0]
        if isinstance(c, float) and 0 <= c <= 1: return (c,) * 3
        elif isinstance(c, (int, float)) and 0 <= c <= 255: return (c / 255,) * 3
        elif isinstance(c, tuple): return to_RGB(*c)
        elif isinstance(c, str):
            if not c.startswith('#'):
                c = mc.BASE_COLORS.get(c.lower(), mc.CSS4_COLORS.get(c.lower(), c))
                if isinstance(c, tuple): return to_RGB(*c)
            if not c.startswith('#'): raise TypeError(f"Unrecognized color {c}.")
            return mc.hex2color(c)
        else: raise TypeError("Unaccepted color type. ")
    elif len(color) == 3:
        if all(isinstance(c, float) and 0 <= c <= 1 for c in color): return color
        elif all(isinstance(c, (int, float)) and 0 <= c <= 255 for c in color): return tuple(c / 255 for c in color)
        else: raise TypeError("Unaccepted color type. ")
    else: raise TypeError("Unaccepted color type. ")

def subplots(n=1, m=None):
    if m is None:
        r = math.floor(math.sqrt(n))
        c = math.ceil(n / r)
        n, m = r, c
    global subplot_rows
    global subplot_cols
    global subplot_id
    subplot_rows = n
    subplot_cols = m
    subplot_id = 1

def clf():
    plt.clf()
    subplots()

def create_subplot():
    global subplot_rows
    global subplot_cols
    global subplot_id
    if subplot_id > subplot_rows * subplot_cols:
        raise TypeError(f"Cannot create subplot {subplot_id} in grid {subplot_rows} x {subplot_cols}, " +
                        "please use `micomputing.plot.subplots(rows, cols)` to specify the number of images. ")
    plt.subplot(subplot_rows, subplot_cols, subplot_id); subplot_id += 1

def plot(*data: (array, null), **kwargs):
    """
    An alias of matplotlib.pyplot.plot, compatible with batorch.
    [n_batch] for multiple lines, while {n_channel} for x and y.
    """
    create_subplot()
    return line(*data, **kwargs)

def line(*data: (array, null), **kwargs):
    """
    [n_batch] for multiple lines, while {n_channel} for x and y.
    """
    if len(data) == 1:
        y = bt.tensor(data[0])
        avouch(y.n_space == 1)
        if y.n_channel == 2: x, y = [u.squeeze() for u in y.split(1, {})]
        else:
            length = y.space[0]
            x = bt.arange(0, length).expand_to(y)
    elif len(data) == 2:
        x, y = data
    else: raise TypeError("Too much arguments for micomputing.plot/line. ")
    if isinstance(x, bt.torch.Tensor):
        if x.device.type == 'cuda': x = x.cpu()
        x = x.detach().cpu().numpy()
    if isinstance(y, bt.torch.Tensor):
        if y.device.type == 'cuda': y = y.cpu()
        y = y.detach().cpu().numpy()
    return plt.plot(x.T, y.T, **kwargs)

@typehint
def imshow(data: (array, null)=None, nslice: (int, null)=None, dim: int=-1, **kwargs):
    """
    An automatic image display function for all kinds of tensors.
    The first image in batched images will be selected to be showed.
    For medical images:
    Displacements with channel dimension identified will be displayed as RGB colored maps.
    If there are no dimension <=3, gray scaled images will be showed.
    Transverse medical image with the right hand side of the subject shown on the left
        and anterior shown at the bottom will be selected for 3D volumes.
    `nslice` and `dim` are used for 3D volumes only, meaning to show the `nslice` slice of dimension `dim`.
    """
    global canvas
    if kwargs.pop('create_subplot', data is not None): create_subplot()
    if data is not None:
        canvas = to_image(data, nslice, dim, kwargs.pop('as_orient', None), kwargs.pop('to_orient', 'ARS'))
    if canvas is None or isinstance(canvas, tuple):
        raise TypeError("Please input data in 'imshow' or 'background' to show. ")
    if isinstance(canvas, bt.torch.Tensor):
        if canvas.device.type == 'cuda': canvas_show = canvas.cpu()
        canvas_show = canvas.detach().cpu().numpy()
    else: canvas_show = canvas
    plt.axis('off')
    kwargs['cmap'] = kwargs.get('cmap', 'gray')
    if canvas_show.ndim >= 3:
        Max = canvas_show.max()
        if canvas_show.dtype.kind == 'i' and Max > 255: canvas_show = canvas_show * 255 // canvas_show.max(); plt.title(f"x {Max / 255:.04f}")
        elif Max > 1: canvas_show = canvas_show / canvas_show.max(); plt.title(f"x {Max:.04f}")
    return plt.imshow(canvas_show, **kwargs)

@typehint
def imsshow(*data: (array, null), nslice: (int, null)=None, dim: int=-1, show=False, **kwargs):
    """
    Display a series of images by 'imshow'. Use argument 'show=True' to directly call plt.show().
    The first image in batched images will be selected to be showed.
    For medical images:
    Displacements with channel dimension identified will be displayed as RGB colored maps.
    If there are no dimension <=3, gray scaled images will be showed.
    Transverse medical image with the right hand side of the subject shown on the left
        and anterior shown at the bottom will be selected for 3D volumes.
    `nslice` and `dim` are used for 3D volumes only, meaning to show the `nslice` slice of dimension `dim`.
    """
    return [imshow(d, nslice, dim, **kwargs) for d in data] + ([plt.show()] if show else [])

@typehint
def volshow(data: (array, null), show=False, **kwargs):
    subplots(2, 2)
    return [imshow(data, None, (d + 2) % 3, **kwargs) for d in range(len(data.shape))]

def background(*color):
    """
    Set a background color by RGB or a gray scale, conflict with imshow.
    """
    global canvas
    canvas = to_RGB(*color)
    create_subplot()

@typehint
def maskshow(*masks, on=None, alpha=0.5, nslice=None, dim=-1, stretch=False, **kwargs):
    global canvas
    orient_as = kwargs.pop('as_orient', None)
    orient_to = kwargs.pop('to_orient', None)
    if on is not None:
        if isinstance(on, (int, tuple)): background(*on)
        elif isarray(on): canvas = to_image(on, nslice, dim, orient_as, orient_to)
        elif isinstance(on, list): canvas = to_image(bt.tensor(on), nslice, dim, orient_as, orient_to)
        else: raise TypeError("Unrecognized argument 'on' for 'maskshow'. ")
    elif canvas is None:
        canvas = (1.,) * 3
    if len(masks) == 1 and not isarray(masks[0]):
        raise TypeError("Non-array mask for function 'maskshow', unzip the tuple if necessary. ")
    if len(masks) == 0: return imshow(create_subplot = on is not None)
    alpha = to_tuple(alpha, len(masks))
    new_masks = []
    new_alpha = []
    for m, a in zip(masks, alpha):
        img = to_image(m, nslice, dim, orient_as, orient_to)
        if img.ndim == 3:
            new_masks.extend(x.squeeze(-1) for x in img.split(1, dim=dim))
            new_alpha.extend([a] * img.size(dim))
        else:
            new_masks.append(img)
            new_alpha.append(a)
    color_mask_map = [(to_RGB(c), m, a) for c, m, a in zip(colors*(len(new_masks) // len(colors) + 1), new_masks, new_alpha)]
    for c, m in kwargs.items():
        if touch(lambda: to_RGB(c)) is not None:
            color_mask_map.append((to_RGB(c), m, alpha[0]))
            kwargs.pop(c)

    if not stretch:
        shapes = [m.ishape for _, m, _ in color_mask_map]
        target_shape = shapes[0]
        if len(set(shapes)) > 1 or not isinstance(canvas, tuple) and target_shape != canvas.shape:
            raise TypeError(f"Please use masks (size: {target_shape}) of the same size as the background image (size: {canvas.shape}), "
                            "or use 'stretch=True' in 'maskshow' to automatically adjust the image sizes. ")
    else:
        def adjust(m, to):
            ms = tuple(m.shape)
            scaling = tuple((a // b, b // a) for a, b in zip(to, ms))
            return m.down_scale([max(v, 1) for u, v in scaling]).up_scale([max(u, 1) for u, v in scaling]).crop_as(to)
        shapes = [m.ishape for _, m, _ in color_mask_map]
        if not isinstance(canvas, tuple): shapes.append(canvas.shape[:2])
        areas = [u * v for u, v in shapes]
        target_shape = shapes[areas.index(max(areas))]
        color_mask_map = [(c, adjust(m, to=target_shape), a) for c, m, a in color_mask_map]
        canvas = adjust(canvas, to=target_shape)

    target_shape = bt.Size(*target_shape, {3})
    if isinstance(canvas, tuple): canvas = bt.tensor(list(canvas)).expand_to(target_shape)
    elif canvas.ndim == 2: canvas = canvas.expand_to(target_shape)
    coeff = prod(1 - a * m for _, m, a in color_mask_map) # color_mask_map: (color, mask, alpha)
    canvas *= coeff
    for i, (c, m, a) in enumerate(color_mask_map):
        coeff = prod(a * m if j == i else 1 - a * m for j, (_, m, a) in enumerate(color_mask_map))
        canvas += coeff.unsqueeze(-1) * m.unsqueeze(-1) * bt.tensor(list(c)).unsqueeze(0, 1)
    if isinstance(canvas, bt.torch.Tensor):
        if canvas.device.type == 'cuda': canvas_show = canvas.cpu()
        canvas_show = canvas.detach().cpu().numpy()
    else: canvas_show = canvas
    plt.axis('off')
    create_subplot()
    kwargs['cmap'] = kwargs.get('cmap', 'gray')
    if canvas_show.ndim >= 3:
        Max = canvas_show.max()
        if canvas_show.dtype.kind == 'i' and Max > 255: canvas_show = canvas_show * 255 // canvas_show.max(); plt.title(f"x {Max / 255:.04f}")
        elif Max > 1: canvas_show = canvas_show / canvas_show.max(); plt.title(f"x {Max:.04f}")
    return plt.imshow(canvas_show, **kwargs)

def smooth(curve):
    """
    curve: 2xn
    """
    middle = (curve[:, :-2] + curve[:, 1:-1] + curve[:, 2:]) / 3
    if all(curve[:, 0] == curve[:, -1]):
        head = (curve[:, :1] + curve[:, 1:2] + curve[:, -2:-1]) / 3
        return bt.cat(head, middle, head, dim=1)
    return bt.cat(curve[:, :1], middle, curve[:, -1:], dim=1)

def sharpen(curve, old_curve):
    """
    sharpen towards old_curve
    curve, old_curve: 2xn
    """
    a = bt.cat(curve[:, -2:-1] - curve[:, -1:], curve[:, :-1] - curve[:, 1:], dim=1)
    b = bt.cat(curve[:, 1:] - curve[:, :-1], curve[:, 1:2] - curve[:, :1], dim=1)
    costheta = (a*b).sum(0) / bt.sqrt((a*a).sum(0)) / bt.sqrt((b*b).sum(0))
    new_curve = curve.clone()
    new_curve[:, costheta > -0.9] = old_curve[:, costheta > -0.9]
    return new_curve

def constraint(new_curve, old_curve, constraint_curve):
    """
    constraint new_curve towards constraint_curve
    curve, old_curve: 2xn
    """
    dis_sqs = bt.sum((new_curve - constraint_curve) ** 2, 0)
    percentile = 2 #min(np.sort(dis_sqs)[98 * len(dis_sqs) // 100], 1)
    return bt.where(dis_sqs <= percentile, new_curve, old_curve)

def border(mask, min_length = 10):
    grid = bt.image_grid(*mask.shape)
    mask = mask > 0.5
    idx = mask[1:, :] ^ mask[:-1, :]
    idx = idx.expand_to(2, -1, mask.size(1))
    locs1 = (grid[:, 1:, :] + grid[:, :-1, :])[idx] / 2
    idx = mask[:, 1:] ^ mask[:, :-1]
    idx = idx.expand_to(2, mask.size(0), -1)
    locs2 = (grid[:, :, 1:] + grid[:, :, :-1])[idx] / 2
    locs = bt.cat(locs1.reshape(2, -1), locs2.reshape(2, -1), dim=1)
    if locs.size == 0: return []
    curves = []
    unvisited = bt.ones(locs.shape[-1])
    while True:
        if not any(unvisited): break
        first = bt.argmax(unvisited).item()
        cloc = locs[:, first:first + 1]
        unvisited[first] = 0
        curve = cloc
        while True:
            dissq = bt.sum((locs - cloc) ** 2, 0)
            inloc = bt.argmax(bt.where((unvisited > 0) & (dissq > 0), 1/dissq.clamp(min=0.1), bt.tensor(0).float()))
            if dissq[inloc] > 2: break
            cloc = locs[:, inloc:inloc + 1]
            curve = bt.cat(curve, cloc, dim=1)
            unvisited[inloc] = 0
            if not any(unvisited): break
        sloc = locs[:, first:first + 1]
        if bt.sum((cloc - sloc) ** 2) <= 2:
            curve = bt.cat(curve, sloc, dim=1)
        if curve.shape[1] <= min_length: continue
        scurve = curve
        for _ in range(100): scurve = constraint(smooth(scurve), scurve, curve)
        ccurve = scurve
        for _ in range(100):
            scurve = constraint(sharpen(scurve, curve), scurve, ccurve)
            scurve = constraint(smooth(scurve), scurve, curve)
        curves.append(scurve)
    return curves

def bordershow(*masks, on=None, mask_alpha=0., nslice=None, dim=-1, stretch=False, min_length = 10, **kwargs):
    global canvas
    orient_as = kwargs.pop('as_orient', None)
    orient_to = kwargs.pop('to_orient', None)
    if on is not None:
        if isinstance(on, (int, tuple)): background(*on)
        elif isarray(on): canvas = to_image(on, nslice, dim, orient_as, orient_to)
        elif isinstance(on, list): canvas = to_image(bt.tensor(on), nslice, dim, orient_as, orient_to)
        else: raise TypeError("Unrecognized argument 'on' for 'maskshow'. ")
    elif canvas is None:
        canvas = (1.,) * 3
    if len(masks) == 0: return
    new_masks = []
    for m in masks:
        img = to_image(m, nslice, dim, orient_as, orient_to)
        if img.ndim == 3:
            new_masks.extend(x.squeeze(-1) for x in img.split(1, dim=dim))
        else:
            new_masks.append(img)
    color_mask_map = [(to_RGB(c), m) for c, m in zip(colors*(len(new_masks) // len(colors) + 1), new_masks)]
    for c, m in kwargs.items():
        if touch(lambda: to_RGB(c)) is not None:
            color_mask_map.append((to_RGB(c), m))
            kwargs.pop(c)

    new_masks = [m for _, m in color_mask_map]
    shapes = [m.ishape for _, m in color_mask_map]
    if not stretch:
        target_shape = shapes[0]
        if len(set(shapes)) > 1 or not isinstance(canvas, tuple) and target_shape != canvas.shape[:2]:
            raise TypeError("Please use masks of the same size as the background image, "
                            "or use 'stretch=True' in 'maskshow' to automatically adjust the image sizes. ")
    else:
        def adjust(m, to):
            ms = tuple(m.shape)
            scaling = tuple((a // b, b // a) for a, b in zip(to, ms))
            return m.down_scale([max(v, 1) for u, v in scaling]).up_scale([max(u, 1) for u, v in scaling]).crop_as(to)
        if not isinstance(canvas, tuple): shapes.append(canvas.shape[:2])
        areas = [u * v for u, v in shapes]
        target_shape = shapes[areas.index(max(areas))]
        color_mask_map = [(c, adjust(m, to=target_shape)) for c, m in color_mask_map]
        canvas = adjust(canvas, to=target_shape)
    if isinstance(canvas, tuple):
        canvas = bt.channel_tensor(canvas).expand_to(target_shape + bt.Size({3},)).with_channeldim(None)

    plots = []
    if on is not None: plots.append(imshow(create_subplot=True))
    if mask_alpha > 0: plots.append(maskshow(*new_masks, alpha=mask_alpha))
    else: plots.append(imshow())
    for color, mask in color_mask_map:
        curves = border(mask, min_length)
        for c in curves:
            if isinstance(c, bt.torch.Tensor):
                if c.device.type == 'cuda': c = c.cpu()
                c = c.detach().cpu().numpy()
            plots.append(plt.plot(c[1], c[0], color = color, **kwargs))
    return plots

def fieldshow(disp, on=None, nslice=None, dim=-1, stretch=False, **kwargs):
    """
    fieldshow shows displacements with RGB channels to form a facke colored map
        indicating different directions by different color and magnitude by brightness.
    """
    global canvas
    orient_as = kwargs.pop('as_orient', None)
    orient_to = kwargs.pop('to_orient', None)
    if on is not None:
        if isinstance(on, (int, tuple)): background(*on)
        elif isarray(on): canvas = to_image(on, nslice, dim, orient_as, orient_to)
        elif isinstance(on, list): canvas = to_image(bt.tensor(on), nslice, dim, orient_as, orient_to)
        else: raise TypeError("Unrecognized argument 'on' for 'maskshow'. ")
    elif canvas is None:
        canvas = (1.,) * 3

    if isinstance(disp, Transformation):
        if canvas is None or isinstance(canvas, tuple) or not hasattr(on, 'ndim') or on.ndim != disp.n_dim:
            raise TypeError("Please input data in 'imshow' or 'background' before trying to display grid. ")
        disp = disp.toDDF(on.space).pick(0, [])
    elif canvas is None or isinstance(canvas, tuple): canvas = bt.zeros(disp.space)
    disp = disp.abs()
    if orient_as is not None: disp = reorient(disp, from_orient=orient_as, to_orient=orient_to)
    disp = to_image(disp, nslice, dim)
    if disp.size(2) > 3: disp = disp[..., :3]
    return imshow(disp, create_subplot = on is not None, **kwargs)

def gridshow(disp, on=None, gap=None, nslice=None, dim=-1, stretch=False, **kwargs):
    """
    gridshow shows grid for the displacements.
    """
    global canvas
    orient_as = kwargs.pop('as_orient', None)
    orient_to = kwargs.pop('to_orient', None)
    if on is not None:
        if isinstance(on, (int, tuple)): background(*on)
        elif isarray(on): canvas = to_image(on, nslice, dim, orient_as, orient_to)
        elif isinstance(on, list): canvas = to_image(bt.tensor(on), nslice, dim, orient_as, orient_to)
        else: raise TypeError("Unrecognized argument 'on' for 'maskshow'. ")
    elif canvas is None:
        canvas = (1.,) * 3

    if isinstance(disp, Transformation):
        if canvas is None or isinstance(canvas, tuple) or not hasattr(on, 'ndim') or on.n_space != disp.n_dim:
            raise TypeError("Please input data for displacement in 'imshow' or 'background' before trying to display grid: need size of the canvas. ")
        disp = disp.toDDF(bt.tensor(on).space)
    elif canvas is None or isinstance(canvas, tuple): canvas = bt.zeros(disp.space)
    if disp.has_batch: disp = disp.sample(random=False, dim=[])
    if orient_as is not None: disp = reorient(disp, from_orient=orient_as, to_orient=orient_to)
    Y = DenseDisplacementField(disp)(bt.image_grid(disp.space)).squeeze([])
    Y = to_image(Y, nslice, dim).movedim(-1, 0)
    if Y.size(0) > 2:
        if dim < 0: dim += Y.size(0)
        Y = bt.cat(Y[:dim], Y[dim+1:], {})

    plots = [imshow(canvas, create_subplot = on is not None)]
    plt.ylim((Y.space[0], 0))
    plt.xlim((0, Y.space[1]))
    if gap is None: gap = Y.space[0] // 8, Y.space[1] // 8
    gap = to_tuple(gap)
    if len(gap) == 1: gap = gap * 2
    if 'color' not in kwargs: kwargs['color'] = 'gold'
    if isinstance(Y, bt.torch.Tensor):
        if Y.device.type == 'cuda': Y = Y.cpu()
        Y_data = Y.detach().cpu().numpy()
    for row in list(bt.arange(0, Y.space[0], gap[0])) + [Y.space[0] - 1]:
        plots.append(plt.plot(Y_data[1, row, :], Y_data[0, row, :], **kwargs))
    for col in list(bt.arange(0, Y.space[1], gap[0])) + [Y.space[1] - 1]:
        plots.append(plt.plot(Y_data[1, :, col], Y_data[0, :, col], **kwargs))
    return plots


import time
import config
import threading
import datetime
import queue
import json
import os

import board as brd
from sgf import make_sgf, parse_sgf
from client import ClientSocketError
from utils import check_and_mkdir, get_html_code

def color_to_char(c):
    if c == brd.BLACK:
        return 'b'
    return 'w'

def get_opp_color(c):
    if c == brd.BLACK:
        return brd.WHITE
    return brd.BLACK

def move_to_vertex(board, move, support_analysis):
    move = move.strip()
    analysis = None
    if support_analysis:
        # Parse and validate analyze info.
        tokens = move.split(None, 1)
        move = tokens[0]
        if len(tokens) > 1:
            try:
                info = json.loads(tokens[1])
                analysis = json.dumps(info, indent=None, separators=(",", ":"))
            except:
                pass
    vertex = brd.NULL_VERTEX
    move = move.lower()

    if move == "pass":
        vertex = brd.PASS
    elif move == "resign":
        vertex = brd.RESIGN
    else:
        x = ord(move[0]) - (ord('A') if ord(move[0]) < ord('a') else ord('a'))
        y = int(move[1:]) - 1
        if x >= 8:
            x -= 1
        vertex = board.get_vertex(x,y)
    return move, vertex, analysis

def write_sgf_and_html(
    setting,
    names,
    date,
    move_history,
    result,
    base_name
):
    black_name, white_name = names
    sgf_store_path = os.path.join(
        *config.DATA_DIR_ROOT, "sgf", setting["store"])

    if os.path.isdir(sgf_store_path):
        # Save the SGF file.
        sgf_name = "{}.sgf".format(base_name)
        sgf = make_sgf(
                  setting["board_size"],
                  setting["komi"],
                  black_name,
                  white_name,
                  setting["main_time"],
                  date,
                  move_history,
                  result
              )
        sgf_full_name = os.path.join(sgf_store_path, sgf_name)
        with open(sgf_full_name, 'w') as f:
            f.write(sgf)

        html_store_path = os.path.join(
            *config.DATA_DIR_ROOT, "html", setting["store"])
        if os.path.isdir(html_store_path) and config.WGO_PATH is not None:
            html_name = "{}.html".format(base_name)

            back_count = 0
            for v in setting["store"].split(os.sep):
                # Assume there is no ".." symbol.
                if v != ".":
                    back_count += 1

            back_path = "."
            for _ in range(back_count+1):
                # Back to data directory root.
                back_path = os.path.join(back_path, "..")

            # Rewrite the WGO and SGF path in the HTML file.
            sgf_name_in_html = os.path.join(
                back_path, "sgf", setting["store"], sgf_name)
            wgo_path_in_html = os.path.join(back_path, config.WGO_PATH)

            html_full_name = os.path.join(html_store_path, html_name)
            if not os.path.isfile(html_full_name):
                # Only wrtie the HTML file once.
                with open(html_full_name, 'w') as f:
                    f.write(get_html_code(wgo_path_in_html, sgf_name_in_html))
        return True
    return False

def play_match_game(game_id, black, white, setting):
    # Play a match game and save the SGF file. The client may
    # crash here. We detect it and guarantee that the client can
    # return back safely. The socket is not closed here when
    # crashing. The master will close it later.

    # Initialize some basic data.
    time_lefts = {
        brd.BLACK : setting["main_time"],
        brd.WHITE : setting["main_time"]
    }
    players = {
        brd.BLACK : black,
        brd.WHITE : white
    }

    # TODO: Add support for Chinese rule and Japanese rule.
    rule = setting["rule"]
    should_superko = rule == "chinese-like"

    # We only record the starting time in order to fix
    # the output file name.
    date = datetime.datetime.now().strftime("%Y-%m-%d-%H:%M:%S")

    #
    sgf_clock_time = time.time()

    # The store path and SGF name.
    base_name = "{}-{}(B)-{}(W)-g{}".format(date, black.name, white.name, game_id)

    move_history = list() # It contains (move, time_left and analysis).

    # Try to read the SGF file. Should start the game from
    # it if the source is not None.
    sgf_source = setting.get("sgf", None)
    if (sgf_source is not None) and (os.path.isfile(sgf_source)):
        with open(sgf_source, 'r') as f:
            sgf = f.read().strip()
            board_size, komi, move_history = parse_sgf(sgf)
            # Rewrite the game setting.
            setting["board_size"] = board_size
            setting["komi"] = komi

    board = brd.Board(setting["board_size"], setting["komi"])
    result_status = dict()

    try:
        for player in players.values():
            # Request each clients to initialize the game also
            # create new socket file.
            player.create_sockfile()
            player.request_setup(
                setting["board_size"],
                setting["komi"],
                setting["main_time"] * 1000,
                players[brd.WHITE].name,
                players[brd.BLACK].name
            )

        while len(move_history) > board.move_num:
            # Play the moves from SGF file.
            side_to_move = board.to_move
            to_move_player = players[side_to_move]
            move, time_left, _ = move_history[board.move_num]

            move, vertex, _ = move_to_vertex(
                                  board, move, False
                              )

            # Always assuem the move is legel.
            board.play(vertex)

            # Both clients should play the move.
            for player in players.values():
                player.request_play(
                    color_to_char(side_to_move),
                    board.vertex_to_text(vertex),
                    int(time_left * 1000)
                )

        while True:
            side_to_move = board.to_move
            opp_to_move = get_opp_color(board.to_move)
            to_move_player = players[side_to_move]

            # Request the engine to genmove the next move.
            clock_time = time.time()
            time_left = time_lefts[side_to_move]
            rep = to_move_player.request_genmove(
                      color_to_char(side_to_move),
                      int(time_left * 1000)
                  )
            time_left -= (time.time() - clock_time)

            if time_left < 0:
                # Game ended by time out.
                winner = opp_to_move
                result_status["winner"] = winner
                result_status["type"] = "time out"
                result_status["info"] = "{}+Time".format(
                                            color_to_char(winner).upper()
                                        )
                break

            # Parse the move and analysis string.
            time_lefts[side_to_move] = time_left
            move, vertex, analysis = move_to_vertex(
                                         board,
                                         rep,
                                         to_move_player.support_analysis
                                     )

            if vertex == brd.RESIGN:
                # Game ended by resign move.
                winner = opp_to_move
                result_status["winner"] = winner
                result_status["type"] = "resign"
                result_status["info"] = "{}+Resign".format(
                                            color_to_char(winner).upper()
                                        )
                break

            is_legal = board.play(vertex)

            if not is_legal or \
                (vertex != brd.PASS and \
                     not should_superko and board.superko()):
                # Game ended by illegal move.
                winner = opp_to_move
                result_status["winner"] = winner
                result_status["type"] = "illegal move"
                result_status["info"] = "{}+Illegal".format(
                                            color_to_char(winner).upper()
                                        )
                break

            move_history.append((move, int(time_left), analysis))

            # Try to save game result into SGF file after updating
            # the move_history. Failed to save it if the client play
            # the move too quick.
            if time.time() - sgf_clock_time > 5:
                if write_sgf_and_html(
                       setting,
                       (black.name, white.name),
                       date,
                       move_history,
                       None,
                       base_name):
                    sgf_clock_time = time.time()

            # Request the opponent to play the move.
            time_left = time_lefts[opp_to_move]
            opp_player = players[opp_to_move]
            opp_player.request_play(
                color_to_char(side_to_move),
                board.vertex_to_text(vertex),
                int(time_left * 1000)
            )

            if board.num_passes >= 2:
                # Game ended by double pass. Now score
                # the final position.

                if rule == "chinese-like":
                    result_status["type"] = "double pass"
                    black_score = board.final_score()
                elif rule == "null":
                    result_status["type"] = "no rule"
                    black_score = 0
                else:
                    # Invalid rules.
                    result_status["type"] = "invalid rule"
                    black_score = 0

                if black_score > 0.001:
                    winner = brd.BLACK
                elif black_score < -0.001:
                    winner = brd.WHITE
                else:
                    winner = brd.EMPTY
                result_status["winner"] = winner

                if winner == brd.EMPTY:
                    result_status["info"] = "0"
                else:
                    result_status["info"] = "{}+{}".format(
                                                color_to_char(winner).upper(),
                                                abs(black_score)
                                            )
                break
    except:
        # TODO: Catch the error and write it into the SGF file.
        result_status["winner"] = brd.EMPTY
        result_status["type"] = "socket error"
        result_status["info"] = "0"

    result = result_status["info"]
    err = str()

    for player in players.values():
        try:
           # Send the last request to server here in order
           # to check whether the socket is still connected.
           player.request_poll()
        except ClientSocketError as e:
            pass

    # Close the socket file because the we can not push socket file
    # onto process queue.
    for player in players.values():
        try:
           player.close_sockfile()
        except ClientSocketError as e:
            pass

    # Always save the SGF file before leaving.
    write_sgf_and_html(
        setting,
        (black.name, white.name),
        date,
        move_history,
        result,
        base_name)

def match_loop(process_id, ready_queue, finished_queue):
    match_threads = dict()

    while True:
        finished_ids = list()
        for k, v in match_threads.items():
            # Collect all finished match games.
            t, _, _, _ = v
            if not t.is_alive():
                finished_ids.append(k)

        for i in finished_ids:
            # The match game is game over. Push the play
            # back to main pooling.
            v = match_threads.pop(i)
            t, i, b, w = v
            t.join()

            task = {
                "black" : b,
                "white" : w,
                "gid"   : i,
                "pid"   : process_id
            }
            finished_queue.put(task)

        try:
            task = ready_queue.get(block=True, timeout=0.1)
            if task["pid"] != process_id:
                # Not correct process id. Reture it to finished
                # queue.
                finished_queue.put(task)
            black = task["black"] # black player
            white = task["white"] # white player
            game_id = task["gid"] # game id
        except queue.Empty:
            continue

        # TODO: Add support for more task type.

        setting = {
            "main_time"  : task.get("main_time", config.DEFAULT_MAIN_SECOND),
            "board_size" : task.get("board_size", config.DEFAULT_BOARD_SIZE),
            "komi"       : task.get("komi", config.DEFAULT_KOMI),
            "sgf"        : task.get("sgf", None),
            "store"      : task.get("store", config.DEFAULT_STORE_DIR),
            "rule"       : task.get("rule", "chinese-like")
        }

        # New game is starting. Each threads hold one game. The threads
        # will be released after the gameover.
        t = threading.Thread(
                target=play_match_game,
                args=(game_id, black, white, setting, ),
                daemon=True
            )
        t.start()
        match_threads[t.ident] = (t, game_id, black, white)


import matplotlib

matplotlib.use('TkAgg')

import numpy as np
import cv2

def toBinary(img):
    '''
    Transform an image so every pixel that is not black becomes white
    :param img: An image
    :return: A copy of the original image, with all the pixel that are not black have 1.0
    '''
    gray = img[:, :, 2] + img[:, :, 1] + img[:, :, 0]
    binary = np.zeros(gray.shape, dtype=np.float32)
    binary[gray > 0] = 1.0
    return binary

def fit_polynomial(warped, left, right):
    '''
    :param binary_warped: a binary image tha represent a bird-eye view of the street
    :return: the lane lines(left and right) that have been detected,
    and an image that can be visualized for debugging purposes
    '''
    binary_warped = toBinary(warped)

    if (not left.detected):
        leftx, lefty, rightx, righty, out_img = find_lane_pixels(binary_warped)

    else:
        leftx, lefty, rightx, righty, out_img = find_lane_pixels_from_prior(binary_warped, left, right)

    ploty = np.linspace(0, warped.shape[0] - 1, warped.shape[0])

    # Fit a second order polynomial to each using `np.polyfit`
    try:
        left_fit = (np.polyfit(lefty, leftx, 2))
        left.update_fitted(left_fit, ploty)
    except Exception:
        left.detected = False

    try:
        right_fit = (np.polyfit(righty, rightx, 2))
        right.update_fitted(right_fit, ploty)
    except Exception:
        right.detected = False

    ## Visualization ##
    # Colors in the left and right lane regions
    out_img[lefty, leftx] = [255, 0, 0]
    out_img[righty, rightx] = [0, 0, 255]

    return left, right, out_img

def find_lane_pixels(binary_warped):
    '''
    :param binary_warped: a binary image tha represent a bird-eye view of the street
    :return: the pixels that are likely to belong to lane lines
    '''
    # Take a histogram of the bottom half of the image
    histogram = np.sum(binary_warped[binary_warped.shape[0] // 2:, :], axis=0)
    # Create an output image to draw on and visualize the result
    out_img = np.dstack((binary_warped, binary_warped, binary_warped))
    # Find the peak of the left and right halves of the histogram
    # These will be the starting point for the left and right lines
    midpoint = np.int(histogram.shape[0] // 2)
    leftx_base = np.argmax(histogram[:midpoint])
    rightx_base = np.argmax(histogram[midpoint:]) + midpoint

    # HYPERPARAMETERS
    # Choose the number of sliding windows
    nwindows = 9
    # Set the width of the windows +/- margin
    margin = 75
    # Set minimum number of pixels found to recenter window
    minpix = 50

    # Set height of windows - based on nwindows above and image shape
    window_height = np.int(binary_warped.shape[0] // nwindows)
    # Identify the x and y positions of all nonzero pixels in the image
    nonzero = binary_warped.nonzero()
    nonzeroy = np.array(nonzero[0])
    nonzerox = np.array(nonzero[1])
    # Current positions to be updated later for each window in nwindows
    leftx_current = leftx_base
    rightx_current = rightx_base

    # Create empty lists to receive left and right lane pixel indices
    left_lane_inds = []
    right_lane_inds = []

    # when no pixels are found in a window, add an artificial point on the center of the window,
    # that will be used for fitting the polynom
    artificial_leftx = []
    artificial_lefty = []
    artificial_rightx = []
    artificial_righty = []

    # Step through the windows one by one
    for window in range(nwindows):
        # Identify window boundaries in x and y (and right and left)
        win_y_low = binary_warped.shape[0] - (window + 1) * window_height
        win_y_high = binary_warped.shape[0] - window * window_height
        win_xleft_low = leftx_current - margin
        win_xleft_high = leftx_current + margin
        win_xright_low = rightx_current - margin
        win_xright_high = rightx_current + margin

        # Draw the windows on the visualization image

        try:
            cv2.rectangle(out_img, (win_xleft_low, win_y_low),
                          (win_xleft_high, win_y_high), (0, 255, 0), 2)
            cv2.rectangle(out_img, (win_xright_low, win_y_low),
                          (win_xright_high, win_y_high), (0, 255, 0), 2)
        except Exception:
            None

        # Identify the nonzero pixels in x and y within the window #
        good_left_inds = ((nonzeroy >= win_y_low) & (nonzeroy < win_y_high) &
                          (nonzerox >= win_xleft_low) & (nonzerox < win_xleft_high)).nonzero()[0]
        good_right_inds = ((nonzeroy >= win_y_low) & (nonzeroy < win_y_high) &
                           (nonzerox >= win_xright_low) & (nonzerox < win_xright_high)).nonzero()[0]

        # Append these indices to the lists
        left_lane_inds.append(good_left_inds)
        right_lane_inds.append(good_right_inds)

        # If you found > minpix pixels, recenter next window on their mean position
        if len(good_left_inds) > minpix:
            leftx_current = np.int(np.mean(nonzerox[good_left_inds]))
        else:
            # otherwise fit a point located on the center of the window
            artificial_leftx.append(int(leftx_current))
            artificial_lefty.append(int((win_y_high + win_y_low) / 2))

        if len(good_right_inds) > minpix:
            rightx_current = np.int(np.mean(nonzerox[good_right_inds]))
        else:
            artificial_rightx.append(int(rightx_current))
            artificial_righty.append(int((win_y_high + win_y_low) / 2))

    # Concatenate the arrays of indices (previously was a list of lists of pixels)
    try:
        left_lane_inds = np.concatenate(left_lane_inds)
        right_lane_inds = np.concatenate(right_lane_inds)
    except ValueError:
        # Avoids an error if the above is not implemented fully
        pass

    # Extract left and right line pixel positions
    leftx = nonzerox[left_lane_inds]
    lefty = nonzeroy[left_lane_inds]
    rightx = nonzerox[right_lane_inds]
    righty = nonzeroy[right_lane_inds]

    leftx = np.append(leftx, np.array(artificial_leftx, dtype=np.int32))
    lefty = np.append(lefty, np.array(artificial_lefty, dtype=np.int32))
    rightx = np.append(rightx, np.array(artificial_rightx, dtype=np.int32))
    righty = np.append(righty, np.array(artificial_righty, dtype=np.int32))

    return leftx, lefty, rightx, righty, out_img

def find_lane_pixels_from_prior(binary_warped, left, right):
    margin = 100

    mask_left = np.zeros_like(binary_warped)
    mask_right = np.zeros_like(binary_warped)

    for y in left.ploty:
        mask_left[int(y), int(left.best_plotx[int(y)]) - margin:int(left.best_plotx[int(y)]) + margin] = 1.0
        mask_right[int(y), int(right.best_plotx[int(y)]) - margin:int(right.best_plotx[int(y)]) + margin] = 1.0

    pixels_left = binary_warped * mask_left
    nonzero_left = pixels_left.nonzero()
    nonzerox_left = np.array(nonzero_left[1])
    nonzeroy_left = np.array(nonzero_left[0])

    pixels_right = binary_warped * mask_right
    nonzero_right = pixels_right.nonzero()
    nonzerox_right = np.array(nonzero_right[1])
    nonzeroy_right = np.array(nonzero_right[0])

    masks = (mask_left + mask_right) * 255
    out_img = np.dstack((masks, masks, masks))

    return nonzerox_left, nonzeroy_left, nonzerox_right, nonzeroy_right, out_img

# convolution
def find_window_centroids(image, window_width, window_height, margin):
    window_centroids = []  # Store the (left,right) window centroid positions per level
    window = np.ones(window_width)  # Create our window template that we will use for convolutions

    # First find the two starting positions for the left and right lane by using np.sum to get the vertical image slice
    # and then np.convolve the vertical image slice with the window template

    # Sum quarter bottom of image to get slice, could use a different ratio
    l_sum = np.sum(image[int(3 * image.shape[0] / 4):, :int(image.shape[1] / 2)], axis=0)
    l_center = np.argmax(np.convolve(window, l_sum)) - window_width / 2
    r_sum = np.sum(image[int(3 * image.shape[0] / 4):, int(image.shape[1] / 2):], axis=0)
    r_center = np.argmax(np.convolve(window, r_sum)) - window_width / 2 + int(image.shape[1] / 2)

    # Add what we found for the first layer
    window_centroids.append((l_center, r_center))

    # Go through each layer looking for max pixel locations
    for level in range(1, (int)(image.shape[0] / window_height)):
        # convolve the window into the vertical slice of the image
        image_layer = np.sum(
            image[int(image.shape[0] - (level + 1) * window_height):int(image.shape[0] - level * window_height), :],
            axis=0)
        conv_signal = np.convolve(window, image_layer)
        # Find the best left centroid by using past left center as a reference
        # Use window_width/2 as offset because convolution signal reference is at right side of window, not center of window
        offset = window_width / 2
        l_min_index = int(max(l_center + offset - margin, 0))
        l_max_index = int(min(l_center + offset + margin, image.shape[1]))
        l_center = np.argmax(conv_signal[l_min_index:l_max_index]) + l_min_index - offset
        # Find the best right centroid by using past right center as a reference
        r_min_index = int(max(r_center + offset - margin, 0))
        r_max_index = int(min(r_center + offset + margin, image.shape[1]))
        r_center = np.argmax(conv_signal[r_min_index:r_max_index]) + r_min_index - offset
        # Add what we found for that layer
        window_centroids.append((l_center, r_center))

    return window_centroids


from django.urls import path

from . import views

urlpatterns = [
    path('', views.index, name='index'),
    path('details_chapter/<int:id_capitulo>/', views.details_chapter, name='details_chapter'),
    path('details_character/<int:id_character>/', views.details_character, name='details_character'),
    path('details_location/<int:id_location>/', views.details_location, name='details_location'),
    path('search/', views.search, name='search'),

]


import numpy as np
import csv
import matplotlib.pyplot as plt
import os
from PIL import Image
from numpy import asarray
import copy
import math
#mwuahhahahahha it works! :)))

#l2 regularization
#dropout
#mini-batch
def main():
    showcost = 200
    #hyperparameters loading data:
    maxDataTest = 100
    imageSize = 120
    maxData = 600

    #hyperparameters network
    mini_batch_size = 200
    keepProb = 1
    lambd = 0.05 #between 0 - 0.1 the bigger lambda the bigger penalty on complicated weights
    layerDims = [imageSize * imageSize * 3, 400, 300, 200, 10,50,30,1]
    iterations = 2000
    learningRate = 0.001

    Xtrain, Ytrain = loadImages("train", imageSize, maxData)
    Xtrain, Ytrain = hussle(Xtrain, Ytrain) #for final training error
    mini_batches = random_mini_batches(Xtrain,Ytrain, mini_batch_size)
    Xtest, Ytest = loadImages("test", imageSize, maxDataTest)
    Xtest, Ytest = hussle(Xtest, Ytest)

    parameters = trainmodel(layerDims, mini_batches, iterations, learningRate, showcost, lambd, Xtest, Ytest, keepProb)
    trainingPrediction = makePrediction(Xtrain, parameters) #no dropout for test time
    testPrediction = makePrediction(Xtest, parameters) #no dropout for test time
    print("trainingAccuracy", accuracy(trainingPrediction, Ytrain))
    print("testAccuracy", accuracy(testPrediction, Ytest))

def random_mini_batches(X,Y, mini_batch_size):
    m = X.shape[1]
    mini_batches = []
    shuffled_X, shuffled_Y = hussle(X,Y)
    inc = mini_batch_size
    num_complete_minibatches = math.floor(m/mini_batch_size)
    for k in range(0, num_complete_minibatches):
        mini_batch_X = shuffled_X[:, k*inc:(k+1) * inc]
        mini_batch_Y = shuffled_Y[:, k*inc:(k+1) * inc]
        mini_batch = (mini_batch_X, mini_batch_Y)
        mini_batches.append(mini_batch)
    #handling non-complete mini_batch at the end:
    if m % mini_batch_size != 0:
        mini_batch_X = shuffled_X[:, num_complete_minibatches * mini_batch_size:]
        mini_batch_Y = shuffled_Y[:, num_complete_minibatches * mini_batch_size:]
        mini_batch = (mini_batch_X, mini_batch_Y)
        mini_batches.append(mini_batch)
    return mini_batches

def loadImages(folder, imageSize, maxData):
    catData = np.zeros((imageSize * imageSize * 3,1))
    counter = 0
    Y = np.full((1, maxData), int(1))
    for images in os.listdir(folder):
        #add 200 cat pictures
        if counter >= maxData:
            break
        counter +=1
        imageVector = processImage(images, folder, imageSize)
        catData = np.append(catData, imageVector, axis = 1)
    catData = catData[:,1:]
    catData = np.append(catData, Y, axis = 0)

    Y = np.full((1,maxData), int(0))
    dogData = np.zeros((imageSize * imageSize * 3,1))
    counter = 0
    for images in reversed(os.listdir(folder)):
        #adds 200 dogg pictures
        if counter >= maxData:
            break
        counter +=1
        imageVector = processImage(images, folder, imageSize)
        dogData = np.append(dogData, imageVector, axis = 1)
    dogData = dogData[:,1:]
    dogData = np.append(dogData, Y, axis = 0)

    Data = np.append(dogData, catData, axis = 1) #glue them together horizontally ofcourse

    #now randomize data:
    transposedData = Data.T
    np.random.shuffle(transposedData)
    Data = transposedData.T
    labels = np.array([Data[imageSize * imageSize * 3]])
    examples = np.delete(Data, imageSize * imageSize * 3, 0)
    examples = examples / 255
    return examples, labels

def processImage(images, folder, imageSize):
    #takes image returns one dimensional (vertical) image vector
        image = Image.open(os.path.join(folder, images)).resize((imageSize, imageSize))
        imageVector = asarray(image)
        imageVector = imageVector.reshape(imageVector.shape[0] * imageVector.shape[1] * imageVector.shape[2], 1)
        return imageVector

def hussle(X,Y):
    size = X.shape[0]
    Data = np.append(X, Y, axis = 0)
    transposedData = Data.T
    np.random.shuffle(transposedData)
    Data = transposedData.T
    labels = np.array([Data[size]])
    examples = np.delete(Data, size, 0)
    return examples, labels

def accuracy(prediction, Y):
    prediction[prediction > 0.5] = 1
    prediction[prediction <= 0.5] = 0

    difference = Y-prediction
    incorrect = np. count_nonzero(difference)
    false = incorrect / Y.shape[1] * 100
    accuracy = 100 - false
    return accuracy

def trainmodel(layerDims, mini_batches, iterations, learningRate, showcost, lambd, Xtest, Ytest, keepProb = 1):
    initialShowCost = showcost
    costs = []
    iterationlist = []
    #parameter initialization (xavier)
    initializeX, initializeY = mini_batches[0]
    parameters, Vdw, Vdb, Sdw, Sdb = parameterInitialization(initializeX, layerDims)
    counter = 0
    for i in range(iterations):
        print(i)
        for mini_batch in mini_batches:
            X,Y = mini_batch
            counter +=1
            iterationlist.append(counter)
            AL, forward_cache, D = L_model_forward(X, parameters, keepProb)

            #compute cost (maybe try out L2)
            thiscost = cost(AL, Y, parameters, lambd )
            costs.append(thiscost)
            #backward propagation (with gradient checking)
            grads = L_model_backward(AL, Y, forward_cache, lambd, parameters, keepProb, D)
            #gradCheck:

            #updating weights
            parameters = updateWeights(grads, parameters, layerDims, learningRate)
        if i == showcost:
            plt.title("cost function")
            plt.xlabel("Iterations")
            plt.ylabel("Cost")
            plt.plot(iterationlist, costs, color ="red")
            plt.show()
            showcost += initialShowCost
            trainingPrediction = makePrediction(X, parameters) #no dropout for test time
            testPrediction = makePrediction(Xtest, parameters) #no dropout for test time
            print("trainingAccuracy", accuracy(trainingPrediction, Y))
            print("testAccuracy", accuracy(testPrediction, Ytest))
    return parameters

def cost(A, Y, parameters, lambd):
    m = Y.shape[1]
    losses = -(np.multiply(np.log(A), Y) + np.multiply((1-Y), np.log(1-A)))
    cost = 1/m * np.sum(losses)
    frobeniusNorm = 0
    for l in range(1, int(len(parameters)/2) + 1):
        frobeniusNorm += np.sum(parameters["W" + str(l)] * parameters["W" + str(l)])
    L2_regularizationCost = lambd / (2*m) * frobeniusNorm
    cost += L2_regularizationCost
    return cost #checked

def makePrediction(X, parameters, keepProb = 1):
    AL, _, D = L_model_forward(X, parameters, keepProb)
    return AL

def updateWeights(grads, parameters, layerDims, learningRate):
    m = len(parameters) / 2
    for l in range(1, len(layerDims)):
        parameters["W" + str(l)] = parameters["W" + str(l)] - learningRate * grads["dW" + str(l)]
        parameters["b" + str(l)] = parameters["b" + str(l)] - learningRate * grads["db" + str(l)]
    return parameters

def parameterInitialization(X, layerDims):
    numPreviousNeurons = X.shape[0]
    parameters = {}
    for l in range(1, len(layerDims)):
        parameters["W" + str(l)] = np.random.randn(layerDims[l], layerDims[l-1]) * np.sqrt(1/layerDims[l-1])
        parameters["b" + str(l)] = np.zeros((layerDims[l], 1))
        numPreviousNeurons = layerDims[l]
    Vdw = 0
    Vdb = 0
    Sdw = 0
    Sdb = 0
    return parameters, Vdw, Vdb, Sdw, Sdb

def L_model_forward(X, parameters, keepProb):
    caches = []
    A = X
    #no dropout in input layer
    L = int(len(parameters) / 2)
    D = []
    for l in range(1, L):
        A_prev = A
        A, cache = linearActivationForward(A_prev, parameters["W" + str(l)], parameters["b" + str(l)], "relu")
        #store d one cache ahead for backward prop
        A,d = dropout(A, keepProb)
        D.append(d)
        #linear activation forward returns the next A for the next layer but also returns a cache containing
        #linear (W, A_PREV, b) and activation (Z, "activation" => "relu") cache
        caches.append(cache)
    AL, cache = linearActivationForward(A, parameters["W" + str(L)], parameters["b" + str(L)], "sigmoid")
    #no dropout in output layer
    caches.append(cache)
    return AL, caches, D

def dropout(A, keepProb):
    d = np.random.rand(A.shape[0], A.shape[1]) < keepProb
    A = np.multiply(A,d)
    A /= keepProb
    return A, d

def L_model_backward(AL, Y, caches, lambd, parameters, keepProb, D):
    grads = {}
    L = len(caches)
    m = Y.shape[1]
    dAl = -np.divide(Y, AL) + np.divide(1-Y, 1-AL)
    currentCache = caches[-1]
    dA_prev_temp, dW_temp, db_temp = linear_activation_backward(dAl, currentCache, "sigmoid")
    currentD = D[-1]
    dA_prev_temp = np.multiply(dA_prev_temp, currentD)/keepProb
    #grads["dA" + str(L-1)] = dA_prev_temp
    grads["dW" + str(L)] = dW_temp
    grads["db" + str(L)] = db_temp
    for l in reversed(range(L-1)):
        dA = dA_prev_temp
        currentCache = caches[l]
        dA_prev_temp, dW_temp, db_temp = linear_activation_backward(dA, currentCache, "relu")
        if l != 0:
            #no dropout for input layer when l == 0
            currentD = D[l - 1]
            dA_prev_temp = np.multiply(dA_prev_temp, currentD) / keepProb
        #grads["dA" + str(l)] = dA_prev_temp
        grads["dW" + str(l+1)] = dW_temp + lambd/m * parameters["W" + str(l + 1)]
        grads["db" + str(l+1)] = db_temp
        #the l-1 makes sense, you give in the A of the l+1 in linearactivationbackward,
        #this function will return the previous a which is therefore l,
        #l-1 because you calculate current a with the next weights (see page 5 notes)
        #same logic goes for why L-1 when calculating dA[l-1]
    return grads

def linear_activation_backward(dA, cache, activation):
    linear_cache, activation_cache = cache
    W = linear_cache["W"]
    b = linear_cache["b"]
    A_prev = linear_cache["A_prev"]
    Z = activation_cache["Z"]
    m = Z.shape[1]

    if activation == "relu":
        A = relu(Z)
        dZ = reluBackward(Z, dA)
    elif activation == "tanh":
        A= tanh(Z)
        dZ = tanhBackward(Z, dA)
    elif activation == "sigmoid":
        A = sigmoid(Z)
        dZ = sigmoidBackward(Z, dA)
    dA_prev = np.dot(W.T, dZ)
    dW = 1/m * np.dot(dZ, A_prev.T)
    db = 1/m * np.sum(dZ, axis = 1, keepdims = True)
    return dA_prev, dW, db

def linearActivationForward(A_prev, W, b, activation):
    linear_cache = {
        "A_prev": A_prev,
        "W": W,
        "b": b
        }
    Z = np.dot(W, A_prev) + b
    activation_cache = {
        "Z": Z,
        "activation": activation
        }
    if activation == "tanh":
        A = tanh(Z)
    if activation == "sigmoid":
        A = sigmoid(Z)
    if activation == "relu":
        A = relu(Z)
    caches = (linear_cache, activation_cache)
    return A, caches

def tanh(Z):
    A = (np.exp(Z) - np.exp(-Z)) / (np.exp(Z) + np.exp(-Z))
    return A
def sigmoid(Z):
    A = 1 / (1+np.exp(-Z))
    return A
def relu(Z):
    Z[Z < 0] = 0 #each value beneath 0 will have activation 0, but Z otherwise
    A = Z
    return A

def tanhBackward(Z, dA):
    dZ = (1-np.power(tanh(Z),2)) * dA
    return dZ
def sigmoidBackward(Z, dA):
    dZ = (sigmoid(Z) * (1 - sigmoid(Z)))* dA
    return dZ
def reluBackward(Z, dA):
    Z[Z<0] = 0
    Z[Z>= 0] = 1
    dZ = Z * dA
    return dZ

def loadDataCSV(path):
    file = open(path)
    csvreader = csv.reader(file)
    dataset = []
    count = 0
    for row in csvreader:
           if count == 0:
               count +=1
               continue
           if count == 1:
               count = 2
               dataset = np.array([row])
               continue
           row = np.array([row])
           dataset = np.concatenate((dataset,row), axis = 0)

    trainSet = dataset[:int(len(dataset) * 0.7)].T
    testSet = dataset[int(len(dataset) * 0.7):].T
    Ytrain = trainSet[-1].astype(np.float).reshape(1, trainSet.shape[1])
    Ytest = testSet[-1].astype(np.float).reshape(1, testSet.shape[1])
    trainSet  = np.delete(trainSet, -1, 0)
    testSet = np.delete(testSet, -1, 0)
    Xtrain = trainSet.astype(np.float)
    Xtest = testSet.astype(np.float)
    means = 1 / Xtrain.shape[0] * np.sum(Xtrain, axis = 1, keepdims = True)
    Xtrain = Xtrain /means
    Xtrain = Xtrain / (np.sqrt(1/Xtrain.shape[0] * np.sum(Xtrain, axis =1, keepdims = True)))
    Xtest = Xtest / means
    Xtest = Xtest / (np.sqrt(1/Xtrain.shape[0] * np.sum(Xtrain, axis =1, keepdims = True)))

    return Xtrain, Ytrain, Xtest, Ytest

if __name__ == "__main__":
    main()


'''
International Morse Code defines a standard encoding where each letter is mapped to a series of dots and dashes, as follows:

    'a' maps to ".-",
    'b' maps to "-...",
    'c' maps to "-.-.", and so on.

For convenience, the full table for the 26 letters of the English alphabet is given below:

[".-","-...","-.-.","-..",".","..-.","--.","....","..",".---","-.-",".-..","--","-.","---",".--.","--.-",".-.","...","-","..-","...-",".--","-..-","-.--","--.."]

Given an array of strings words where each word can be written as a concatenation of the Morse code of each letter.

    For example, "cab" can be written as "-.-..--...", which is the concatenation of "-.-.", ".-", and "-...". We will call such a concatenation the transformation of a word.

Return the number of different transformations among all words we have.

Example 1:

Input: words = ["gin","zen","gig","msg"]
Output: 2
Explanation: The transformation of each word is:
"gin" -> "--...-."
"zen" -> "--...-."
"gig" -> "--...--."
"msg" -> "--...--."
There are 2 different transformations: "--...-." and "--...--.".

'''
from typing import List
def uniqueMorseRepresentations(words: List[str]) -> int:
    alpha = [chr(i) for i in range(97, 123)]
    code = [".-","-...","-.-.","-..",".","..-.","--.","....","..",".---","-.-",".-..","--","-.","---",".--.","--.-",".-.","...","-","..-","...-",".--","-..-","-.--","--.."]
    uni_code = dict(zip(alpha, code))

    def code(letter):
        result = ""
        for key, value in uni_code.items():
            if key == letter:
                result = result + value
        return result

    def word_to_code(word):
        return "".join([code(letter) for letter in word])

    result = [word_to_code(word) for word in words]
    return len(set(result))

words = ["gin","zen","gig","msg"]
print(uniqueMorseRepresentations(words))