import os.path
import cv2 as cv
from argparse import ArgumentParser
from stack.factory import StackFactory
from rbm.status import Status
from stack.deep import StackDeep
from rbm.matrix import Mat, np, convert, read_armadillo
from rbm.entity import Entity

def cv_show(name: str, vec: Mat, shape):
	img = cv.Mat(convert(np.resize(vec, shape)))
	img_n = cv.normalize(src=img, dst=None, alpha=255, beta=0, norm_type=cv.NORM_MINMAX, dtype=cv.CV_8U)
	cv.imshow(f"{name}", img_n)

class MyStatus(Status):
	def __init__(self, _stack: StackDeep, _batch: Mat):
		Status.__init__(self, update_interval=10)
		self.stack = _stack
		self.batch = _batch
		self.index = 0

	def on_report(self, entity: Entity, status: dict) -> bool:
		do_continue = True
		# print status values
		Status.print_status(entity, status)
		# Shape of training vector
		shape = self.stack.from_index(0).shape[0:2] + (1,)

		# User input
		max_index = self.batch.shape[0] - 1
		key = cv.waitKeyEx(1)
		if key == ord('q'):
			do_continue = False
		if key == ord('-'):
			self.index = max(0, self.index-1)
		if key == ord('+'):
			self.index = min(max_index, self.index+1)

		# Show reconstruction
		img = self.stack.pass_down_up(self.batch[self.index,:])
		cv_show("img", img, shape)

		return do_continue


def main(prj_name: str = "test"):
	work_dir = "../../results"
	prj_root = "/home/jens/work/repos/Rbm"
	prj_path = os.path.join(prj_root, f"{prj_name}.prj")
	# Create stack from project file
	stack = StackFactory.from_file(prj_path, work_dir=work_dir)

	# Init state
	stack.state_init(0.01)

	# Load state
	stack.state_load()

	# Load train data
	training_data = read_armadillo(os.path.join(prj_root, f"{prj_name}.training.dat"))
	try:
		test_data = read_armadillo(os.path.join(prj_root, f"{prj_name}.test.dat"))
	except FileNotFoundError:
		test_data = training_data

	# Prepare status listener
	my_status = MyStatus(stack, test_data)

	# Train
	stack.train(training_data, status=my_status)

	# Save state
	stack.state_save()

if __name__ == "__main__":
	ap = ArgumentParser()
	ap.add_argument("name", type=str,
		default='default',
		help="Name of project")

	args = ap.parse_args()
	var_args = vars(args)

	main(var_args["name"])
	cv.destroyAllWindows()

	print("Test: [passed]")