import argparse
import os
from PIL import Image
import time
import json

parser = argparse.ArgumentParser(description="""
Frames are split into segments so SMALL_FACES actually detects _really_ small faces. This requires enabling --segments on opencv-webcam-demo (so it detects/writes segment%06d files) split_and_merge_output.py also remerges these segment%06.json files into a frame%06d.json file.
""")
parser.add_argument('--frameOutput', '-o', required=True, help='directory to look for frames & json')
parser.add_argument('--segmentDir', required=True, help='directory to write segfor frames & json')

args = parser.parse_args()

if not os.path.exists(args.frameOutput):
	raise Exception(args.frameOutput + " does not exist")
if not os.path.exists(args.segmentDir):
	raise Exception(args.segmentDir + " does not exist")

def getSectionsForSize(size):
	"""
	Get sections based on image size, in a tuple format (width, height)
	"""
	return [
		( 0,				0,					int(size[0]/1.85),	int(size[1]/1.85) ),
		( int(size[0]/2.1),	0,					size[0], 			int(size[1]/1.85) ),
		( 0,				int(size[1]/2.1),	int(size[0]/1.85),	size[1]  	),
		( int(size[0]/2.1),	int(size[1]/2.1),	size[0], 			size[1]  	),
	]

nextInNr = 1
segmentCount = 4 # cannot really be changed for now
nextOutNr = segmentCount
while True:
	nextJpg = os.path.join(args.frameOutput, "frame%06d.jpg" % nextInNr)
	nextInJson = os.path.join(args.frameOutput, "frame%06d.json" % nextInNr)
	while os.path.exists(nextJpg) and not os.path.exists(nextInJson):		
		print("(SPLIT) Found: {}".format(nextJpg))
		img = Image.open(nextJpg)
		
		segments = getSectionsForSize(img.size)
		
		for i, segment in enumerate(segments):
			segmentImg = img.crop(segment)
			segmentNr = (nextInNr-1) * segmentCount + i + 1 # start at 1, increment from there
			segmentFilename = os.path.join(args.segmentDir, "segment%06d.jpg" % segmentNr)
			segmentImg.save(segmentFilename, quality=90)
			print("(SPLIT) Created: {}".format(segmentFilename))
		
		nextInNr += 1
		nextJpg = os.path.join(args.frameOutput, "frame%06d.jpg" % nextInNr)
		nextInJson = os.path.join(args.frameOutput, "frame%06d.jpg" % nextInNr)
	
	nextJson = os.path.join(args.segmentDir, "segment%06d.json" % nextOutNr)
	while os.path.exists(nextJson):
		print("(MERGE) Found: {}".format(nextJson))
		frameOutNr = (nextOutNr - 1) / segmentCount + 1 # 1,2,3,4 -> 1; 5,6,7,8 -> 2 (uses trick to divide ints: 7/4=1)
		img = Image.open( os.path.join(args.frameOutput, "frame%06d.jpg" % frameOutNr) )
		segments = getSectionsForSize(img.size)
		creationTime = img._getexif()[36867]
		
		faces = []
		times = []
		segmentImages = []
		
		for i, segment in enumerate(segments):
			segmentNr = nextOutNr - segmentCount + i + 1 # nextOutNr: 4 => 1,2,3,4
			segmentImages.append(os.path.join(args.segmentDir, "segment%06d.jpg" % segmentNr))
			with open( os.path.join(args.segmentDir, "segment%06d.json" % segmentNr) ) as fp:
				j = json.load(fp)
				
			for f in j['faces']:
				f['rect']['x'] += segment[0]
				f['rect']['y'] += segment[1]
				f['id'] = "%d%02d" % ((i+1), f['id'])
				faces.append(f)
			times.append(j['t'])
		
		data = {
			'nr': int(frameOutNr),
			't': creationTime,
			'faces': faces,
			'sectionTimes': times,
		}
		jsonOutFilename = os.path.join(args.frameOutput, "frame%06d.json" % frameOutNr)
		with open(jsonOutFilename, 'w') as fp:
			json.dump(data, fp)
			print("(MERGE) Wrote: {}".format(jsonOutFilename))
		print("(MERGE) Cleanup segments for frame {}".format(frameOutNr))
		for imgName in segmentImages:
			print("(MERGE) \t Unlink {}".format(imgName))
			os.unlink(imgName)
		
		nextOutNr += segmentCount
		nextJson = os.path.join(args.segmentDir, "segment%06d.json" % nextOutNr)
		
	time.sleep(0.2)