from .state import RbmState
from .matrix import prob, Mat
from enum import Enum

class TrainingParams:
	def __init__(self,
				 learning_rate: float = 0.1,
				 momentum: float = 0.5,
				 weight_decay: float = 0.0,
				 num_epochs: int = 1000,
				 num_gibbs_samples: int = 1,
				 mini_batch_size: int = 0,
				 do_rao_blackwell: bool = False,
				 do_gibbs_sample_visible: bool = False,
				 do_gibbs_sample_hidden: bool = False,
				 do_batch_sample: bool = False
				 ):
		# Training parameters
		self.learning_rate = learning_rate
		self.momentum = momentum
		self.weight_decay = weight_decay
		self.num_epochs = num_epochs
		self.num_gibbs_samples = num_gibbs_samples
		self.mini_batch_size = mini_batch_size
		self.do_rao_blackwell = do_rao_blackwell
		self.do_gibbs_sample_visible = do_gibbs_sample_visible
		self.do_gibbs_sample_hidden = do_gibbs_sample_hidden
		self.do_batch_sample = do_batch_sample

	@classmethod
	def from_dict(cls, params: dict):
		obj = TrainingParams()
		obj.learning_rate = params["learningRate"]
		obj.momentum = params["momentum"]
		obj.weight_decay = params["weightDecay"]
		obj.num_epochs = params["numEpochs"]
		obj.num_gibbs_samples = params["numGibbs"]
		obj.mini_batch_size = params["miniBatchSize"]
		obj.do_rao_blackwell = params["doRaoBlackwell"]
		obj.do_gibbs_sample_visible = params["gibbsDoSampleVisible"]
		obj.do_gibbs_sample_hidden = params["gibbsDoSampleHidden"]
		obj.do_batch_sample = params["doSampleBatch"]

		return obj

class EntityParams:
	def __init__(self, do_gaussian_visible: bool = False, do_gaussian_hidden: bool = False, num_gibbs_samples: int = 1):
		# Entity parameters
		self.do_gaussian_visible = do_gaussian_visible
		self.do_gaussian_hidden = do_gaussian_hidden
		self.num_gibbs_samples = num_gibbs_samples

	@classmethod
	def from_dict(cls, params: dict):
		obj = EntityParams()
		if "doGaussianVisible" in params:
			obj.do_gaussian_visible = params["doGaussianVisible"]
		if "doGaussianHidden" in params:
			obj.do_gaussian_hidden = params["doGaussianHidden"]
		if "num_gibbs_samples" in params:
			obj.num_gibbs_samples = params["num_gibbs_samples"]

		return obj

class Entity:
	class Type(Enum):
		BB_RBM = "BB-RBM"
		BG_RBM = "BG-RBM"
		GB_RBM = "GB-RBM"
		GG_RBM = "GG-RBM"

	def __init__(self, shape: tuple[int, int], params: EntityParams, training_params: TrainingParams|None = None):
		self.shape = shape
		self.params = params
		self.training_params = training_params
		self.state = RbmState.from_layer_params(shape)
		self.grad = RbmState.from_layer_params(shape)
		self.type = None
		if params.do_gaussian_visible:
			if params.do_gaussian_hidden:
				self.type = Entity.Type.GG_RBM
			else:
				self.type = Entity.Type.GB_RBM
		else:
			if params.do_gaussian_hidden:
				self.type = Entity.Type.BG_RBM
			else:
				self.type = Entity.Type.BB_RBM


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

	def grad_zero(self):
		self.grad = RbmState.from_layer_params(self.shape)

	def grad_compute(self, d_bv: Mat, d_bh: Mat, d_whv: Mat, learning_rate: float, momentum: float, weight_decay: float):
		# Compute gradient
		self.grad.b_v = (momentum * self.grad.b_v + learning_rate * d_bv)
		self.grad.b_h = (momentum * self.grad.b_h + learning_rate * d_bh)
		self.grad.w_hv = (momentum * self.grad.w_hv + learning_rate * d_whv - learning_rate * weight_decay * self.state.w_hv)

		return self.grad

	def state_adjust(self, grad: RbmState):
		# Adjust state
		self.state.b_v += grad.b_v
		self.state.b_h += grad.b_h
		self.state.w_hv += grad.w_hv

	def forward(self, v: Mat, num_gibbs: int = 0) -> Mat:
		num_gibbs = self.params.num_gibbs_samples if num_gibbs == 0 else num_gibbs
		h = self._v_to_ph(v)
		for i in range(num_gibbs-1):
			h = self._h_to_pv(h)
			h = self._v_to_ph(h)

		return h

	def reconstruct(self, h: Mat, num_gibbs: int = 0) -> Mat:
		num_gibbs = self.params.num_gibbs_samples if num_gibbs == 0 else num_gibbs
		v = self._h_to_pv(h)
		for i in range(num_gibbs-1):
			v = self._v_to_ph(v)
			v = self._h_to_pv(v)

		return v

	def _v_to_ph(self, v: Mat) -> Mat:
		state = self.state.v_to_h(v)
		if self.params.do_gaussian_hidden:
			return state

		return prob(state)


	def _h_to_pv(self, h: Mat) -> Mat:
		state = self.state.h_to_v(h)
		if self.params.do_gaussian_visible:
			return state

		return prob(state)

	def h_given_v(self, v: Mat) -> Mat:
		state = self.state.v_to_h(v)
		return state

	def v_given_h(self, h: Mat) -> Mat:
		state = self.state.h_to_v(h)
		return state

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