import sys

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

import gzip

import json

from typing import MutableMapping

from analyzer.epoint import EPoint

from analyzer.corners import Corners

class Board:

	def __init__(self,**kwargs):

		self.board=Corners(kwargs.get("board") or [])

		self.grid=Corners(kwargs.get("grid") or [])

	def isEmpty(self):

class DataFile(MutableMapping):

	"""self.data: {filename: [Board, ...}"""

	def __init__(self,filename):

		self.filename=filename

		try:

			with gzip.open(filename,mode="rt",encoding="utf8") as f:

				self._data=json.load(f,object_hook=self.deserialize)

		except OSError:

			self._data=dict()

	def save(self):

		with gzip.open(self.filename,mode="wt",encoding="utf8") as f:

			json.dump(self._data,f,default=self.serialize,indent="\t")

	def serialize(self,obj):

		if isinstance(obj,EPoint):

			return {"type": "EPoint", "val": [obj.x,obj.y]}

		elif isinstance(obj,Board):

		raise TypeError(obj)

	def deserialize(self,obj):

		type=obj.get("type")

		if type=="EPoint": return EPoint(*obj["val"])

		elif type=="Board": return Board(**obj["val"])

		return obj

	def __getitem__(self, key): return self._data[key]

	def __setitem__(self, key, val): self._data[key]=val

	def __delitem__(self, key): del self._data[key]

	def __iter__(self): return iter(self._data)

	def __len__(self): return len(self._data)

def computeBoundingBox(corners):

	x1=min(p.x for p in corners)

	x2=max(p.x for p in corners)

	y1=min(p.y for p in corners)

	y2=max(p.y for p in corners)

	return (x1,y1,x2,y2)

	def showImage(self):

		filename=self.filenames[self.k]

		self.img=Image.open(os.path.join(self.dirname,filename))

		self.root.title(filename)

		(w,h)=self.img.size

		self.photo=ImageTk.PhotoImage(self.img)

		self.canvas.delete("all")

		self.canvas.create_image(0,0,image=self.photo,anchor="nw")

		self.canvas.configure(width=w,height=h)

		self.m=0

		self.isGrid=False

		f=self.filename

		if f in self.annotations:

			self._boards=self.annotations[f]+[Board()]

		else:

			self._boards=[Board()]

		self._markCorners()

	def switchImage(self,step):

		self.save()

		n=len(self.filenames)

		self.k=(self.k+step+n)%n

		self.showImage()

	def switchCorners(self,step):

		n=len(self._boards)

		self.m=(self.m+step+n)%n

		self.isGrid=False

		self._markCorners()

	def setLetter(self,c):

		self.letter=c

		self.letterLabel.configure(text="Letter: {0}".format(c))

	def switchAsterisk(self):

		self.asterisk=not self.asterisk

		self.asteriskLabel.configure(text="*" if self.asterisk else "()")

		self._markCorners()

	def switchGrid(self):

		self.isGrid=not self.isGrid

		self._markCorners()

	def deleteBoard(self):

		self._boards.pop(self.m)

		self._flush()

		self.showImage()

	def addCorner(self,e):

		self.corners.add(e.x,e.y)

		self._flush()

		self._markCorners()

	def save(self):

		if self._dirty:

			self.annotations.save()

			self._dirty=False

	def _flush(self):

		self.annotations[self.filename]=[b for b in self._boards if not b.isEmpty()]

	def _createGUI(self):

		self.root=root=tk.Tk()

		frame=tk.Frame(root)

		frame.grid(column=0,row=0,sticky=(N,S,E,W))

		self.canvas=tk.Canvas(frame)

		self.canvas.grid(row=1,column=0)

		bar=self._createBar(frame)

		bar.grid(row=2,column=0,sticky=(E,W))

		self.showImage()

		self.canvas.bind('<1>',lambda e: self.printSample())

		self.canvas.bind('<3>',self.addCorner)

		self.canvas.bind("<Motion>",self.sample)

		root.bind("<Left>",lambda e: self.switchImage(-1))

		root.bind("<Right>",lambda e: self.switchImage(1))

		root.bind("<Up>",lambda e: self.switchCorners(-1))

		root.bind("<Down>",lambda e: self.switchCorners(1))

		root.bind("<b>",lambda e: self.setLetter("b"))

		root.bind("<e>",lambda e: self.setLetter("e"))

		root.bind("<w>",lambda e: self.setLetter("w"))

		root.bind("<Key-period>",lambda e: self.switchAsterisk())

		root.bind("<Key-comma>",lambda e: self.switchGrid())

		root.bind("<d>",lambda e: self.deleteBoard())

		root.mainloop()

	def _createBar(self,frame):

		bar=tk.Frame(frame,height=20,borderwidth=1,relief="sunken")

		self.letterLabel=tk.Label(bar,width=8,text="Letter: _")

		self.letterLabel.pack(side=LEFT)

		self.posLabel=tk.Label(bar,width=16,text="(,)")

		self.posLabel.pack(side=LEFT)

		self.colorLabel=tk.Label(bar,width=20,text="HSV: (,,)")

		self.colorLabel.pack(side=LEFT)

		self.asteriskLabel=tk.Label(bar,width=4,text="()")

		self.asteriskLabel.pack(side=LEFT)

		return bar

	def _markCorners(self):

		self.canvas.delete("mark")

		for b in self._boards:

			for corners in (b.board,b.grid):

				for c in corners:

					(x,y)=(c.x,c.y)

					self.canvas.create_oval(x-2,y-2,x+2,y+2,fill="#00ff00",tags="mark")

				if corners.is_canon():

					(a,b,c,d)=corners

					dash=tuple() if corners is self.corners else (4,4)

					self.canvas.create_line(a.x,a.y,b.x,b.y,fill="#00ff00",dash=dash,tags="mark")

					self.canvas.create_line(b.x,b.y,c.x,c.y,fill="#00ff00",dash=dash,tags="mark")

					self.canvas.create_line(c.x,c.y,d.x,d.y,fill="#00ff00",dash=dash,tags="mark")

					self.canvas.create_line(d.x,d.y,a.x,a.y,fill="#00ff00",dash=dash,tags="mark")

	@property

	def corners(self):

		return b.grid if self.isGrid else b.board

	@property

	def filename(self):

		return "." if self.asterisk else self.filenames[self.k]

if __name__=="__main__":

	dirname=sys.argv[1]

	annotations=sys.argv[2]

	s=Sampler(dirname,annotations)

import os

import sys

import re

import random

import numpy as np

import cv2 as cv

sys.path.append("..")

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

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):

		""":param img: a greyscale image as a 2D np.uint8

		:param grid: iterable of 4 EPoints, ie. Corners"""

		self.img=img

		self.grid=grid

	def transform(self):

		""":return: (img, grid), where img is a 2D np.float32 with values in (0,1),

		grid [(float) x, (float) y, ...], with x, y in (-1,1)"""

		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))

		img=np.uint8(img)

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

		grid=list(map(lambda p: list(2*p/self.SIDE-EPoint(1,1)), grid))

		return (img,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=cv.cvtColor(self.img,cv.COLOR_GRAY2BGR)

		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):

			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)

			# sample.show()

			(transformedImg,label)=sample.transform()

			# Sample(np.uint8((transformedImg+1)*128),map(lambda c: (c+EPoint(1,1))*Sample.SIDE/2,transformedGrid)).show()

			yield (transformedImg,label)

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

	)