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)