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

DEBUG = True
registrationMarkThreshold = 0.7
methodIndex = 0
marks = []

#clockwise from top left
order = [ 0, 2, 3, 5, 4, 1 ]
matchMethods = [cv2.TM_CCOEFF, cv2.TM_CCOEFF_NORMED, cv2.TM_CCORR, cv2.TM_CCORR_NORMED, cv2.TM_SQDIFF, cv2.TM_SQDIFF_NORMED]

#
# TEMPLATE
#

parser = argparse.ArgumentParser(description='Create a template from normalized scan file.')
parser.add_argument('normalized',  type=str, help='Normalized scan to build template from.')
parser.add_argument('registration', type=str, nargs='?', help='Registration mark to use.');
#parser.add_argument('output',  type=str, help='Normalized file to output', default=None)
args = parser.parse_args()

normalImage = args.normalized

if not exists(normalImage) :
	print('Normalized scan does not exist, please provide one that does')
	exit(2)

normalText = normalImage + '.json'

if not exists(normalText) :
	print('Corresponding normalized scan text does not exist, please generate one')
	exit(3)

if args.registraion is not None :
	if not exists(args.registraion) :
		registrationMark = cv2.imread(args.registraion)
		print(f'Registration mark {args.registraion} does not exist')
		exit(4)
	print(f'Using registration mark {args.registraion}')

def get_dpi (width) :
	if width == 2550 :
		return 300
	if width == 5100 :
		return 600
	if width == 10200 :
		return 1200

print(f'Building template from scan {basename(normalImage)}')

holePunches = read_json(normalText)
original = cv2.imread(normalImage)
img = original.copy()
height, width = img.shape[:2]
orientation = height > width
dpi = get_dpi(width)

try: 
	registrationMark
except NameError: 
	registrationMark = None
if registrationMark is None :
	registrationMark = cv2.imread(f'./registrationMark/{dpi}.png', 0)
w, h = registrationMark.shape[:2]

if not orientation :
	print(f'Scan is not in portrait mode, exiting...')
	exit(3)

#print(holePunches)

def get_distance(ref, point):
	# print('ref: {} , point: {}'.format(ref, point))
	x1, y1 = ref[0], ref[1]
	x2, y2 = point[0], point[1]
	return math.hypot(x2 - x1, y2 - y1)

def group_points(points):
    groups = {}
    groupnum = 0
    while len(points) > 1:
        groupnum += 1
        key = str(groupnum)
        groups[key] = []
        ref = points.pop(0)
        for i, point in enumerate(points):
            d = get_distance(ref, point)
            if d < 30:
                groups[key].append(points[i])
                points[i] = None
        points = list(filter(lambda x: x is not None, points))
    return list([[int(np.mean(list([x[0] for x in groups[arr]]))), int(np.mean(list([x[1] for x in groups[arr]])))] for arr in groups])

def find_closest (pt, pts) :
	return pts[min(range(len(pts)), key = lambda i: get_distance(pts[i], pt))]

def find_in_half (half) :
	halfGray = cv2.cvtColor(half, cv2.COLOR_BGR2GRAY)
	res = cv2.matchTemplate(halfGray, registrationMark, matchMethods[methodIndex])
	threshold = registrationMarkThreshold
	loc = np.where( res >= threshold )
	for pt in zip(*loc[::-1]):
		cv2.rectangle(half, pt, (pt[0] + w, pt[1] + h), (0,0,255), 2)
	return list(zip(*loc[::-1]))

ttly = holePunches['0']['y']-round(height*0.05)
ttlx = holePunches['0']['x']
topHalf = img[ttly:holePunches['1']['y']+round(height*0.1), ttlx:holePunches['2']['x']]
while True: 
	topHalfPts = find_in_half(topHalf)
	if len(topHalfPts) < 12 :
		if methodIndex == len(matchMethods) :
			registrationMarkThreshold -= 0.005
			methodIndex = 0
			print(f'Found only {len(topHalfPts)}, decreasing threshold to {registrationMarkThreshold}')
		else :
			methodIndex += 1
			print(f'Found only {len(topHalfPts)}, switching to method {matchMethods[methodIndex]}')
		
	elif len(topHalfPts) > 12 :
		if methodIndex == len(matchMethods) :
			registrationMarkThreshold += 0.005
			methodIndex = 0
			print(f'Found {len(topHalfPts)}, increasing threshold to {registrationMarkThreshold}')
		else :
			methodIndex += 1
			print(f'Found {len(topHalfPts)}, switching to method {matchMethods[methodIndex]}')
		
	else :
		break

thpts = group_points(topHalfPts)
for pt in thpts :
	#print(f'{ttlx + pt[0]},{ttly + pt[1]}')
	marks.append((ttlx + pt[0] + round(w / 2), ttly + pt[1] + round(h / 2),))

print(f'Found {len(thpts)} points')
#display(topHalf)

btly = holePunches['4']['y']-round(height*0.1)
btlx = holePunches['3']['x']
bottomHalf = img[btly:holePunches['5']['y']+round(height*0.05), btlx:holePunches['4']['x']]
bottomHalfPts = find_in_half(bottomHalf)
bhpts = group_points(bottomHalfPts)
for pt in bhpts :
	#print(f'{btlx + pt[0]},{btly + pt[1]}')
	marks.append((btlx + pt[0] + round(w / 2), btly + pt[1] + round(h / 2), ))

clean = original.copy()
for pt in marks :
	#print(pt)
	cv2.circle(clean, pt, 50, (0,0,255), -1)

print(f'Found {len(bhpts)} points')

if len(marks) != 16 :
	print(f'{len(marks)} != 16 marks')
	exit(1)

##
# TOP LEFT
##

