import numpy as np
from params import RbmParams
from state import RbmState
from status import Status
from cd_train import cd_jens
from entity import Entity

class Layer:
	def __init__(self, name: str, shape: tuple[int, int, int, int], params: RbmParams):
		self.name = name
		self.shape = shape
		self.entity = Entity((shape[0]*shape[1]+shape[2], shape[3]), params)
		self.state_filename = f"{self.name}_state.npz"

	def init(self, std: float):
		self.entity.state.init(mu=0, std=std)

	def save(self, filename: str = None):
		if filename is None:
			filename = self.state_filename
		self.entity.state.to_file(filename)

	def load(self, filename: str = None):
		if filename is None:
			filename = self.state_filename
		state = RbmState.from_file(filename)
		if state is not None:
			self.entity.state = state

def xor():
	# Create params
	params = RbmParams()
	params.do_rao_blackwell = True
	params.num_gibbs_samples = 3

	# Create layer
	layer = Layer("Layer_0", (3, 1, 0, 16), params)

	# Init weights
	layer.init(0.01)

	# Load weights (if exists)
	layer.load()

	# Prepare training data
	training_batch = np.array([[0,1,1], [0,0,0], [1,1,0], [1,0,1]], dtype=np.float64)

	# Train layer
	layer.entity.train(training_batch, cd_jens, Status())

	# Save weights
	layer.save()

	# Test with test data
	test_batch = np.array([[0,0,0], [0,1,0], [1,0,0], [1,1,0]], dtype=np.float64)
	for pattern in test_batch:
		h = layer.entity.gibbs_v_to_h(pattern)
		v = layer.entity.gibbs_h_to_v(h)
		print(f"P{pattern} : {v}")


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