import sys
sys.path.append("../src")

import os
import random

import PIL.Image
import PIL.ImageDraw
import PIL

from diagram import createDiagram
from annotations import DataFile

finalPosCount=5000
empty=True
perspective=False
random.seed(361)
inputDir=sys.argv[1]
outputDir=sys.argv[2]
negDir=sys.argv[3]
negFiles=os.listdir(negDir)
negFiles.sort()

annotations=DataFile("annotations.json.gz")
posCount=len(annotations)
multiple=int(round(finalPosCount/posCount+0.5))


def pickRandom():
	file=random.choice(negFiles)
	return PIL.Image.open(os.path.join(negDir,file))


def cropBackground(bgImg,w,h):
	(bgw,bgh)=bgImg.size
	scale=min(bgw/w,bgh/h)
	if scale<1: # minimal upscale
		bgImg=bgImg.resize((int(round(bgw/scale+0.5)),int(round(bgh/scale+0.5))),PIL.Image.BICUBIC)
		scale=1
	else: # random downscale
		r=1+random.random()*(scale-1)
		bgImg=bgImg.resize((int(bgw/r+0.5),int(bgh/r+0.5)),PIL.Image.BICUBIC)
	(bgw,bgh)=bgImg.size
	wRange=bgw-w
	hRange=bgh-h
	x=random.randint(0,wRange)
	y=random.randint(0,hRange)
	bg=bgImg.crop((x,y,x+w,y+h))
	if random.random()<0.5:
		bg=bg.transpose(PIL.Image.FLIP_LEFT_RIGHT)
	return bg


def extract(image,cornerList,i,filename):
	img=image.copy()
	x1=min(p.x for p in cornerList)
	x2=max(p.x for p in cornerList)
	y1=min(p.y for p in cornerList)
	y2=max(p.y for p in cornerList)
	w=x2-x1
	h=y2-y1

	bg=cropBackground(pickRandom(),w,h)

	mask=PIL.Image.new("1",bg.size,255)
	d=PIL.ImageDraw.Draw(mask)
	d.polygon([(p.x-x1,p.y-y1) for p in cornerList], fill=0, outline=0)
	img.paste(bg,(x1,y1),mask=mask)

	relevantArea=img.crop((x1,y1,x2,y2))
	if random.random()<0.5:
		relevantArea=relevantArea.transpose(PIL.Image.FLIP_LEFT_RIGHT)
	outputPath=os.path.join(outputDir,filename.replace(".","-{0}.".format(i)))
	relevantArea.save(outputPath)
	print(outputPath,1,0,0,x2-x1,y2-y1)


def generate(background,i,outputDir):
	bg=background.copy()
	img=createDiagram([[]])
	maxSize=min(*img.size,*bg.size)
	size=random.randint(min(80,maxSize),maxSize)
	img=img.resize((size,size),PIL.Image.BICUBIC)
	wRange=bg.size[0]-img.size[0]
	hRange=bg.size[1]-img.size[1]
	x=random.randint(0,wRange)
	y=random.randint(0,hRange)
	bg.paste(img,(x,y))
	outputPath=os.path.join(outputDir,"{0:04}.jpg".format(i))
	bg.save(outputPath)
	rect=[
		max(x-5,0),
		max(y-5,0),
		img.size[0]+min(10,bg.size[0]-x),
		img.size[1]+min(10,bg.size[1]-y)
	]
	print(outputPath,1,*rect,flush=True)

# for (filename,cornerList) in annotations.items():
# 	path=os.path.join(inputDir, filename)
# 	if not os.path.isfile(path):
# 		print("{0} not found, ignored".format(path))
# 		continue
#
# 	img=PIL.Image.open(path)
# 	for i in range(multiple):
# 		extract(img,cornerList,i,filename)

for i in range(finalPosCount):
	generate(pickRandom(),i,outputDir)