import os
from random import *
import math

from PIL import Image, ImageDraw, ImageOps

no_cv = True
export_path = "output/out.svg"
draw_contours = True
draw_hatch = False

try:
	import numpy as np
	import cv2
except:
	print("Cannot import numpy/openCV. Switching to NO_CV mode.")
	no_cv = True


F_Blur = {
	(-2,-2):2,(-1,-2):4,(0,-2):5,(1,-2):4,(2,-2):2,
	(-2,-1):4,(-1,-1):9,(0,-1):12,(1,-1):9,(2,-1):4,
	(-2,0):5,(-1,0):12,(0,0):15,(1,0):12,(2,0):5,
	(-2,1):4,(-1,1):9,(0,1):12,(1,1):9,(2,1):4,
	(-2,2):2,(-1,2):4,(0,2):5,(1,2):4,(2,2):2,
}
F_SobelX = {
	(-1,-1): 1,
	(0,-1):  0,
	(1,-1):  -1,
	(-1,0):  2,
	(0,0):   0,
	(1,0):   -2,
	(-1,1):  1,
	(0,1):   0,
	(1,1):   -1}
F_SobelY = {
	(-1,-1): 1,
	(0,-1):  2,
	(1,-1):  1,
	(-1,0):  0,
	(0,0):   0,
	(1,0):   0,
	(-1,1):  -1,
	(0,1):   -2,
	(1,1):   -1}

def appmask(IM,masks):
	PX = IM.load()
	w,h = IM.size
	NPX = {}
	for x in range(0,w):
		for y in range(0,h):
			a = [0]*len(masks)
			for i in range(len(masks)):
				for p in masks[i].keys():
					if 0<x+p[0]<w and 0<y+p[1]<h:
						a[i] += PX[x+p[0],y+p[1]] * masks[i][p]
				if sum(masks[i].values())!=0:
					a[i] = a[i] / sum(masks[i].values())
			NPX[x,y]=int(sum([v**2 for v in a])**0.5)
	for x in range(0,w):
		for y in range(0,h):
			PX[x,y] = NPX[x,y]


