import os
import sys
import re
import random

import numpy as np
import cv2 as cv

sys.path.append("../exp")
from annotations import DataFile,computeBoundingBox,Corners
from geometry import Line
from keras.transformation_matrices import getIdentity,getRotation,getTranslation,getScale,getMirroring,getProjection

random.seed(361)


class Sample:
	SIDE=256

	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)
		Sample(img,grid).show()

	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)
		for b in boards:
			sample=Sample(img,b.grid)
			sample.transform()


def show(img,filename="x"):
	cv.imshow(filename,img)
	cv.waitKey(0)
	cv.destroyAllWindows()


if __name__=="__main__":
	root=sys.argv[1]
	for d in traverseDirs(root):
		harvestDir(d)