import torch
import torch.nn as nn
from torch.distributions import Normal as NormalDistribution
# from .. import framework as fm
from ..op import framework as fm

from omnibelt import InitWall, unspecified_argument
import omnifig as fig

from .layers import make_MLP

#################
# region General
#################

class MLPBase(fm.FunctionBase, InitWall, nn.Sequential):
	def __init__(self, din, dout, hidden=None, initializer=None,
	             nonlin='elu', output_nonlin=None,
	             bias=True, output_bias=None, make_kwargs={}, _req_args=(), **kwargs):

		assert len(_req_args) == 0, str(_req_args)

		net = make_MLP(din=din, dout=dout, hidden=hidden, initializer=initializer,
		               nonlin=nonlin, output_nonlin=output_nonlin,
		               bias=bias, output_bias=output_bias, **make_kwargs)
		super().__init__(din=net.din, dout=net.dout, _req_args=tuple(net), **kwargs)

@fig.Component('mlp')
class MLP(fm.Function, MLPBase):
	def __init__(self, A, din=unspecified_argument, dout=unspecified_argument,
	             hidden=unspecified_argument, initializer=unspecified_argument,
	             nonlin=unspecified_argument, output_nonlin=unspecified_argument,
	             bias=unspecified_argument, output_bias=unspecified_argument,
	             make_kwargs=unspecified_argument, **kwargs):

		# region fill in args

		if din is unspecified_argument:
			din = A.pull('_din', '<>din')
		if dout is unspecified_argument:
			dout = A.pull('_dout', '<>dout')
		if hidden is unspecified_argument:
			hidden = A.pull('hidden_dims', '<>hidden_fc', '<>hidden', None)
		if initializer is unspecified_argument:
			initializer = A.pull('initializer', None)
		if nonlin is unspecified_argument:
			nonlin = A.pull('nonlin', 'elu')
		if output_nonlin is unspecified_argument:
			output_nonlin = A.pull('output_nonlin', '<>out_nonlin', None)
		if bias is unspecified_argument:
			bias = A.pull('bias', True)
		if output_bias is unspecified_argument:
			output_bias = A.pull('output_bias', '<>out_bias', None)
		if make_kwargs is unspecified_argument:
			make_kwargs = A.pull('make_kwargs', {})

		# endregion

		super().__init__(A, din=din, dout=dout,
		                 hidden=hidden, initializer=initializer,
		                 nonlin=nonlin, output_nonlin=output_nonlin,
		                 bias=bias, output_bias=output_bias,
		                 make_kwargs=make_kwargs, **kwargs)

@fig.Component('multihead')
class Multihead(fm.FunctionBase): # currently, the input dim for each head must be the same (but output can be different) TODO: generalize
	def __init__(self, A):
		
		din = A.pull('din')
		pin, N = din
		
		n_first = A.pull('n_first', False) # to change the "pin" and "N" dimensions
		if n_first:
			raise NotImplementedError
		
		merge = A.pull('merge', 'concat')
		if merge != 'concat':
			raise NotImplementedError
		
		A.push('heads.din', pin, silent=True)
		create_head = A.pull('heads')
		
		head_douts = A.pull('head_douts', None)
		try:
			len(head_douts)
		except TypeError:
			idouts = iter([head_douts]*len(create_head))
		else:
			idouts = iter(head_douts)
		
		heads = []
		
		nxt = create_head.view()
		if head_douts is not None:
			nxt.push('dout', next(idouts), silent=True)
		for head in create_head:
			heads.append(head)
			try:
				nxt = create_head.view()
			except StopIteration:
				break
			else:
				if head_douts is not None:
					nxt.push('dout', next(idouts), silent=True)
		
		assert N == len(heads), f'{N} vs {len(heads)}'
		
		douts = [head.dout for head in heads]
		
		if merge == 'concat':
			dout = sum(douts)
		
		super().__init__(din, dout)
		
		self.heads = nn.ModuleList(heads)
		
		self.merge = merge
		
	def forward(self, xs):
		
		xs = xs.permute(2,0,1)
		
		ys = [head(x) for head, x in zip(self.heads, xs)]
		
		return torch.cat(ys, dim=1)

@fig.Component('multilayer') # used for CNNs
class MultiLayer(fm.Function):

	def __init__(self, A, layers=None, din=None, dout=None, in_order=None, **kwargs):

		if layers is None:
			layers = self._create_layers(A, din=din, dout=dout, in_order=in_order)

		super().__init__(A, din=layers[0].din, dout=layers[-1].dout, **kwargs)

		self.layers = nn.ModuleList(layers)

	def _create_layers(self, A, din=None, dout=None, in_order=None):

		if din is None:
			din = A.pull('_din', '<>din', None)
		if dout is None:
			dout = A.pull('_dout', '<>dout', None)

		assert din is not None or dout is not None, 'need some input info'

		if in_order is None:
			in_order = A.pull('in_order', din is not None)

		self._current = din if in_order else dout
		self._req_key = 'din' if in_order else 'dout'
		self._find_key = 'dout' if in_order else 'din'
		end = dout if in_order else din

		pre, post = ('first', 'last') if in_order else ('last', 'first')

		pre_layer_info = A.pull(pre, None, raw=True)
		pre_layer = None if pre_layer_info is None else self._create_layer(pre_layer_info)

		mid = []
		create_layers = A.pull('layers', None, as_iter=True)
		if create_layers is not None:
			create_layers.set_auto_pull(False)  # prevents layer from being created before din/dout info is updated
			create_layers.set_reversed(not in_order)
	
			mid = [self._create_layer(layer) for layer in create_layers]

		post_layer = None
		post_layer_info = A.pull(post, None, raw=True)
		if end is not None or post_layer_info is not None:
			if post_layer_info is None:
				A.push('output_nonlin', None)
				post_layer_info = A.sub(post)
			mytype = post_layer_info.push('_type', 'dense-layer', silent=True, overwrite=False)
			if end is not None:
				post_layer_info.push(self._find_key, end, silent=True)
			if mytype == 'dense-layer' and in_order:
				post_layer_info.push('nonlin', '<>output_nonlin', silent=False, overwrite=False)
			post_layer = self._create_layer(post_layer_info, empty_find=end is None)

		layers = []
		if pre_layer is not None:
			layers.append(pre_layer)
		layers.extend(mid)
		if post_layer is not None:
			layers.append(post_layer)
		return layers if in_order else layers[::-1]

	def _create_layer(self, info, empty_find=True):
		current = self._current
		info.push(self._req_key, current, silent=True, overwrite=True)
		if empty_find:
			info.push(self._find_key, None, silent=True, overwrite=True)
		layer = info.pull_self()
		# c, n = self._current, getattr(layer, self._find_key)
		self._current = getattr(layer, self._find_key)
		return layer

	def forward(self, x):
		for i, layer in enumerate(self.layers):
			x = layer(x)
		return x

# endregion
#################








