from .entity import Entity, EntityParams, TrainingParams

class Layer:
	def __init__(self, name: str, shape: tuple[int, int, int, int], entity_params: EntityParams, training_params: TrainingParams):
		self.name = name
		self.shape = shape
		self.entity = Entity((shape[0]*shape[1]+shape[2], shape[3]), entity_params, training_params)
		self.training_params = training_params

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

	def save(self, filename: str = None):
		self.entity.state.save(filename)

	def load(self, filename: str = None):
		self.entity.state.load(filename)