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

WORK_DIR = "../../results"
USE_OPTIMIZER = True
N_VIS = 3000
N_CASES = 1000
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((N_VIS, 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((N_VIS, 1000), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False),
								TrainingParams(learning_rate=0.005, momentum=0.9, num_epochs=1000, mini_batch_size=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.randn(N_CASES, N_VIS, dtype=np.float64)

	# Normalize training data
	training_batch_norm = normalize(training_batch)

	# Train layer
	model.train(training_batch_norm)

	# 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]")