import os
import sys
import random

import PIL.Image
import PIL.ImageDraw
import PIL

from color_sampler import DataFile


finalPosCount=5000
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)


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)