import os
from enum import Enum
from .layer import Layer

class StackType(Enum):
	Deep = "Deep",
	Rnn = "Rnn"

class StackException(Exception):
	pass

class Stack:
	def __init__(self, stack_type: StackType, name: str, work_dir: str = '.'):
		self.stack_type = stack_type
		self.name = name
		self.work_dir = work_dir
		self.layers: list[Layer] = []
		os.makedirs(self.work_dir, exist_ok=True)

	def num_layers(self):
		return len(self.layers)

	def append(self, layer: Layer):
		self.layers.append(layer)
		return self.num_layers()-1

	def remove(self, layer: Layer):
		for index, lay in enumerate(self.layers):
			if lay == layer:
				del(self.layers[index])
				return
		raise StackException(f"Layer \"{layer.name}\" does not exist")

	def from_index(self, index : int) -> Layer:
		result = self.layers[index]
		return result

	def from_name(self, name: str):
		for layer in self.layers:
			if layer.name in name:
				return layer
		return None

	def state_init(self, std: float):
		for layer in self.layers:
			layer.init(std)

	def state_save(self):
		for index, layer in enumerate(self.layers):
			filepath = os.path.join(self.work_dir, f"{self.name}-{index}-state.npz")
			layer.save(filepath)

	def state_load(self):
		for index, layer in enumerate(self.layers):
			filepath = os.path.join(self.work_dir, f"{self.name}-{index}-state.npz")
			layer.load(filepath)