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
from rbm.train import train, Status
from label.label import Label

class TestModel(Model):
	def __init__(self, name: str, work_dir: str = '.'):
		super().__init__(name, work_dir)
#		self.unit_image = Entity((96*96, 16), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False, num_gibbs_samples=3), TrainingParams(learning_rate=0.00001, momentum=0.9, do_rao_blackwell=False, num_epochs=1000))
		self.unit_image = Entity((96 * 96, 16), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False))
		self.unit_combine = Entity((32, 64), EntityParams(), TrainingParams(learning_rate=0.01, momentum=0.9, do_rao_blackwell=True, num_epochs=1000))

	def forward2(self, images: Mat, labels: Mat):
		h = self.unit_image.forward(images)
		h = np.ndarray.flatten(h)
		v_combine = np.concat((h, labels), axis=0)
		h_res = self.unit_combine.forward(v_combine)
		return h_res

	def reconstruct(self, h: Mat):
		v_combine = self.unit_combine.reconstruct(h)
		v_image = self.unit_image.reconstruct(np.transpose(np.transpose(v_combine)[0:self.unit_image.shape[1]]))
		v_label = np.transpose(np.transpose(v_combine)[self.unit_image.shape[1]:self.unit_image.shape[1]+self.unit_image.shape[1]])
		return np.ndarray.flatten(v_image), np.ndarray.flatten(v_label)

	def train2(self, images: Mat, labels: Mat):
		train(self.unit_image, images, Status())
		h11 = self.unit_image.forward(images)
		x2 = np.concat((h11, labels), axis=1)
		train(self.unit_combine, x2, Status())

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

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

	# Init state
	model.init(0.01)

	# 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
	# Create Labeler
	enc = Label(20, 16, Label.EncodingType.OneHot, work_dir)

	# Load
	enc.fitter.load()

	# Encode test labels
	train_labels = np.array([1, 2, 3, 4, 5])
	enc.fit(train_labels)

	encoded, _ = enc.encode(train_labels)

	model.train2(train_batch, encoded)

	# save state
	model.save()

	cmap = 'Grays'
	fig, axes = plt.subplots(3, len(train_batch), figsize=(12, 3))
	for index, image in enumerate(train_batch):
		label_zero = np.zeros(shape=16)
		image_zero = np.zeros(shape=96*96)
		image_reconst, labels_reconst = model.reconstruct(model.forward2(image, label_zero))
#		image_reconst, labels_reconst = model.reconstruct(model.forward2(image_zero, encoded[index, :]))
		img = 1-2*(image_reconst + 0.5)
		img = np.reshape(img, (96, 96))
		axes[0][index].imshow(img, cmap=cmap)
		axes[0][index].axis('off')
		axes[1][index].imshow(np.reshape(labels_reconst, (4,4)), cmap=cmap)
		axes[1][index].axis('off')
		axes[2][index].imshow(np.reshape(encoded[index, :], (4,4)), cmap=cmap)
		axes[2][index].axis('off')

	plt.show()