import os

import sys

import re

import random

import itertools

import numpy as np

import cv2 as cv

sys.path.append("..")

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

from annotations import DataFile,computeBoundingBox,Corners

from geometry import Line

from kerokero.transformation_matrices import getIdentity,getRotation,getTranslation,getScale,getMirroring,getProjection

random.seed(361)

class Sample:

	SIDE=224

	def __init__(self,img,grid):

		self.img=img

		self.grid=grid

	def transform(self):

		center=self._getCenter()

		m=getIdentity()

		t1=getTranslation(-center.x,-center.y)

		proj=getProjection()

		rot=getRotation()

		mir=getMirroring()

		for mi in [t1,mir,proj,rot]:

			m=np.matmul(mi,m)

		m=np.matmul(self._computeCrop(m),m)

		img=cv.warpPerspective(self.img,m,(self.SIDE,self.SIDE))

		grid=Corners(c.transform(m) for c in self.grid)

		return (img,list(itertools.chain.from_iterable(grid)))

	def _getCenter(self):

		(a,b,c,d)=self.grid

		p=Line.fromPoints(a,c)

		q=Line.fromPoints(b,d)

		return p.intersect(q)

	def _computeCrop(self,m):

		grid=Corners(c.transform(m) for c in self.grid)

		(x1,y1,x2,y2)=computeBoundingBox(grid)

		(wg,hg)=(x2-x1,y2-y1)

		(left,top,right,bottom)=[random.uniform(0.05,0.2) for i in range(4)]

		t2=getTranslation(left*wg-x1, top*hg-y1)

		scale=getScale(self.SIDE/(wg*(1+left+right)), self.SIDE/(hg*(1+top+bottom)))

		return np.matmul(scale,t2)

	def show(self):

		img=np.copy(self.img)

		for c in self.grid:

			cv.circle(img,(int(c.x),int(c.y)),3,[0,255,0],-1)

		show(img)

def traverseDirs(root):

	stack=[root]

	while len(stack)>0:

		d=stack.pop()

		contents=sorted(os.scandir(d),key=lambda f: f.name,reverse=True)

		if any(f.name=="annotations.json.gz" for f in contents):

			print(d)

			yield d

		for f in contents:

			if f.is_dir(): stack.append(f.path)

def harvestDir(path):

	annotations=DataFile(os.path.join(path,"annotations.json.gz"))

	imgFilter=lambda f: f.is_file() and re.match(r".*\.(jpg|jpeg|png|gif)$", f.name.lower())

	files=sorted(filter(imgFilter,os.scandir(path)),key=lambda f: f.name)

	boards=annotations["."]

	for f in files:

		img=cv.imread(f.path)

		img=cv.cvtColor(img,cv.COLOR_BGR2GRAY)

		for b in boards:

			sample=Sample(img,b.grid)

def loadDataset(root):

	testRatio=0.1

	trainRatio=1-testRatio

	images=[]

	labels=[]

	for d in traverseDirs(root):

		for (img,label) in harvestDir(d):

			images.append(img)

			labels.append(label)

	n=len(images)

	keys=list(range(n))

	random.shuffle(keys)

	images=[images[k] for k in keys]

	labels=[labels[k] for k in keys]

	m=int(n*trainRatio)

	return (

		(np.uint8(images[:m]),np.float32(labels[:m])),

		(np.uint8(images[m:]),np.float32(labels[m:]))

	)

def show(img,filename="x"):

	cv.imshow(filename,img)

	cv.waitKey(0)

	cv.destroyAllWindows()

if __name__=="__main__":

	((trainImages,trainLabels),(testImages,testLabels))=loadDataset(sys.argv[1])

	np.savez_compressed(

		sys.argv[2],

		trainImages=trainImages,

		trainLabels=trainLabels,

		testImages=testImages,

		testLabels=testLabels

	)