import matplotlib.pyplot as plt

from rbm.matrix import np
from label.label import Label

if __name__ == "__main__":
	voc_size = 100
	do_train = True
	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.Binary, 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, binary_labels = enc.encode(test_labels)

	# Decode
	decoded_soft = enc.decode(encoded)
	decode_hard = (decoded_soft > 0.90).astype(int)

	print(f"Soft Error: {np.sum((decoded_soft - binary_labels) ** 2)}")
	print(f"Hard Error: {np.sum((decode_hard - binary_labels) ** 2)}")

	fig, axes = plt.subplots(1, len(encoded), figsize=(12, 3))
	for index, inp in enumerate(encoded):
		img = np.reshape(inp, (label_w, label_h))
		axes[index].imshow(img)
		axes[index].axis('off')
		axes[index].set_title(f'{test_labels[index]}')

	plt.show()

	print("Test: [passed]")