import argparse
import sys
import cv2
import numpy as np
import math
from json import dumps
from os.path import exists
from common import image_resize, display, normalize_angle

#clockwise from top left
order = [ 1, 3, 4, 6, 5, 2 ]

def get_center (contour) : 
	M = cv2.moments(contour)
	cX = int(M["m10"] / M["m00"])
	cY = int(M["m01"] / M["m00"])
	return cX, cY

def draw_line (image, hps, a, b) :
	print(f'{a} -> {b}')
	lA = (hps[a-1]['x'], hps[a-1]['y'])
	lB = (hps[b-1]['x'], hps[b-1]['y'])
	cv2.line(image, lA, lB, [0, 255, 0], 10)
	return (lA, lB)

def horiz_angle (line, rotate = 0) :
	deltaY = line[1][1] - line[0][1] #P2_y - P1_y
	deltaX = line[1][0] - line[0][0] #P2_x - P1_x

	angleInDegrees = normalize_angle(math.degrees(math.atan2(deltaY, deltaX) + rotate))
	return angleInDegrees

def verts_angle (line) :
	angleInDegrees = normalize_angle(horiz_angle(line, math.pi/2))
	return angleInDegrees

def is_close (point, points) :
	for pt in points :
		if math.dist(point, pt) < 100 :
			return True
	return False

def mean (lst):
    return sum(lst) / len(lst)

def find_hole_punches (img) :
	left=-1
	right=-1
	top=-1
	bottom=-1

	if orientation :
		left = width * 0.2
		right = width * 0.8
	else :
		top = height * 0.2
		bottom = height * 0.8

	gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
	blur = cv2.medianBlur(gray, 31)
	ret, thresh = cv2.threshold(blur, 200, 255, cv2.THRESH_BINARY)
	canny = cv2.Canny(thresh, 75, 200)
	contours, hierarchy = cv2.findContours(canny, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)

	contourList = []
	areaList = []
	for contour in contours:
		approx = cv2.approxPolyDP(contour, 0.03 * cv2.arcLength(contour, True), True)
		if cv2.isContourConvex(approx) :
			cX, cY = get_center(contour)
			if (orientation and ( cX < left or cX > right) ) or ( not orientation and ( cY < top or cY > bottom)) :
				area = cv2.contourArea(contour)
				areaList.append(area)
				contourList.append(contour)

	maxArea=0
	maxIndex=0

	#reduce to lambda
	for i in range(len(areaList)) :
		area = areaList[i]
		if area > maxArea:
			maxArea = area
			maxIndex = i

	count = 0
	holePunches = []
	centers = []
	areaRange = 0
	topLeft = None
	minDist = 1000000

	# pretty good
	# add position constraint

	while count < 6 :
		areaRange+=1
		for i in range(len(areaList)) :
			area = areaList[i]
			if area == maxArea or area * ((100 + areaRange) / 100) > maxArea  :
				cX, cY = get_center(contourList[i])
				if is_close((cX, cY), centers) :
					continue
				centers.append((cX, cY))
				print(f'{cX},{cY}')
				hp = {
					'x' : cX,
					'y' : cY,
					'contour' : contourList[i],
					'dist' : math.dist((cX, cY), (0, 0)),
					'order': -1
				}
				if hp['dist'] < minDist :
					minDist = hp['dist']
					topLeft = hp
				holePunches.append(hp)
				count+=1

	for hp in holePunches :
		hp['dist'] = math.dist( (topLeft['x'], topLeft['y']), (hp['x'], hp['y']) )

	print(f'Hole punches: {len(holePunches)}')
	print(f'Found hole punches within {areaRange}% of largest')

	if len(holePunches) != 6:
		print(f'Wrong number of hole punches, exiting...')
		exit(4)

	holePunches = sorted(holePunches, key = lambda hp: hp['dist'])
	
	i = 0
	for hp in holePunches :
		hp['order'] = i 
		#cv2.putText(img, str(i + 1), (hp['x'], hp['y']), cv2.FONT_HERSHEY_SIMPLEX, 20, (0, 0, 255), 5, cv2.LINE_AA, False)
		i+=1

	return holePunches

