import os

import sys

import re

import math

import random

import logging as log

import numpy as np

import cv2 as cv

import config as cfg

sys.path.append("..")

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

from annotations import DataFile,computeBoundingBox,Corners,EPoint,Board

from geometry import Line

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

random.seed(361)

class Stats:

	counts=[0,0,0,0]

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 rectify(self):

		x1=self.SIDE*0.1

		x2=self.SIDE*0.9

		abcd=list(map(list,self.grid))

		destPoints=[(x1,x1),(x1,x2),(x2,x2),(x2,x1)]

		abcd_=list(map(list, (EPoint(x,y)+self._createNoise() for (x,y) in destPoints)))

		m=cv.getPerspectiveTransform(np.float32(abcd),np.float32(abcd_))

		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)

		alpha=random.uniform(0,math.pi*2)

		return EPoint(math.cos(alpha)*d, math.sin(alpha)*d)

	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)

		img=cv.resize(img,(self.SIDE*2,self.SIDE*2))

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

			log.info(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:

		grade=annotations.get(f.name,[Board()])[0].grade

		Stats.counts[grade]+=1

		if not Board.UNSET<grade<=Board.GOOD: continue

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

			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)

	log.info("clear images: %s",Stats.counts[1])

	log.info("good images: %s",Stats.counts[2])

	log.info("poor images: %s",Stats.counts[3])

	log.info("unset images: %s",Stats.counts[0])

	log.info("total: %s",sum(Stats.counts))

	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

	)

import argparse

import logging as log

import numpy as np

from keras.models import load_model

import keras.metrics

from prepare_data import loadDataset,Sample

from analyzer.epoint import EPoint

from analyzer.corners import Corners

from k_util import averageDistance,generateData

import config as cfg

keras.losses.averageDistance=averageDistance

keras.metrics.averageDistance=averageDistance

parser=argparse.ArgumentParser()

parser.add_argument("model")

parser.add_argument("data")

args=parser.parse_args()

model=load_model(args.model)

model.summary()

log.info("loading data...")

with np.load(args.data) as data:

	testImages=(np.float32(data["testImages"])/128-1).reshape((-1,224,224,1))

log.info("done")

log.info(model.evaluate(testImages,testLabels))

for img in testImages:

	label=model.predict(np.reshape(img,(1,224,224,1)))

	print(label)

	points=[]

		points.append(EPoint((label[0][i][0]+1)*112,(label[0][i][1]+1)*112))

	sample=Sample(np.uint8((img+1)*128),corners)

	sample.show()

import os

import math

from time import time

import argparse

import logging as log

import numpy as np

from keras.layers import Conv2D,Dropout,Dense,Flatten,MaxPooling2D,BatchNormalization,GlobalAveragePooling2D,Reshape

from keras.models import Sequential,load_model

from keras.callbacks import TensorBoard,ModelCheckpoint

import keras.metrics

import config as cfg

from k_util import averageDistance,generateData

keras.losses.averageDistance=averageDistance

keras.metrics.averageDistance=averageDistance

parser=argparse.ArgumentParser()

parser.add_argument("data")

parser.add_argument("--load_model")

parser.add_argument("--save_model",default="/tmp/gogo-{0:03}.h5")

parser.add_argument("--log_dir",default="/tmp/tflogs")

parser.add_argument("--epochs",type=int,default=100)

parser.add_argument("--initial_epoch",type=int,default=0)

args=parser.parse_args()

def createFullyConnected():

	model=Sequential([

		Flatten(input_shape=(224,224)),

		Dense(128, activation="relu"),

		Dropout(0.1),

		Dense(64, activation="relu"),

		Dense(8)

	])

	model.compile(

		optimizer='adam',

		loss='mse',

		metrics=['mae','accuracy']

	)

	return model

def createCNN():

	model=Sequential()

	model.add(BatchNormalization(input_shape=(224,224,1)))

	model.add(Conv2D(24,(5,5),padding="same",kernel_initializer="he_normal",activation="relu",input_shape=(224,224,1),data_format="channels_last"))

	model.add(MaxPooling2D(pool_size=(2,2),strides=(2,2),padding="valid"))

	model.add(Conv2D(36,(5,5),activation="relu"))

	model.add(MaxPooling2D(pool_size=(2,2),strides=(2,2),padding="valid"))

	model.add(Conv2D(48,(5,5),activation="relu"))

	model.add(MaxPooling2D(pool_size=(2,2),strides=(2,2),padding="valid"))

	model.add(Conv2D(64,(3,3),activation="relu"))

	model.add(MaxPooling2D(pool_size=(2,2),strides=(2,2),padding="valid"))

	model.add(Conv2D(64,(3,3),activation="relu"))

	model.add(GlobalAveragePooling2D())

	model.add(Dense(500,activation="relu"))

	model.add(Dense(90,activation="relu"))

	return model

if args.load_model:

	model=load_model(args.load_model)

else:

	model=createCNN()

	model.summary()

log.info("loading data...")

with np.load(args.data) as data:

	trainImages=(np.float32(data["trainImages"])/128-1).reshape((-1,224,224,1))

	testImages=(np.float32(data["testImages"])/128-1).reshape((-1,224,224,1))

log.info("done")

n=len(trainImages)

k=round(n*0.9)

n_=n-k

(trainImages,valImages)=(np.float32(trainImages[:k]),np.float32(trainImages[k:]))

(trainLabels,valLabels)=(np.float32(trainLabels[:k]),np.float32(trainLabels[k:]))

tensorboard=TensorBoard(log_dir=os.path.join(args.log_dir,"{}".format(time())))

checkpoint=ModelCheckpoint(args.save_model,monitor="val_loss",period=10)

model.fit_generator(

	generateData(trainImages,trainLabels,batch_size=20),

	epochs=args.epochs,

	initial_epoch=args.initial_epoch,

	steps_per_epoch=math.ceil(n_/20),

	validation_data=generateData(valImages,valLabels,batch_size=20),

	validation_steps=math.ceil(k/20),

	callbacks=[tensorboard,checkpoint]

)

print(model.evaluate(testImages,testLabels))