import matplotlib.pyplot as plt


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

class LabelLearner(Model):
	def __init__(self, name: str, dim, work_dir: str = '.'):
		super().__init__(name, work_dir)
		self.unit1 = Entity(dim, EntityParams(do_gaussian_hidden=False), TrainingParams(learning_rate=0.01, momentum=0.9, do_rao_blackwell=True, num_epochs=1000))

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

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

if __name__ == "__main__":
	voc_size = 100
	do_train = False
	work_dir = "../../results"
	prj_root = "/home/jens/work/repos/Rbm"

	# Create Labeler
	label_w = 5
	label_h = 5
	enc = Label(voc_size, label_w*label_h, Label.EncodingType.OneHot, work_dir)

	# Load
	enc.fitter.load()

	# Train
	if do_train:
		train_labels = list(range(voc_size))
		enc.fit(train_labels)
		enc.fitter.save()

	# Encode test labels
	test_labels = np.array([1, 23, 99, 37, 55, 7, 31, 10, 19, 70])
	encoded, _ = enc.encode(test_labels)

	# Create model
	model = LabelLearner("Label-Learner", (label_w*label_h, 16), work_dir)

	# Init state
	model.init(0.01)

	# load state
	model.load()

	# Train
	model.train(encoded)

	# save state
	model.save()

	# Plot
	cmap = 'Grays'
	fig, axes = plt.subplots(3, len(encoded), figsize=(12, 3))
	for index, inp in enumerate(encoded):
		hidden = model.forward(inp)
		recon = model.backward(hidden)
		img_hidden = np.reshape(hidden, (4, 4))
		img_recon = np.reshape(recon, (label_w, label_h))
		axes[0][index].imshow(img_hidden.get(), cmap=cmap)
		axes[0][index].axis('off')
		axes[0][index].set_title(f'{test_labels[index]}')
		axes[1][index].imshow(img_recon.get(), cmap=cmap)
		axes[1][index].axis('off')
		axes[1][index].set_title(f'{test_labels[index]}')
		axes[2][index].imshow(np.reshape(inp.get(), (label_w, label_h)), cmap=cmap)
		axes[2][index].axis('off')
		axes[2][index].set_title(f'{test_labels[index]}')

	plt.show()

	print("Test: [passed]")