import json
from .stack import StackType
from .layer import Layer
from .entity import EntityParams, TrainingParams
from .stack_deep import StackDeep
from .stack_rnn import StackRnn

class StackFactory:
	@classmethod
	def from_dict(cls, project: dict, work_dir: str = ".") -> StackDeep|StackRnn|None:
		name = project["stack"]["name"]
		layers = project["stack"]["layers"]

		try:
			stack_type = StackType[project["stack"]["type_string"]]
		except KeyError:
			stack_type = StackType.Deep

		obj = None
		if stack_type == StackType.Deep:
			obj = StackDeep(name, work_dir)
		if stack_type == StackType.Rnn:
			obj = StackRnn(name, work_dir)

		if obj is None:
			return obj

		for layer in layers:
			layer_id = layer["id"]
			layer_name = layer["name"]
			num_visible_x = layer["numVisibleX"]
			num_visible_y = layer["numVisibleY"]
			num_hidden = layer["numHidden"]
			try:
				num_context = layer["numContext"]
			except KeyError:
				num_context = 0

			# Determine version by existence of keys
			params = layer["rbm"]["params"]

			# signal that doSampleBatch will be ignored
			if params["doSampleBatch"] == 1:
				raise Exception(f"Error: Parameter \"doSampleBatch\" will be ignored. To continue, set \"doSampleBatch = False\"")

			# import params
			entity_params = EntityParams.from_dict(params)
			training_params = TrainingParams.from_dict(params)

			# Create layer
			layer_obj = Layer(f"{layer_name}-{layer_id}", (num_visible_x, num_visible_y, num_context, num_hidden), entity_params, training_params)

			# Add layer to stack
			obj.append(layer_obj)

		return obj

	@classmethod
	def from_file(cls, filename: str, work_dir: str = ".") -> StackDeep|StackRnn:
		obj = None
		with open(filename, "r") as fp:
			prj = json.load(fp)
			obj = StackFactory.from_dict(prj, work_dir)

		return obj


if __name__ == "__main__":
	stack = StackFactory.from_file("/home/jens/work/repos/Rbm/test.prj", work_dir="../../results")

	print("Test: [passed]")
