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

class TestModel(Model):
	def __init__(self, name: str, work_dir: str = '.'):
		super().__init__(name, work_dir)
		self.unit1 = Entity((1024, 333), EntityParams(), TrainingParams(learning_rate=0.1, momentum=0.9, do_rao_blackwell=True, num_epochs=1000))
		self.unit2 = Entity((333, 64), EntityParams(), TrainingParams(learning_rate=0.1, momentum=0.9, do_rao_blackwell=True, num_epochs=1000))
		self.unit3 = Entity((64, 128), EntityParams(), TrainingParams(learning_rate=0.1, momentum=0.9, do_rao_blackwell=True, num_epochs=1000))
		self.unit4 = Entity((128, 128), EntityParams(), TrainingParams(learning_rate=0.1, momentum=0.9, do_rao_blackwell=True, num_epochs=1000))

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

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

if __name__ == "__main__":
	# Create model
	model = TestModel("TestModel", "results")

	# Init state
	model.init(0.1)

	# load state
	model.load()

	# create batch
	batch = (np.random.rand(64, 1024) > 0.5).astype(np.float64)

	# Train
	model.train(batch)

	# save state
	model.save()

	for index, inp in enumerate(batch):
		out = model.backward(model.forward(inp))[0]
		print(f"- Pattern {index} -------------------------")
		print(f"Input  : {inp}")
		print(f"Output : {(out > 0.9).astype(np.float64)}")
		print(f"Error  : {np.mean((out - inp)**2):0.3f}")