def simplify_hole_punches (holePunches) :
	simple = {}
	for hp in holePunches :
		simple[hp['order']] = {
			'x' : hp['x'],
			'y' : hp['y']
		}
	return simple

def correct_rotation (img, original, holePunches) :
	horizLines = [
		(3, 1),
		(6, 4),
		(5, 2)
	]
	vertsLines = [
		(1, 2),
		(1, 4), #double long left
		(1, 4), #
		(3, 5),
		(3, 6), #double long right
		(3, 6),  #
		(2, 4),
		(5, 6)
	]

	rotations = []

	for h in horizLines :
		line = draw_line(img, holePunches, h[0], h[1])
		angle = horiz_angle(line)
		print(angle)
		rotations.append(angle)

	for v in vertsLines :
		line = draw_line(img, holePunches, v[0], v[1])
		angle = verts_angle(line)
		print(angle)
		rotations.append(angle)

	correctionRotation = mean(rotations) - 180
	print(f'Mean rotation: {correctionRotation}')

	(cX, cY) = (width // 2, height // 2)
	M = cv2.getRotationMatrix2D((cX, cY), correctionRotation, 1.0)

	#create rotation of original
	return cv2.warpAffine(original, M, (width, height))

def create_blank (w, h, rgb_color = (255, 255, 255)) :
	blank = np.zeros([h, w, 3], dtype=np.uint8)
	color = tuple(reversed(rgb_color))
	blank[:] = color
	return blank

def get_mean_rect (holePunches) :
	left = 0
	right = 0
	top = 0
	bottom = 0
	for hp in holePunches :
		if hp['order'] == 0 :
			left += float(hp['x'])
			top += float(hp['y'])
		elif hp['order'] == 2 :
			right += float(hp['x'])
			top += float(hp['y'])
		elif hp['order'] == 3 :
			left += float(hp['x'])
			bottom += float(hp['y'])
		elif hp['order'] == 5 :
			right += float(hp['x'])
			bottom += float(hp['y'])
	w = round((right / 2.0) - (left / 2.0))
	h = round((bottom / 2.0) - (top / 2.0))
	return (w, h)


def center_within (larger, smaller) :
	w1 = larger[0]
	h1 = larger[1]
	w2 = smaller[0]
	h2 = smaller[1]
	x = ((w1 - w2) / 2)
	y = ((h1 - h2) / 2)
	return (int(x), int(y))

# If we consider (0,0) as top left corner of image called 
# im with left-to-right as x direction and top-to-bottom 
# as y direction. and we have (x1,y1) as the top-left vertex 
# and (x2,y2) as the bottom-right vertex of a rectangle 
# region within that image, then:
#
# roi = im[y1:y2, x1:x2]

def crop (img, xoffset, yoffset, w, h) :
	#crop_img = img[y:y+h, x:x+w].copy()
	return im[yoffset:yoffset+w, xoffset:xoffset+w].copy()

def normalize_image(blank, rotated, offset, tl) :
	rotatedHeight, rotatedWidth = rotated.shape[:2]
	normalHeight, width = blank.shape[:2]

	diffX = offset[0] - tl["x"]
	diffY = offset[1] - tl["y"]

	#print(f'diffX  : {diffX}')
	#print(f'diffY  : {diffY}')

	crop = rotated.copy()
	if diffX < 0 :
		crop = crop[0:rotatedHeight, abs(diffX):rotatedWidth]
		rotatedHeight, rotatedWidth = crop.shape[:2]
		#print('Cropped X')
		#print(f'Rotated: {rotatedWidth},{rotatedHeight}')
		diffX = 0

	if diffY < 0 :
		crop = crop[abs(diffY):rotatedHeight, 0:rotatedWidth]
		rotatedHeight, rotatedWidth = crop.shape[:2]
		#print('Cropped Y')
		#print(f'Rotated: {rotatedWidth},{rotatedHeight}')
		diffY = 0

	if rotatedWidth > width :
		crop = crop[0:rotatedHeight, 0:rotatedWidth-(rotatedWidth - width)]
		rotatedHeight, rotatedWidth = crop.shape[:2]
		#print('Cropped X')
		#print(f'Rotated: {rotatedWidth},{rotatedHeight}')

	if  rotatedHeight > normalHeight :
		crop = crop[0:rotatedHeight-(rotatedHeight - normalHeight), 0:width]
		rotatedHeight, rotatedWidth = crop.shape[:2]
		#print('Cropped Y')
		#print(f'Rotated: {rotatedWidth},{rotatedHeight}')

	#print(f'diffX  : {diffX}')
	#print(f'diffY  : {diffY}')

	#print(f'Rotated: {rotatedWidth},{rotatedHeight}')
	#print(f'Blank  : {width},{normalHeight}')

	cropHeight = normalHeight
	cropWidth = width 

	if normalHeight > rotatedHeight :
		cropHeight = rotatedHeight

	if width > rotatedWidth :
		cropWidth = rotatedWidth

	blank[diffY:cropHeight, diffX:cropWidth] = crop[0:cropHeight-diffY, 0:cropWidth - diffX]

	return blank

#
# NORMALIZE
#
# The purpose of this script is to normalize scanned hole punch images to a standard
# size, with standard padding around the holes and oriented straight up and down.
# Useful for converting scanned plotter-drawn pages into 4 frame sheets and for 
# generating images to print.
#

parser = argparse.ArgumentParser(description='Normalize a scanned hole punch image.')
parser.add_argument('input',  type=str, help='Scan to normalize')
parser.add_argument('output',  type=str, help='Normalized file to output')
args = parser.parse_args()

scanImage = args.input
normalImage = args.output

if not exists(scanImage) :
	print('Scan provided does not exist')
	exit(1)

pageDim = (11, 8.5)
pageRatio = pageDim[1] / pageDim[0]

print(f'Normalizing {scanImage} as {normalImage}')

original = cv2.imread(scanImage)
img = original.copy()

height, width = img.shape[:2]
orientation = height > width
if not orientation :
	print(f'Scan is not in portrait mode, exiting...')
	exit(3)

normalHeight = round(float(width) / pageRatio)

holePunches = find_hole_punches(img)
rotated = correct_rotation(img, original, holePunches)
rotatedHeight, rotatedWidth = rotated.shape[:2]

holePunches = find_hole_punches(rotated)
blank = create_blank(width, normalHeight)

tl = None
for hp in holePunches :
	if tl is None :
		tl = hp
	#print(f'{hp["order"] + 1} {hp["x"]},{hp["y"]}')

# the mean rectangle is the average width and height
# determined by the four corner hole punches
meanRect = get_mean_rect(holePunches)
print(f'Mean rectangle: {meanRect[0]},{meanRect[1]}')

# offset is the position within the new normal image
# the top left hole punch should be centered to
offset = center_within((width, normalHeight), meanRect)
#print(f'Offset : {offset[0]},{offset[1]}')
#print(f'Topleft: {tl["x"]},{tl["y"]}')
#print(f'Rotated: {rotatedWidth},{rotatedHeight}')
print(f'Normal : {width},{normalHeight}')

#cv2.rectangle(blank, offset, (offset[0]+meanRect[0], offset[1]+meanRect[1]), (255, 0, 0), thickness=20)

normal = normalize_image(blank, rotated, offset, tl)

print(f'Writing normalized image to {normalImage}')
cv2.imwrite(normalImage, normal)

evaluation = find_hole_punches(normal)
jsonOut = simplify_hole_punches(evaluation)

with open(f'{normalImage}.json', 'w') as output:
	output.write(dumps(jsonOut, sort_keys = True, indent = 4))
	print(f'Wrote hole punch definition file to {normalImage}.json')