def distsum(*args):
	"""
	Takes a list of pairs of points, finds the distance between each point
	pair, and returns the sum of all distances.
	"""
	dists = []
	for i in range(1, len(args):
		a = args[i][0]-args[i-1][0]
		b = args[i][1]-args[i-1][1]
		dist = (a**2 + b**2)**0.5
		dists.append(dist)
	return sum(dists)


def sortlines(lines):
	print("optimizing stroke sequence...")
	clines = lines[:]
	slines = [clines.pop(0)]
	while clines != []:
		x,s,r = None,1000000,False
		for l in clines:
			d = distsum(l[0],slines[-1][-1])
			dr = distsum(l[-1],slines[-1][-1])
			if d < s:
				x,s,r = l[:],d,False
			if dr < s:
				x,s,r = l[:],s,True

		clines.remove(x)
		if r == True:
			x = x[::-1]
		slines.append(x)
	return slines



def auto_canny(img, sigma=0.33):
	"""
	Automatically determines appropriate upper and lower boundries for the Canny function.
	"""
	med = np.median(img)
	lower = int(max(0, (1.0 - sigma) * med))
	upper = int(min(255, (1.0 + sigma) * med))
	edges = cv2.Canny(img, lower, upper)
	return edges


def find_edges(IM):
	print("finding edges...")
	no_cv = True
	if no_cv:
		#appmask(IM,[F_Blur])
		appmask(IM,[F_SobelX,F_SobelY])
	else:
		im = np.array(IM) 
		im = cv2.GaussianBlur(im,(3,3),0)
		#im = cv2.Canny(im,100,200)
		im = auto_canny(im)
		IM = Image.fromarray(im)
	return IM.point(lambda p: p > 128 and 255)  


def getdots(IM):
	print("getting contour points...")
	PX = IM.load()
	dots = []
	w,h = IM.size
	for y in range(h-1):
		row = []
		for x in range(1,w):
			if PX[x,y] == 255:
				if len(row) > 0:
					if x-row[-1][0] == row[-1][-1]+1:
						row[-1] = (row[-1][0],row[-1][-1]+1)
					else:
						row.append((x,0))
				else:
					row.append((x,0))
		dots.append(row)
	return dots


def connectdots(dots):
	print("connecting contour points...")
	contours = []
	for y in range(len(dots)):
		for x,v in dots[y]:
			if v > -1:
				if y == 0:
					contours.append([(x,y)])
				else:
					closest = -1
					cdist = 100
					for x0,v0 in dots[y-1]:
						if abs(x0-x) < cdist:
							cdist = abs(x0-x)
							closest = x0

					if cdist > 3:
						contours.append([(x,y)])
					else:
						found = 0
						for i in range(len(contours)):
							if contours[i][-1] == (closest,y-1):
								contours[i].append((x,y,))
								found = 1
								break
						if found == 0:
							contours.append([(x,y)])
		for c in contours:
			if c[-1][1] < y-1 and len(c)<4:
				contours.remove(c)
	return contours


def getcontours(IM,sc=2):
	print("generating contours...")
	IM = find_edges(IM)
	IM1 = IM.copy()
	IM2 = IM.rotate(-90,expand=True).transpose(Image.FLIP_LEFT_RIGHT)
	dots1 = getdots(IM1)
	contours1 = connectdots(dots1)
	dots2 = getdots(IM2)
	contours2 = connectdots(dots2)

	for i in range(len(contours2)):
		contours2[i] = [(c[1],c[0]) for c in contours2[i]]	
	contours = contours1+contours2

	for i in range(len(contours)):
		for j in range(len(contours)):
			if len(contours[i]) > 0 and len(contours[j])>0:
				if distsum(contours[j][0],contours[i][-1]) < 8:
					contours[i] = contours[i]+contours[j]
					contours[j] = []

	for i in range(len(contours)):
		contours[i] = [contours[i][j] for j in range(0,len(contours[i]),8)]


	contours = [c for c in contours if len(c) > 1]

	for i in range(0,len(contours)):
		contours[i] = [(v[0]*sc,v[1]*sc) for v in contours[i]]

	for i in range(0,len(contours)):
		for j in range(0,len(contours[i])):
			contours[i][j] = int(contours[i][j][0]+10),int(contours[i][j][1]+10)

	return contours


def hatch(IM,sc=16):
	print("hatching...")
	PX = IM.load()
	w,h = IM.size
	lg1 = []
	lg2 = []
	for x0 in range(w):
		for y0 in range(h):
			x = x0*sc
			y = y0*sc
			if PX[x0,y0] > 144:
				pass
				
			elif PX[x0,y0] > 64:
				lg1.append([(x,y+sc/4),(x+sc,y+sc/4)])
			elif PX[x0,y0] > 16:
				lg1.append([(x,y+sc/4),(x+sc,y+sc/4)])
				lg2.append([(x+sc,y),(x,y+sc)])

			else:
				lg1.append([(x,y+sc/4),(x+sc,y+sc/4)])
				lg1.append([(x,y+sc/2+sc/4),(x+sc,y+sc/2+sc/4)])
				lg2.append([(x+sc,y),(x,y+sc)])

	lines = [lg1,lg2]
	for k in range(0,len(lines)):
		for i in range(0,len(lines[k])):
			for j in range(0,len(lines[k])):
				if lines[k][i] != [] and lines[k][j] != []:
					if lines[k][i][-1] == lines[k][j][0]:
						lines[k][i] = lines[k][i]+lines[k][j][1:]
						lines[k][j] = []
		lines[k] = [l for l in lines[k] if len(l) > 0]
	lines = lines[0]+lines[1]

	for i in range(0,len(lines)):
		for j in range(0,len(lines[i])):
			lines[i][j] = int(lines[i][j][0]+sc),int(lines[i][j][1]+sc)-j
	return lines


def sketch(path, export_path=None, resolution=1024, hatch_size=16, contour_simplify=2):
	image = Image.open(path)

	image = image.convert("L")
	image = ImageOps.autocontrast(image ,10)

	lines = []
	if draw_contours:
		width = int(resolution/contour_simplify)
		height = int(resolution/contour_simplify*image.size[0]/image.size[1])
		lines += getcontours(image.resize((width, height)), contour_simplify)
	if draw_hatch:
		width = int(resolution/hatch_size)
		height = int(resolution/hatch_size*image.size[0]/image.size[1])
		lines += hatch(image.resize((width, height)), hatch_size)

	lines = sortlines(lines)

	export_path = path.replace("received", "converted")
	export_path = os.path.splitext(export_path)[0]
	with open(export_path + ".svg", "w") as file:
		file.write(makePathSvg(lines))
	with open(export_path + ".ngc", "w") as file:
		file.write(makeGcode(lines))

	print(len(lines), "strokes.")
	print("done.")
	return lines


def makePolySvg(lines):
	"""
	Uses the provided set of contours to generate an SVG file and returns it as a string. Contours get
	written as polyline as objects.
	"""
	print("generating svg file...")
	out = '<svg xmlns="http://www.w3.org/2000/svg" version="1.1">\n'
	for l in lines:
		l = ",".join([str(p[0]*0.5)+","+str(p[1]*0.5) for p in l])
		out += '<polyline points="'+l+'" stroke="black" stroke-width="2" fill="none" />\n'
	out += '</svg>'
	return out


def makePathSvg(lines):
	"""
	Uses the provided set of contours to generate an SVG file and returns it as a string. Contours get
	written as path objects.
	"""
	print("generating svg file...")
	out = '<svg xmlns="http://www.w3.org/2000/svg" version="1.1">\n'
	for l in lines:
		l = ",".join([str(p[0]*0.5)+","+str(p[1]*0.5) for p in l])
		out += '<path d="M'+l+'" stroke="black" stroke-width="2" fill="none" />\n'
	out += "</svg>"
	return out


def makeGcode(lines, paper_size="letter"):
	"""
	Converts the provided contour lines into G-code commands and returns them as a string.
	"""
	paper_sizes = {
		"letter": (67.46875, 87.3125),
		"ledger": (87.3125, 139.2903226),
		"max": (90, 140)}
	tot = []
	for line in lines:
		tot += line
	max_x = max([p[0] for p in tot])
	max_y = max([p[1] for p in tot])
	d_x = paper_sizes[paper_size][0] / max_x
	d_y = paper_sizes[paper_size][1] / max_y
	scale = min(d_x, d_y)

	print("generating gcode file...")
	out = "$X\n$32=1\nM03\nF600\nG17 G21 G90 G54\nG01\n\n"
	for line in lines:
		start = line.pop(0)
		out += "S1000\n"
		#out += f"X{start[0]*0.5} Y{start[1]*0.5}\n"
		out += "X" + str(start[0]*scale) + " Y" + str(start[1]*scale) + "\n"
		out += "S0\n"
		for point in line:
			#out += f"X{point[0]*0.5} Y{point[1]*0.5}\n"
			out += "X" + str(point[0]*scale) + " Y" + str(point[1]*scale) + "\n"
		out += "\n"
	out += "S1000\nX0 Y0\nM2\n"
	return out


if __name__ == "__main__":
	import argparse

	parser = argparse.ArgumentParser(
		description="Convert image to vectorized line drawing for plotters.")
	parser.add_argument(
		'-i',
		'--input',
		nargs='?',
		help='Input path')
	parser.add_argument(
		'-o',
		'--output',
		nargs='?',
		help='Output path.')
	parser.add_argument(
		'--no_contour',
		action='store_true',
		help="Don't draw contours.")
	parser.add_argument(
		'--hatch',
		action='store_true',
		help='Enable hatching.')
	parser.add_argument(
		'--no_cv',
		action='store_true',
		help="Don't use OpenCV.")
	parser.add_argument(
		"--resolution",
		default=1024,
		type=int,
		help="Resolution. eg. 512, 1024, 2048")
	parser.add_argument(
		'--contour_simplify',
		default=2,
		type=int,
		help='Level of contour simplification. eg. 1, 2, 3')
	parser.add_argument(
		'--hatch_size',
		default=16,
		type=int,
		help='Patch size of hatches. eg. 8, 16, 32')
	args = parser.parse_args()

	export_path = args.output
	draw_hatch = args.hatch
	draw_contours = not args.no_contour
	no_cv = args.no_cv

	sketch(args.input, args.output, args.resolution, args.contour_simplify, args.hatch_size)