import tensorflow as tf
from tensorflow.keras.layers import Input, Multiply, Add
from tensorflow.keras.models import Model

from .stnn_layers import TTL, SoftmaxEmbeddingLayer


def build_stnn(config):
	"""
	Constructs a Stacked Tensorial Neural Network (STNN) as a TensorFlow model based on
	the provided configuration dictionary.

	Args:
		config (dict): Configuration dictionary for the STNN model. Must contain the following entries:
						- 'K' (int): Number of tensor networks to be stacked
						- 'd' (int): Number of dense layers in the model's SoftmaxEmbeddingLayer.
						- 'nx1', 'nx2', 'nx3' (int): Dimensions of the finite-difference grid
						- All other required entries for the TTL class, not already listed above.

	Returns:
		tf.keras.Model: The constructed STNN model.

	Raises:
		ValueError: If the config dictionary does not contain positive integers 'K', 'd', 'nx2', 'nx3';
					also if config['nx3'] is not divisible by 2.
	"""
	required_keys = ['nx1', 'nx2', 'nx3', 'K', 'd', 'shape1','shape2','ranks','W']
	missing_keys = [key for key in required_keys if key not in config]
	if missing_keys:
		raise KeyError(f"Missing keys in config: {', '.join(missing_keys)}")

	for key in ['nx1', 'nx2', 'nx3', 'K', 'd']:
		if not isinstance(config[key], int):
			raise TypeError(f"{key} must be an integer.")

	for key in ['nx1', 'nx2', 'nx3', 'K', 'd']:
		if config[key] <= 0:
			raise ValueError(f"{key} must be positive.")

	if config['nx3'] % 2 == 1:
		raise ValueError('Config error: nx3 must be divisible by 2.')

	K = config['K'] # Number of tensor networks
	d = config['d'] # Number of dense layers in SoftmaxEmbeddingLayer
	input_shape = (2 * config['nx2'], config['nx3'] // 2, 1)
	input_tensor = Input(shape = input_shape)

	# Process parameter array (ell, a1, a2) and output weights for stacking the tensor networks
	preprocess_layer = SoftmaxEmbeddingLayer(K, d)
	params_input = Input(shape = (3,))
	stack_weights = preprocess_layer(params_input)[:, tf.newaxis, tf.newaxis, :]

	# Build the tensor networks using the custom keras layer class TLL
	models = [TTL(config) for _ in range(K)]

	# Combine the tensor networks based on the weights outputted by 'preprocess_layer'
	weighted_outputs = []
	for i, model in enumerate(models):
		processed_output = model(input_tensor)
		weighted_output = Multiply()([processed_output, stack_weights[..., i]])
		weighted_outputs.append(weighted_output)
	final_output = Add()(weighted_outputs)

	model = Model(inputs = [params_input, input_tensor], outputs = final_output)

	return model