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

WORK_DIR = "../../results"
USE_OPTIMIZER = True

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

		if do_gaussian_hidden:
			# Hidden gaussian
			self.unit1 = Entity((3, 64), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=True),
								TrainingParams(learning_rate=0.01, momentum=0.9, num_epochs=1000))
		else:
			# Hidden binary
			self.unit1 = Entity((3, 64), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False),
								TrainingParams(learning_rate=0.001, momentum=0.9, num_epochs=1000))

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

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

def linear():
	work_dir = "results"
	prj_name = "linear"
	prj_root = "/home/jens/work/repos/Rbm"

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

	# Init weights
	model.init(0.1)

	# Load weights (if exists)
	model.load()

	# Prepare training data
	training_batch = (np.random.rand(150, 3, dtype=np.float64) - 0.5)

	# Normalize training data
	mean_training_batch = np.reshape(np.repeat(np.mean(training_batch, axis=1), 3, axis=0), training_batch.shape)
	var_training_batch = np.reshape(np.repeat(np.std(training_batch, axis=1), 3, axis=0), training_batch.shape)
	training_batch = (training_batch - mean_training_batch) / var_training_batch

	# Train layer
	model.train(training_batch)

	# Save weights
	model.save()

	# Test with test data
	test_batch = training_batch
	for pattern in test_batch:
		h = model.forward(pattern)
		v = model.reconstruct(h)
		print(f"P{pattern} : {v}")

if __name__ == "__main__":
	linear()
	print("Test: [passed]")