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

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

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

		return obj

class Entity:
	def __init__(self, shape: tuple[int, int], params: EntityParams):
		self.shape = shape
		self.params = params
		self.state = RbmState.from_layer_params(shape)

	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 forward(self, v: Mat) -> Mat:
		h = self.v_to_ph(v)
		for i in range(self.params.num_gibbs_samples-1):
			h = self.h_to_pv(h)
			h = self.v_to_ph(h)

		return h

	def reconstruct(self, h: Mat) -> Mat:
		v = self.h_to_pv(h)
		for i in range(self.params.num_gibbs_samples-1):
			v = self.v_to_ph(v)
			v = self.h_to_pv(v)

		return v

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