from .matrix import uniform, Mat, np

class RbmState:
	def __init__(self, w_hv: Mat, b_v: Mat, b_h: Mat):
		self.num_visible, self.num_hidden = w_hv.shape
		self.w_hv = w_hv
		self.b_v = b_v
		self.b_h = b_h

	@classmethod
	def from_layer_params(cls, dim: tuple[int, int]):
		w_hv = np.zeros(dim)
		b_v = np.zeros((1, dim[0]))
		b_h = np.zeros((1, dim[1]))
		obj = cls(w_hv, b_v, b_h)
		return obj

	@classmethod
	def from_file(cls, filename: str):
		obj = None
		try:
			with np.load(filename) as X:
				w_hv, b_v, b_h = [X[i] for i in ('whv', 'bv', 'bh')]
				obj = cls(w_hv, b_v, b_h)
				print(f"{filename} loaded successfully!")
		except FileNotFoundError:
			pass
		except KeyError:
			pass

		return obj

	def v_to_h(self, visible: Mat) -> Mat:
		return np.dot(visible, self.w_hv) + self.b_h

	def h_to_v(self, hidden: Mat) -> Mat:
		return np.dot(hidden, np.transpose(self.w_hv)) + self.b_v

	def load(self, filename: str):
		try:
			with np.load(filename) as X:
				w_hv, b_v, b_h = [X[i] for i in ('whv', 'bv', 'bh')]
				if w_hv.shape == self.w_hv.shape and b_v.shape == self.b_v.shape and b_h.shape == self.b_h.shape:
					self.w_hv = w_hv
					self.b_v = b_v
					self.b_h = b_h
				print(f"{filename} loaded successfully!")
		except FileNotFoundError:
			pass
		except KeyError:
			pass


	def save(self, filename: str):
		np.savez(filename, whv=self.w_hv, bv=self.b_v, bh=self.b_h)
		print(f"{filename} saved successfully!")

	def init(self, mu: float = 0.0, std: float = 1.0):
		self.w_hv = uniform(self.w_hv.shape, mu, std)
		self.b_v = uniform(self.b_v.shape, 0, 0)
		self.b_h = uniform(self.b_h.shape, 0, 0)
