import os
from abc import ABC, abstractmethod

from .entity import Entity
from .train import train
from .matrix import Mat, np
from .status import Status

class Model(ABC):
	known_classes = [Entity]
	def __init__(self, name: str = "myStack", work_dir: str = "."):
		self.name = name
		self.work_dir = work_dir
		os.makedirs(self.work_dir, exist_ok=True)

	def objects(self, obj_type: type = Entity) -> list[Entity]:
		obj_list: list[type[obj_type]] = []
		for name, value in self.__dict__.items():
			if isinstance(value, obj_type):
				obj_list.append(value)
		return obj_list

	def train(self, batch: Mat, status: Status = None):
		entities = self.objects(Entity)
		_batch = np.array(batch)   # np.array() converts numpy→CuPy; np.copy() does not
		for entity in entities:
			if entity.enable_training:
				train(entity, _batch, status if status is not None else Status())
			_batch = entity.forward(_batch)

	def forward(self, x: Mat) -> Mat:
		pass

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

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

	def init(self, std: float):
		for index, entity in enumerate(self.objects(Entity)):
			entity.state.init(std=std)