from .status import Status
from .train import train
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 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.training_params.num_epochs} epochs")
			train(layer.entity, _batch, status=status)
			_batch = layer.entity.forward(_batch)

	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.forward(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.reconstruct(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