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

WORK_DIR = "../../results"
USE_OPTIMIZER = True
N_VIS = 3000
N_CASES = 1000
class LinearModel(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

class _SmallLinearModel(Model):
	def __init__(self):
		super().__init__("linear_pytest", "results")
		self.unit1 = Entity((64, 32), EntityParams(do_gaussian_visible=True),
		                    TrainingParams(learning_rate=0.005, momentum=0.9, num_epochs=100))

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

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


def test_linear():
	model = _SmallLinearModel()
	model.init(0.1)
	batch = normalize(np.random.randn(50, 64, dtype=np.float64))
	model.train(batch)
	for pattern in batch:
		recon = model.reconstruct(model.forward(pattern))
		assert recon.size == pattern.size


def main():
	model = LinearModel("linear", "results")
	model.init(0.1)
	model.load()
	training_batch = np.random.randn(N_CASES, N_VIS, dtype=np.float64)
	training_batch_norm = normalize(training_batch)
	model.train(training_batch_norm)
	model.save()
	for pattern in training_batch:
		model.reconstruct(model.forward(pattern))


if __name__ == "__main__":
	main()