import os
import matplotlib.pyplot as plt

from model.model import Model
from rbm.entity import Entity, EntityParams, TrainingParams
from rbm.matrix import Mat, np, read_armadillo

DO_HIDDEN_GAUSSIAN = True

class TestModel(Model):
	def __init__(self, name: str, work_dir: str = '.'):
		super().__init__(name, work_dir)

		if DO_HIDDEN_GAUSSIAN:
			# Hidden gaussian
			self.unit1 = Entity((96*96, 16), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=True),
								TrainingParams(learning_rate=0.00001, momentum=0.9, num_epochs=1000))
		else:
			# Hidden binary
			self.unit1 = Entity((96 * 96, 333), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False),
								TrainingParams(learning_rate=0.001, momentum=0.9, num_epochs=1000, l2_lambda=0.001))

	def forward(self, x: Mat):
		x = self.unit1.forward(x)
		return x

	def reconstruct(self, x: Mat):
		x = self.unit1.reconstruct(x)
		return x

if __name__ == "__main__":
	work_dir = "results"
	prj_name = "norb_small_16h_v2"
	prj_root = "/home/jens/work/repos/Rbm"

	# Create model
	model = TestModel(prj_name, "results")

	# Init state
	if DO_HIDDEN_GAUSSIAN:
		model.init(0.01)
	else:
		model.init(0.1)

	# load state
#	model.load()

	# Load train data
	train_batch = read_armadillo(os.path.join(prj_root, f"{prj_name}.training.dat"))

	# Load test data
	test_batch = read_armadillo(os.path.join(prj_root, f"{prj_name}.test.dat"))

	# Train
	model.train(train_batch)

	# save state
	model.save()

	fig, axes = plt.subplots(1, len(test_batch), figsize=(12, 3))
	for index, inp in enumerate(test_batch):
		out_normalized = model.reconstruct(model.forward(inp))
		img = 2*(out_normalized + 0.5)
		img = np.reshape(img, (96, 96))
		axes[index].imshow(np.asnumpy(img))
		axes[index].axis('off')

	plt.show()