from rbm.entity import Entity
from rbm.train import cd_binary_binary, cd_binary_gaussian, cd_gaussian_binary, cd_gaussian_gaussian
from rbm.matrix import Mat, rms_error_accu, np

class LossFunction:
	def __init__(self, entity_type: Entity.Type):
		self.func = None
		if entity_type == Entity.Type.GB_RBM:
			self.func = cd_gaussian_binary
		if entity_type == Entity.Type.BB_RBM:
			self.func = cd_binary_binary
		if entity_type == Entity.Type.GG_RBM:
			self.func = cd_gaussian_gaussian
		if entity_type == Entity.Type.BG_RBM:
			self.func = cd_binary_gaussian

	def __call__(self, entity: Entity, data: Mat):
		# Contrastive divergence learning: calculate gradients
		dwhv, dbv, dbh = self.func(entity, data)

		return dwhv, dbv, dbh

class Optimizer:
	def __init__(self, entity: Entity):
		self.loss = LossFunction(entity.type)
		self.entity = entity

	def zero_grad(self):
		self.entity.grad_zero()

	def step(self, data: Mat):
		params = self.entity.training_params

		# Call loss function
		dwhv, dbv, dbh = self.loss(self.entity, data)

		# Adjust weight and biases
		grad = self.entity.grad_compute(dbv, dbh, dwhv, learning_rate=params.learning_rate / data.shape[0],
								   momentum=params.momentum, weight_decay=params.weight_decay)
		# Adjust weights
		self.entity.state_adjust(grad)