# pylint: disable=E1101
import numpy as np
from apdt.ml.tf_model import mlp_weight, MLP

try:
    import tensorflow as tf
except:
    pass

def deepsets_weight(**kwarg):
    '''Construct weight for deepsets.

    Parameter
    ---------
        name: str
            the variable name scope to be used.
        input_dim: int, default 10
            the number of input neurals
        output_dim: int, default 10
            the number of output neurals
        n_hidden: list of int, default 128
            the dimension of hidden states.
        mlp_n_hidden: list of int, default [128,128]
            the number of mlp hidden units, the length of list decide the number of layers.
    
    Return
    ------
        dict
            the dictionary of weight tensor.
    '''
    if 'name' not in kwarg:
        kwarg['name'] = 'mlp'
    if 'input_dim' not in kwarg:
        kwarg['input_dim'] = 10
    if 'output_dim' not in kwarg:
        kwarg['output_dim'] = 10
    if 'n_hidden' not in kwarg:
        kwarg['n_hidden'] = 128
    if 'mlp_n_hidden' not in kwarg:
        kwarg['mlp_n_hidden'] = [128, 128]   
    if 'trainable' not in kwarg:
        kwarg['trainable'] = True

    w = {}
    with tf.variable_scope(kwarg['name']):
        mlp_weight_encoder = mlp_weight(name="encoder_mlp", input_dim=kwarg['input_dim'], output_dim=kwarg['n_hidden'],
                                        n_hidden=kwarg['mlp_n_hidden'],trainable=kwarg['trainable'])
        mlp_weight_decoder = mlp_weight(name="decoder_mlp", input_dim=kwarg['n_hidden'], output_dim=kwarg['output_dim'],
                                        n_hidden=kwarg['mlp_n_hidden'],trainable=kwarg['trainable'])
    w.update({'pf0'+key: value for key, value in mlp_weight_encoder.items()})
    w.update({'pf1'+key: value for key, value in mlp_weight_decoder.items()})

    return w

def deepsets(x, axis, weights=None, name='deepsets', **kwarg):
    '''Map a set into one feature using deepsets.

    Parameter
    ---------
        x: tensor
            the input tensor
        axis: int
            along which axis is the set going
        weights: dict
            a dictionary of weights generated by apdt.ml.deepsets_weight
    
    Return
    ------
        tensor
            The output tensor
    
    '''
    
    if weights is None:
        if "output_dim" not in kwarg:
            raise Exception("output_dim should be given if weights is None.")
        weights = deepsets_weight(name=name, input_dim=x.shape[-1].value, output_dim=kwarg["output_dim"])
    
    mlp_weight_encoder = {key[3:]: value for key, value in weights.items() if key.startswith('pf0')}
    mlp_weight_decoder = {key[3:]: value for key, value in weights.items() if key.startswith('pf1')}
    with tf.variable_scope(name):
        a = MLP(x, weights=mlp_weight_encoder, name="encoder_mlp")
        h = tf.reduce_mean(a, axis)
        z = MLP(h, weights=mlp_weight_decoder, name="decoder_mlp")

    return z