topLeftTR = find_closest((0,0,), marks)
#print(topLeftTR)
topLeftSorted = sorted(marks, key=lambda e: get_distance(e, topLeftTR))
topLeftSorted.remove(topLeftTR)

if abs(topLeftSorted[0][1] - topLeftTR[1]) < abs(topLeftSorted[1][1] - topLeftTR[1]) :
	topLeftBR = topLeftSorted[0]
else :
	topLeftBR = topLeftSorted[1]

topLeftTL = topLeftSorted[2]
topLeftBL = topLeftSorted[4]

if DEBUG :
	cv2.line(clean, topLeftTR, topLeftBR, (0,0,255,), 3)
	cv2.line(clean, topLeftTR, topLeftTL, (0,0,255,), 3)
	cv2.line(clean, topLeftBR, topLeftBL, (0,0,255,), 3)
	cv2.line(clean, topLeftTL, topLeftBL, (0,0,255,), 3)

##
# TOP RIGHT
##

topRightTL = find_closest((width,0,), marks)
#print(topRightTL)

topRightSorted = sorted(topLeftSorted, key=lambda e: get_distance(e, topRightTL))
topRightSorted.remove(topRightTL)
topRightSorted.remove(topLeftBR)
topRightSorted.remove(topLeftTL)
topRightSorted.remove(topLeftBL)

topRightBL = topRightSorted[0]
topRightTR = topRightSorted[1]
if abs(topRightTL[0] - topRightSorted[2][0]) > abs(topRightTL[0] - topRightSorted[3][0]) :
	topRightBR = topRightSorted[2]
else :
	topRightBR = topRightSorted[3]

if DEBUG :
	cv2.line(clean, topRightTR, topRightBR, (0,0,255,), 3)
	cv2.line(clean, topRightTR, topRightTL, (0,0,255,), 3)
	cv2.line(clean, topRightBR, topRightBL, (0,0,255,), 3)
	cv2.line(clean, topRightTL, topRightBL, (0,0,255,), 3)

##
# BOTTOM LEFT
##

bottomLeftTL = find_closest((0,height,), marks)

bottomLeftSorted = sorted(topRightSorted, key=lambda e: get_distance(e, bottomLeftTL))
bottomLeftSorted.remove(bottomLeftTL)
bottomLeftSorted.remove(topRightBR)
bottomLeftSorted.remove(topRightTR)
bottomLeftSorted.remove(topRightBL)

bottomLeftBL = bottomLeftSorted[0]
bottomLeftTR = bottomLeftSorted[2]
bottomLeftBR = bottomLeftSorted[3]

if DEBUG :
	cv2.line(clean, bottomLeftTR, bottomLeftBR, (0,0,255,), 3)
	cv2.line(clean, bottomLeftTR, bottomLeftTL, (0,0,255,), 3)
	cv2.line(clean, bottomLeftBR, bottomLeftBL, (0,0,255,), 3)
	cv2.line(clean, bottomLeftTL, bottomLeftBL, (0,0,255,), 3)

##
# BOTTOM RIGHT
##

bottomRightTR = find_closest((width,height,), marks)

bottomRightSorted = sorted(bottomLeftSorted, key=lambda e: get_distance(e, bottomRightTR))
bottomRightSorted.remove(bottomRightTR)
bottomRightSorted.remove(bottomLeftBR)
bottomRightSorted.remove(bottomLeftTR)
bottomRightSorted.remove(bottomLeftBL)

bottomRightBR = bottomRightSorted[0]
bottomRightTL = bottomRightSorted[1]
bottomRightBL = bottomRightSorted[2]

if DEBUG :
	cv2.line(clean, bottomRightTR, bottomRightBR, (0,0,255,), 3)
	cv2.line(clean, bottomRightTR, bottomRightTL, (0,0,255,), 3)
	cv2.line(clean, bottomRightBR, bottomRightBL, (0,0,255,), 3)
	cv2.line(clean, bottomRightTL, bottomRightBL, (0,0,255,), 3)

jsonOut = {
	'width' : width,
	'height' : height,
	'holePunches' : holePunches,
	'0' : {
		'0' : {
			'x' : bottomLeftTL[0],
			'y' : bottomLeftTL[1]
		},
		'1' : {
			'x' : bottomLeftTR[0],
			'y' : bottomLeftTR[1]
		},
		'2' : {
			'x' : bottomLeftBR[0],
			'y' : bottomLeftBR[1]
		},
		'3' : {
			'x' : bottomLeftBL[0],
			'y' : bottomLeftBL[1]
		}
	},
	'1' : {
		'0' : {
			'x' : topLeftTL[0],
			'y' : topLeftTL[1]
		},
		'1' : {
			'x' : topLeftTR[0],
			'y' : topLeftTR[1]
		},
		'2' : {
			'x' : topLeftBR[0],
			'y' : topLeftBR[1]
		},
		'3' : {
			'x' : topLeftBL[0],
			'y' : topLeftBL[1]
		}
	},
	'2' : {
		'0' : {
			'x' : topRightTL[0],
			'y' : topRightTL[1]
		},
		'1' : {
			'x' : topRightTR[0],
			'y' : topRightTR[1]
		},
		'2' : {
			'x' : topRightBR[0],
			'y' : topRightBR[1]
		},
		'3' : {
			'x' : topRightBL[0],
			'y' : topRightBL[1]
		}
	},
	'3' : {
		'0' : {
			'x' : bottomRightTL[0],
			'y' : bottomRightTL[1]
		},
		'1' : {
			'x' : bottomRightTR[0],
			'y' : bottomRightTR[1]
		},
		'2' : {
			'x' : bottomRightBR[0],
			'y' : bottomRightBR[1]
		},
		'3' : {
			'x' : bottomRightBL[0],
			'y' : bottomRightBL[1]
		}
	}
}

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

display(clean)