from status import Status
from cd_train import cd_jens
from stack import Stack, StackType
from matrix import Mat, np

class StackDeep(Stack):
	def __init__(self, name: str, work_dir: str = '.'):
		Stack.__init__(self, StackType.Deep, name, work_dir)

	def batch_from(self, batch: Mat, from_layer_id: int = 0):
		_batch = np.copy(batch)
		for index, layer in enumerate(self.layers):
			if index == from_layer_id:
				break
			_batch = layer.entity.v_to_ph(_batch)

		return _batch

	def train(self, batch: Mat, status=Status()):
		_batch = np.copy(batch)
		for index, layer in enumerate(self.layers):
			print(f"Train layer {index} for {layer.entity.params.num_epochs} epochs")
			_batch = self.batch_from(batch, index)
			layer.entity.train(_batch, cd_func=cd_jens, status=status)


	def pass_up(self, visible: Mat, from_layer_id: int = 0):
		h = np.copy(visible)
		for layer in self.layers[from_layer_id:]:
			h = layer.entity.gibbs_v_to_h(h)
		return h

	def pass_down(self, hidden: Mat, from_layer_id: int = 0):
		v = np.copy(hidden)
		for layer in list(reversed(self.layers))[from_layer_id:]:
			v = layer.entity.gibbs_h_to_v(v)
		return v

	def pass_down_up(self, visible: Mat, from_layer_id: int = 0):
		h = self.pass_up(visible, from_layer_id)
		v = self.pass_down(h)
		return v