#!/usr/bin/env python
'''
component_pt
Created by Seria at 05/02/2019 1:41 PM
Email: zzqsummerai@yeah.net

                    _ooOoo_
                  o888888888o
                 o88`_ . _`88o
                 (|  0   0  |)
                 O \   。   / O
              _____/`-----‘\_____
            .’   \||  _ _  ||/   `.
            |  _ |||   |   ||| _  |
            |  |  \\       //  |  |
            |  |    \-----/    |  |
             \ .\ ___/- -\___ /. /
         ,--- /   ___\<|>/___   \ ---,
         | |:    \    \ /    /    :| |
         `\--\_    -. ___ .-    _/--/‘
   ===========  \__  NOBUG  __/  ===========

'''
# -*- coding:utf-8 -*-
import torch
import torch.nn as nn
from ..cockpit.engine_pt import EnginePT

__all__ = ('Craft', 'Rudder', 'Nozzle',
           'XAVIER', 'XAVIER_NORMAL', 'XAVIER_UNIFORM', 'NORMAL', 'UNIFORM', 'ZERO', 'ONE',
           'STEP', 'POLY', 'COSINE', 'EXP', 'WAVY',
           'Conv', 'MaxPool', 'AvgPool', 'Dense', 'Reshape',
           'Relu',
           'SftmXE',
           'AccCls',
           'Momentum')



XAVIER = 10
XAVIER_NORMAL = 11
XAVIER_UNIFORM = 12
NORMAL = 13
UNIFORM = 14
ZERO = 15
ONE = 16

STEP = 20
POLY = 21
COSINE = 22
EXP = 23
WAVY = 24



class Craft(nn.Module):
    def __init__(self, scope):
        super(Craft, self).__init__()
        self.scope = scope
        self.__pods = []
        self.__dict = {}
        self.__formulated = False

    def run(self, *args):
        raise NotImplementedError

    def forward(self, *args):
        # assert all([isinstance(a, Tensor) for a in args])
        return self.run(*args)

    def gear(self, gr):
        if isinstance(gr, bool):
            if gr:
                self.train()
            else:
                self.eval()
        elif isinstance(gr, EnginePT):
            if gr.param['device'] == 'gpu':
                if gr.param['rank'] < 0:
                    self.cuda()
                else:
                    self.to(gr.device[gr.param['rank']])
            elif self.param['device'] == 'cpu':
                self.cpu()
        else:
            raise Exception('NEBULAE ERROR ⨷ %s is not a valid type of gear.' % type(gr))

    @property
    def pods(self):
        return self.__pods

    @pods.setter
    def pods(self, pods):
        if not self.__formulated:
            self.__pods = pods
            self.__formulated = True

    def __getitem__(self, key):
        paths = key.split('/')
        craft = self
        for p in paths[:-1]:
            craft = getattr(craft, p)
        return craft.__dict.get(paths[-1], craft)

    def __setitem__(self, key, value):
        self.__dict[key] = value



class Rudder(object):
    def __init__(self):
        self.grader = torch.enable_grad()

    def __enter__(self):
        self.grader.__enter__()
        return True

    def __exit__(self, *args):
        self.grader.__exit__()

class Nozzle(object):
    def __init__(self):
        self.grader = torch.no_grad()

    def __enter__(self):
        self.grader.__enter__()
        return False

    def __exit__(self, *args):
        self.grader.__exit__()


def _initialize(layer, initializer, parameter):
    layer = (layer.weight, layer.bias)
    for elt, iniz, param in zip(layer, initializer, parameter):
        if iniz == XAVIER_NORMAL:
            nn.init.xavier_normal_(elt)
        elif iniz == XAVIER_UNIFORM:
            nn.init.xavier_uniform_(elt)
        elif iniz == UNIFORM:
            if param is None:
                param = (-1, 1)
            nn.init.uniform_(elt, a=param[0], b=param[1])
        elif iniz == NORMAL:
            if param is None:
                param = (0, 1)
            nn.init.normal_(elt, param[0], param[1])
        elif iniz == ZERO:
            nn.init.zeros_(elt)
        elif iniz == ONE:
            nn.init.ones_(elt)
        elif iniz is None:
            continue
        else:
            raise Exception('NEBULAE ERROR ⨷ the invoked initializer is not defined or supported.')



# ------------------------------------ Computation ------------------------------------ #

class Conv(Craft):
    def __init__(self, in_chs, out_chs, kernel: tuple, stride=1, padding=0, dilation=1, group=1,
                 w_init=XAVIER_NORMAL, w_param=None, b_init=ZERO, b_param=None, scope='CONV'):
        '''
        Args:
        - in_chs: input channel
        - out_chs: output channel
        - kernel: kernel size (must be a tuple)
        - stride: moving stride
        - padding: padding size
        - dilation: stride in atrous convolution
        - group: number of groups to be divided
        - w_init: weight initializer
        - w_param: options for initializing weight
        - b_init: bias initializer
        - b_param: options for initializing bias
        - scope: name scope
        '''
        super(Conv, self).__init__(scope)
        self.scope = scope
        dim = len(kernel)
        if dim == 1:
            conv_fn = nn.Conv1d
            pad_fn = nn.ConstantPad1d
        elif dim == 2:
            conv_fn = nn.Conv2d
            pad_fn = nn.ConstantPad2d
        elif dim == 3:
            conv_fn = nn.Conv3d
            pad_fn = nn.ConstantPad3d
        else:
            raise Exception('NEBULAE ERROR ⨷ %d-d convolution is not supported.' % dim)

        if isinstance(stride, int):
            stride = dim * [stride]
        if isinstance(dilation, int):
            dilation = dim * [dilation]

        self.pad = pad_fn(padding, 0)
        self.conv = conv_fn(in_chs, out_chs, kernel, stride=stride, dilation=dilation,
                            groups=group, bias=False if b_init is None else True)
        _initialize(self.conv, (w_init, b_init), (w_param, b_param))

    def run(self, x):
        y = self.conv(self.pad(x))
        return y



class MaxPool(Craft):
    def __init__(self, kernel: tuple, stride=2, padding=0, is_global=False, scope='MPOOL'):
        super(MaxPool, self).__init__(scope)
        self.scope = scope
        dim = len(kernel)
        if dim == 1:
            if is_global:
                pool_fn = nn.AdaptiveMaxPool1d
            else:
                pool_fn = nn.MaxPool1d
            pad_fn = nn.ConstantPad1d
        elif dim == 2:
            if is_global:
                pool_fn = nn.AdaptiveMaxPool2d
            else:
                pool_fn = nn.MaxPool2d
            pad_fn = nn.ConstantPad2d
        elif dim == 3:
            if is_global:
                pool_fn = nn.AdaptiveMaxPool3d
            else:
                pool_fn = nn.MaxPool3d
            pad_fn = nn.ConstantPad3d
        else:
            raise Exception('NEBULAE ERROR ⨷ %d-d pooling is not supported.' % dim)

        if isinstance(stride, int):
            stride = dim * [stride]

        self.pad = pad_fn(padding, 0)
        if is_global:
            assert padding == 0
            self.pool = pool_fn(tuple(dim*[1]))
        else:
            self.pool = pool_fn(kernel_size=kernel, stride=stride)

    def run(self, x):
        y = self.pool(self.pad(x))
        return y



class AvgPool(Craft):
    def __init__(self, kernel: tuple, stride=2, padding=0, is_global=False, scope='APOOL'):
        super(AvgPool, self).__init__(scope)
        self.scope = scope
        dim = len(kernel)
        if dim == 1:
            if is_global:
                pool_fn = nn.AdaptiveAvgPool1d
            else:
                pool_fn = nn.AvgPool1d
            pad_fn = nn.ConstantPad1d
        elif dim == 2:
            if is_global:
                pool_fn = nn.AdaptiveAvgPool2d
            else:
                pool_fn = nn.AvgPool2d
            pad_fn = nn.ConstantPad2d
        elif dim == 3:
            if is_global:
                pool_fn = nn.AdaptiveAvgPool3d
            else:
                pool_fn = nn.AvgPool3d
            pad_fn = nn.ConstantPad3d
        else:
            raise Exception('NEBULAE ERROR ⨷ %d-d pooling is not supported.' % dim)

        if isinstance(stride, int):
            stride = dim * [stride]

        self.pad = pad_fn(padding, 0)
        if is_global:
            assert padding==0
            self.pool = pool_fn(tuple(dim*[1]))
        else:
            self.pool = pool_fn(kernel_size=kernel, stride=stride)

    def run(self, x):
        y = self.pool(self.pad(x))
        return y



class Dense(Craft):
    def __init__(self, in_chs, out_chs, axis=-1,
                 w_init=XAVIER_NORMAL, w_param=None, b_init=ZERO, b_param=None, scope='DENSE'):
        super(Dense, self).__init__(scope)
        self.scope = scope
        if axis == 0:
            raise Exception('NEBULAE ERROR ⨷ you cannot apply dense layer along batch axis.')
        else:
            self.axis = axis

        if not b_init is None:
            use_bias = True
        else:
            use_bias = False

        self.fc = nn.Linear(in_chs, out_chs, bias=use_bias)
        _initialize(self.fc, (w_init, b_init), (w_param, b_param))

    def run(self, x):
        if self.axis == -1:
            y = self.fc(x)
        else:
            dim = x.ndim()
            permuted = [i for i in range(dim)]
            permuted = permuted[:self.axis] + permuted[self.axis + 1:] + [self.axis]
            x = x.transpose(*permuted)
            y = self.fc(x)
            permuted = [i for i in range(dim)]
            permuted = permuted[:self.axis] + [dim - 1] + permuted[self.axis:-1]
            y = y.transpose(*permuted)
        return y



class Reshape(Craft):
    def __init__(self, scope='RESHAPE'):
        super(Reshape, self).__init__(scope)
        self.scope = scope

    def run(self, x, shape):
        y = torch.reshape(x, shape)
        return y



# ------------------------------------ Activation ------------------------------------ #

class Relu(Craft):
    def __init__(self, scope='RELU'):
        super(Relu, self).__init__(scope)
        self.scope = scope
        self.actv = nn.ReLU()

    def run(self, x):
        y = self.actv(x)
        return y



# ------------------------------------ Distributing ------------------------------------ #

class BN(Craft):
    def __init__(self, out_chs, dim, mmnt=0.9, beta=False, gamma=False, scope='BN'):
        super(BN, self).__init__(scope)
        self.scope = scope
        if dim == 1:
            bn_fn = nn.BatchNorm1d
        elif dim == 2:
            bn_fn = nn.BatchNorm2d
        elif dim == 3:
            bn_fn = nn.BatchNorm3d
        else:
            raise Exception('NEBULAE ERROR ⨷ %d-d BN is not supported.' % dim)
        self.bn = bn_fn(out_chs, momentum=1-mmnt, affine=beta or gamma, eps=1e-5)

    def run(self, x):
        y = self.bn(x)
        return y



# ------------------------------------ Loss ------------------------------------ #

class SftmXE(Craft):
    def __init__(self, is_one_hot, scope='SFTMXE'):
        super(SftmXE, self).__init__(scope)
        self.scope = scope
        self.ioh = is_one_hot
        self.cost = nn.CrossEntropyLoss()

    def run(self, x, y):
        if self.ioh:
            y = torch.argmax(y, dim=-1)
        z = self.cost(x, y)
        return z



# ------------------------------------ Metric ------------------------------------ #

class AccCls(Craft):
    def __init__(self, multi_class, is_one_hot, scope='ACCCLS'):
        super(AccCls, self).__init__(scope)
        self.scope = scope
        if multi_class:
            assert not is_one_hot
        self.mulcls = multi_class
        self.ioh = is_one_hot

    def run(self, x, y):
        if self.mulcls:
            x = torch.round(x)
            correct = torch.mean((x == y).float(), dim=-1)
            z = torch.mean((correct == 1).float())
        else:
            if self.ioh:
                y = torch.argmax(y, dim=-1)
            x = torch.argmax(x, dim=-1)
            z = torch.mean((x == y).float())
        return z


# ------------------------------------ Optimizer ------------------------------------ #

def _lrStrategy(optimizer, lr, lr_decay, lr_params):
    if lr_decay == STEP:
        # lr_params: [Period of decay, Multiplicative factor of decay]
        return torch.optim.lr_scheduler.StepLR(optimizer, lr_params[0], gamma=lr_params[1])
    elif lr_decay == POLY:
        # lr_params: [End of decay, Power of decay]
        lr_update = lambda mile: (1.001 - min(mile / lr_params[0], 1)) ** lr_params[1]
        return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_update)
    elif lr_decay == COSINE:
        # lr_params: Maximum number of iterations
        return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=lr_params, eta_min=lr*0.001)
    elif lr_decay == EXP:
        # lr_params: [Period of decay, Base of decay]
        lr_update = lambda mile: lr_params[1] ** (mile / lr_params[0])
        return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_update)
    elif lr_decay == WAVY:
        # lr_params: [Period of wave, Dampening function]
        lr_update = lambda x: lr_params[1] ** (x - 1)
        return torch.optim.lr_scheduler.CyclicLR(optimizer, lr/2, lr,
             step_size_up=lr_params[0]//2, scale_fn=lr_update, scale_mode='cycle')
    else:
        raise KeyError('NEBULAE ERROR ⨷ %s decay is not supported or defined.' % lr_decay)



class Momentum(Craft):
    def __init__(self, hull, lr, mmnt=0.9, wd=0,
                 lr_decay=None, lr_params=None, grad_limit=-1,
                 update_scope=None, scope='MOMENTUM'):
        super(Momentum, self).__init__(scope)
        self.scope = scope
        # select parameters await updating
        if update_scope is None:
            update_var = hull.parameters()
        else:
            if isinstance(update_scope, str):
                update_scope = [update_scope]
            update_var = []
            for us in update_scope:
                paths = us.split('/')
                craft = hull
                for p in paths:
                    craft = getattr(craft, p)
                update_var.append(craft.parameters())

        if grad_limit>0:
            nn.utils.clip_grad_value_(update_var, grad_limit)
        self.optz = torch.optim.SGD(update_var, lr=lr, momentum=mmnt, weight_decay=wd)
        if isinstance(lr_decay, int):
            self.lrer = _lrStrategy(self.optz, lr, lr_decay, lr_params)
        else:
            self.lrer = None

    def run(self, target):
        self.optz.zero_grad()
        target.backward()
        self.optz.step()
        if self.lrer is not None:
            self.lrer.step()


# def conv(self, **kwargs):
#     message = kwargs.get('inputs', '_INPUTS')
#     return Pod(partial(self._conv, **kwargs), ['inputs'], kwargs['name'], [message],
#                {'type': 'CONV', 'cmajor': self.channel_major, 'out_chs': kwargs.get('out_chs', None),
#                 'kernel': kwargs.get('kernel', None), 'stride': kwargs.get('stride', (1, 1)),
#                 'auto_pad': kwargs.get('auto_pad', True)})
# def _conv(self, name, inputs, out_chs, kernel, w_init='xavier', w_param=None, b_init=None, b_param=None,
#             w_reg='l2', b_reg='l2', stride=(1, 1), dilation=(1, 1), group=1, auto_pad=True):
#     ch_err = Exception('NEBULAE ERROR ⨷ %s NHWC-type tensor is not supported in PyTorch core.')
#     dim = len(kernel)
#     if dim == 1:
#         conv_func = nn.Conv1d
#         if self.channel_major:
#             dform = 'NCW'
#         else:
#             dform = 'NWC'
#             raise ch_err
#     elif dim == 2:
#         conv_func = nn.Conv2d
#         if self.channel_major:
#             dform = 'NCHW'
#         else:
#             dform = 'NHWC'
#             raise ch_err
#     elif dim == 3:
#         conv_func = nn.Conv3d
#         if self.channel_major:
#             dform = 'NCDHW'
#         else:
#             dform = 'NDHWC'
#             raise ch_err
#     else:
#         raise Exception('NEBULAE ERROR ⨷ %d-d convolution is not supported.' % dim)
#
#     if self.channel_major:
#         size = inputs[2:]
#         in_chs = inputs[1]
#     else:
#         size = inputs[1:-1]
#         in_chs = inputs[-1]
#     if auto_pad:
#         padding = []
#         for d in range(dim):
#             padding.append(ceil(((ceil(size[d] / stride[d]) - 1) * stride[d] + kernel[d] + (dilation[d] - 1)
#                                  * (kernel[d] - 1) - size[d]) / 2))
#     else:
#         padding = dim * [0]
#     convolution = conv_func(in_chs, out_chs, kernel, stride=stride, padding=padding,
#                             dilation=dilation, groups=group, bias=False if b_init is None else True)
#     self._initialize(convolution, (w_init, b_init), (w_param, b_param), (w_reg, b_reg))
#     return convolution
#
# def trans_conv(self, **kwargs):
#     message = kwargs.get('inputs', '_INPUTS')
#     return Pod(partial(self._trans_conv, **kwargs), ['inputs'], kwargs['name'], [message],
#                {'type': 'CONV', 'cmajor': self.channel_major, 'output': kwargs.get('output', None),
#                 'kernel': kwargs.get('kernel', None), 'stride': kwargs.get('stride', (1, 1)),
#                 'auto_pad': kwargs.get('auto_pad', True)})
# def _trans_conv(self, name, inputs, output, kernel, w_init='xavier', w_param=None, b_init=None, b_param=None,
#             w_reg='l2', b_reg='l2', stride=(1, 1), dilation=(1, 1), group=1, auto_pad=True):
#     ch_err = Exception('NEBULAE ERROR ⨷ %s NHWC-type tensor is not supported in PyTorch core.')
#     dim = len(kernel)
#     if dim == 1:
#         conv_func = nn.ConvTranspose1d
#         if self.channel_major:
#             dform = 'NCW'
#         else:
#             dform = 'NWC'
#             raise ch_err
#     elif dim == 2:
#         conv_func = nn.ConvTranspose2d
#         if self.channel_major:
#             dform = 'NCHW'
#         else:
#             dform = 'NHWC'
#             raise ch_err
#     elif dim == 3:
#         conv_func = nn.ConvTranspose3d
#         if self.channel_major:
#             dform = 'NCDHW'
#         else:
#             dform = 'NDHWC'
#             raise ch_err
#     else:
#         raise Exception('NEBULAE ERROR ⨷ %d-d convolution is not supported.' % dim)
#
#     if self.channel_major:
#         in_size = inputs[2:]
#         in_chs = inputs[1]
#         out_chs = output[1]
#         out_size = output[2:]
#     else:
#         in_size = inputs[1:-1]
#         in_chs = inputs[-1]
#         out_chs = output[-1]
#         out_size = output[1:-1]
#     if auto_pad:
#         padding = []
#         for d in range(dim):
#             padding.append(ceil(((ceil(out_size[d] / stride[d]) - 1) * stride[d] + kernel[d] + (dilation[d] - 1)
#                                  * (kernel[d] - 1) - out_size[d]) / 2))
#     else:
#         padding = dim * [0]
#     compensation = []
#     for d in range(dim):
#         compensation.append(out_size[d] + 2 * padding[d] - kernel[d] - (in_size[d] - 1) * stride[d])
#     convolution = conv_func(in_chs, out_chs, kernel, stride, padding, output_padding=compensation,
#                             groups=group, bias=False if b_init is None else True, dilation=dilation)
#     self._initialize(convolution, (w_init, b_init), (w_param, b_param), (w_reg, b_reg))
#     return convolution
#
# # ------------------------------------ Activation ------------------------------------ #
#
# def sigmoid(self, **kwargs):
#     message = kwargs.get('inputs', '_INPUTS')
#     return Pod(partial(self._sigmoid, **kwargs), ['inputs'], kwargs['name'], [message], {'type': 'ACTIVATION'})
# def _sigmoid(self, name, inputs):
#     return nn.Sigmoid()
#
# def tanh(self, **kwargs):
#     message = kwargs.get('inputs', '_INPUTS')
#     return Pod(partial(self._tanh, **kwargs), ['inputs'], kwargs['name'], [message], {'type': 'ACTIVATION'})
# def _tanh(self, name, inputs):
#     return nn.Tanh()
#
# def softmax(self, **kwargs):
#     message = kwargs.get('inputs', '_INPUTS')
#     return Pod(partial(self._tanh, **kwargs), ['inputs'], kwargs['name'], [message], {'type': 'ACTIVATION'})
# def _softmax(self, name, inputs):
#     return nn.Softmax()
#
# def relu(self, **kwargs):
#     message = kwargs.get('inputs', '_INPUTS')
#     return Pod(partial(self._relu, **kwargs), ['inputs'], kwargs['name'], [message], {'type': 'ACTIVATION'})
# def _relu(self, name, inputs):
#     return nn.ReLU()
#
# def relu_leaky(self, **kwargs):
#     message = kwargs.get('inputs', '_INPUTS')
#     return Pod(partial(self._relu_leaky, **kwargs), ['inputs'], kwargs['name'], [message], {'type': 'ACTIVATION'})
# def _relu_leaky(self, name, inputs, alpha=0.2):
#     return nn.LeakyReLU(alpha)
#
# # ------------------------------------ Distributing ------------------------------------ #
#
# def dropout(self, **kwargs):
#     message = []
#     message.append(kwargs.get('inputs', '_INPUTS'))
#     message.append(kwargs.get('is_train', '_IS_TRAIN'))
#     return Pod(partial(self._dropout, **kwargs), ['inputs', 'is_train'], kwargs['name'], message, {'type': 'DROP'})
# def _dropout(self, name, inputs, p_drop, is_train):
#     return nn.Dropout(p_drop)
#
# def batch_norm(self, **kwargs):
#     message = []
#     message.append(kwargs.get('inputs', '_INPUTS'))
#     message.append(kwargs.get('is_train', '_IS_TRAIN'))
#     return Pod(partial(self._batch_norm, **kwargs), ['inputs', 'is_train'], kwargs['name'], message,
#                {'type': 'BN', 'cmajor': self.channel_major})
# def _batch_norm(self, name, inputs, is_train, mmnt=0.9, beta=False, gamma=False):
#     if self.channel_major:
#         num_feat = inputs[1]
#     else:
#         raise Exception('NEBULAE ERROR ⨷ %s NHWC-type tensor is not supported in PyTorch core.')
#     dim = len(inputs)-2
#     if dim == 1:
#         bn_func = nn.BatchNorm1d
#     elif dim == 2:
#         bn_func = nn.BatchNorm2d
#     elif dim == 3:
#         bn_func = nn.BatchNorm3d
#     else:
#         raise Exception('NEBULAE ERROR ⨷ %d-d BN is not supported.' % dim)
#     return bn_func(num_feat, momentum=1-mmnt, affine=beta or gamma, eps=1e-5)
#
# # def embedding(self, **kwargs):
# #     message = kwargs.get('inputs', '_INPUTS')
# #     return Pod(partial(self._embedding, **kwargs), ['inputs'], kwargs['name'], [message])
# # def _embedding(self, name, inputs, vocabulary, vec_dims, w_init='xavier', w_param=None):
# #     embd_vec = self._createVar(name+'_vec', [vocabulary, vec_dims], w_init, w_param)
# #     return tf.nn.embedding_lookup(embd_vec, inputs, name=name)
#
# # ------------------------------------ Pooling ------------------------------------ #
#
# def max_pool(self, **kwargs):
#     message = kwargs.get('inputs', '_INPUTS')
#     return Pod(partial(self._max_pool, **kwargs), ['inputs'], kwargs['name'], [message],
#                {'type': 'POOL', 'cmajor': self.channel_major, 'kernel': kwargs.get('kernel', (2, 2)),
#                 'stride': kwargs.get('stride', (2, 2)), 'auto_pad': kwargs.get('auto_pad', True),
#                 'if_global': kwargs.get('if_global', False)})
# def _max_pool(self, name, inputs, kernel=(2, 2), stride=(2, 2), auto_pad=True, if_global=False):
#     dim = len(kernel)
#     if dim == 1:
#         pool_func = nn.MaxPool1d
#         if self.channel_major:
#             dform = 'NCW'
#         else:
#             dform = 'NWC'
#             raise Exception('NEBULAE ERROR ⨷ %s NHWC-type tensor is not supported in PyTorch core.')
#     elif dim == 2:
#         pool_func = nn.MaxPool2d
#         if self.channel_major:
#             dform = 'NCHW'
#         else:
#             dform = 'NHWC'
#             raise Exception('NEBULAE ERROR ⨷ %s NHWC-type tensor is not supported in PyTorch core.')
#     elif dim == 3:
#         pool_func = nn.MaxPool3d
#         if self.channel_major:
#             dform = 'NCDHW'
#         else:
#             dform = 'NDHWC'
#             raise Exception('NEBULAE ERROR ⨷ %s NHWC-type tensor is not supported in PyTorch core.')
#     else:
#         raise Exception('NEBULAE ERROR ⨷ %d-d convolution is not supported.' % dim)
#
#     if self.channel_major:
#         size = inputs[2:]
#     else:
#         size = inputs[1:-1]
#     if if_global:
#         kernel = size
#         padding = dim * [0]
#     else:
#         if auto_pad:
#             padding = []
#             for d in range(dim):
#                 padding.append(ceil(((ceil(size[d] / stride[d]) - 1) * stride[d] + kernel[d] - size[d]) / 2))
#         else:
#             padding = dim * [0]
#     if dim == 1:
#         kernel = kernel[0]
#         stride = stride[0]
#         padding = padding[0]
#     return pool_func(kernel_size=kernel, stride=stride, padding=padding)
#
# def avg_pool(self, **kwargs):
#     message = kwargs.get('inputs', '_INPUTS')
#     return Pod(partial(self._avg_pool, **kwargs), ['inputs'], kwargs['name'], [message],
#                {'type': 'POOL', 'cmajor': self.channel_major, 'kernel': kwargs.get('kernel', (2, 2)),
#                 'stride': kwargs.get('stride', (2, 2)), 'auto_pad': kwargs.get('auto_pad', True),
#                 'if_global': kwargs.get('if_global', False)})
# def _avg_pool(self, name, inputs, kernel=(2, 2), stride=(2, 2), auto_pad=True, if_global=False):
#     dim = len(kernel)
#     if dim == 1:
#         pool_func = nn.AvgPool1d
#         if self.channel_major:
#             dform = 'NCW'
#         else:
#             dform = 'NWC'
#             raise Exception('NEBULAE ERROR ⨷ %s NHWC-type tensor is not supported in PyTorch core.')
#     elif dim == 2:
#         pool_func = nn.AvgPool2d
#         if self.channel_major:
#             dform = 'NCHW'
#         else:
#             dform = 'NHWC'
#             raise Exception('NEBULAE ERROR ⨷ %s NHWC-type tensor is not supported in PyTorch core.')
#     elif dim == 3:
#         pool_func = nn.AvgPool3d
#         if self.channel_major:
#             dform = 'NCDHW'
#         else:
#             dform = 'NDHWC'
#             raise Exception('NEBULAE ERROR ⨷ %s NHWC-type tensor is not supported in PyTorch core.')
#     else:
#         raise Exception('NEBULAE ERROR ⨷ %d-d convolution is not supported.' % dim)
#
#     if self.channel_major:
#         size = inputs[2:]
#     else:
#         size = inputs[1:-1]
#     if if_global:
#         kernel = size
#         padding = dim * [0]
#     else:
#         if auto_pad:
#             padding = []
#             for d in range(dim):
#                 padding.append(ceil(((ceil(size[d] / stride[d]) - 1) * stride[d] + kernel[d] - size[d]) / 2))
#         else:
#             padding = dim * [0]
#     if dim == 1:
#         kernel = kernel[0]
#         stride = stride[0]
#         padding = padding[0]
#     return pool_func(kernel_size=kernel, stride=stride, padding=padding)
#
# # ------------------------------------ Reshape ------------------------------------ #
#
# def reshape(self, **kwargs):
#     message = kwargs.get('inputs', '_INPUTS')
#     return Pod(partial(self._reshape, **kwargs), ['inputs'], kwargs['name'], [message],
#                {'type': 'RESHAPE', 'shape': kwargs.get('shape', None)})
# def _reshape(self, name, inputs, shape):
#     return PTreshape(shape)
#
# def flat(self, **kwargs):
#     message = kwargs.get('inputs', '_INPUTS')
#     return Pod(partial(self._flat, **kwargs), ['inputs'], kwargs['name'], [message], {'type': 'FLAT'})
# def _flat(self, name, inputs):
#     return PTflat()
# #
# # def slice(self, **kwargs):
# #     message = kwargs.get('inputs', '_INPUTS')
# #     return Pod(partial(self._slice, **kwargs), ['inputs'], kwargs['name'], [message])
# # def _slice(self, name, inputs, indices):
# #     for d in range(len(indices)):
# #         output = inputs[indices[d][0]:indices[d][1]]
# #     return tf.identity(output, name)
# #
# # def pad(self, **kwargs):
# #     message = kwargs.get('inputs', '_INPUTS')
# #     return Pod(partial(self._pad, **kwargs), ['inputs'], kwargs['name'], [message])
# # def _pad(self, name, inputs, margin, fill_in=0):
# #     return tf.pad(inputs, margin, constant_values=fill_in, name=name)
# #
# # # -------------------------------------- Arithmetic -------------------------------------- #
# #
# # def clip(self, **kwargs):
# #     message = kwargs.get('inputs', '_INPUTS')
# #     return Pod(partial(self._clip, **kwargs), ['inputs'], kwargs['name'], [message])
# # def _clip(self, name, inputs, min_max_vals):
# #     return tf.clip_by_value(inputs, min_max_vals[0], min_max_vals[1], name)
#
# # def argmax(self, **kwargs):
# #     message = kwargs.get('inputs', '_INPUTS')
# #     return Pod(partial(self._argmax, **kwargs), ['inputs'], kwargs['name'], [message])
# # def _argmax(self, name, inputs, axis):
# #     return tf.argmax(inputs, axis, name=name)
# #
# # # ------------------------------------ Image Manipulation ------------------------------------ #
# #
# # def resize(self, **kwargs):
# #     message = kwargs.get('inputs', '_INPUTS')
# #     return Pod(partial(self._resize, **kwargs), ['inputs'], kwargs['name'], [message])
# # def _resize(self, name, inputs, size, method='bilinear'):
# #     if method == 'bilinear':
# #         return tf.image.resize_bilinear(inputs, size, name=name)
# #     elif method == 'bicubic':
# #         return tf.image.resize_bicubic(inputs, size, name=name)
# #     elif method == 'crop':
# #         return tf.image.resize_image_with_crop_or_pad(inputs, size, name=name)
# #     else:
# #         raise KeyError('NEBULAE ERROR ⨷ %s is not a legal resize method.' % method)
# #
# #
# # # ------------------------------------ Redefine or Rename ------------------------------------ #
# #
# # def convert(self, **kwargs):
# #     message = kwargs.get('inputs', '_INPUTS')
# #     return Pod(partial(self._convert, **kwargs), ['inputs'], kwargs['name'], [message])
# # def _convert(self, name, inputs, dtype, trainable=False):
# #     '''
# #     convert data type or convert list/numpy array to tensor
# #     :param name:
# #     :param inputs: inputs tensor / list / numpy array
# #     :param dtype:
# #     :param trainable: if tensor is trainable
# #     :return: tensor
# #     '''
# #     if isinstance(inputs, (tf.Tensor, tf.SparseTensor, tf.Variable)):
# #         return tf.cast(inputs, tf.as_dtype(dtype), name=name)
# #     else:
# #         return tf.Variable(inputs, trainable=trainable, name=name)
#
# # ------------------------------------ Loss ------------------------------------ #
#
# def weight_decay(self, **kwargs):
#     return Pod(partial(self._weight_decay, **kwargs), ['penalty'], kwargs['name'],
#                [torch.ones(1) * kwargs['penalty']], {'type': 'LOSS'})
# def _weight_decay(self, name, penalty, decay_scope=None):
#     if not decay_scope is None:
#         raise Exception('NEBULAE ERROR ⨷ weight decay cannot be manipulated explicitly in PyTorch core.')
#     return PTwd()
#
# def sigm_xentropy(self, **kwargs):
#     message = []
#     message.append(kwargs.get('inputs', '_INPUTS'))
#     message.append(kwargs.get('label', '_LABEL'))
#     return Pod(partial(self._sigm_xentropy, **kwargs), ['inputs', 'label'], kwargs['name'], message, {'type': 'LOSS'})
# def _sigm_xentropy(self, name, inputs, label):
#     return nn.BCEWithLogitsLoss()
#
# def sftm_xentropy(self, **kwargs):
#     message = []
#     message.append(kwargs.get('inputs', '_INPUTS'))
#     message.append(kwargs.get('label', '_LABEL'))
#     return Pod(partial(self._sftm_xentropy, **kwargs), ['inputs', 'label'], kwargs['name'], message, {'type': 'LOSS'})
# def _sftm_xentropy(self, name, inputs, label, is_one_hot):
#     return PTsftmxe(is_one_hot)
#
# def mse(self, **kwargs):
#     message = []
#     message.append(kwargs.get('inputs', '_INPUTS'))
#     message.append(kwargs.get('label', '_LABEL'))
#     return Pod(partial(self._mse, **kwargs), ['inputs', 'label'], kwargs['name'], message, {'type': 'LOSS'})
# def _mse(self, name, inputs, label):
#     return nn.MSELoss()
#
# def mae(self, **kwargs):
#     message = []
#     message.append(kwargs.get('inputs', '_INPUTS'))
#     message.append(kwargs.get('label', '_LABEL'))
#     return Pod(partial(self._mae, **kwargs), ['inputs', 'label'], kwargs['name'], message, {'type': 'LOSS'})
# def _mae(self, name, inputs, label):
#     return nn.L1Loss()
#
# # ------------------------------------ Optimizer ------------------------------------ #
#
# def _lrStrategy(self, optimizer, lr, lr_decay, miles, param):
#     if lr_decay == 'step':
#         return torch.optim.lr_scheduler.StepLR(optimizer, miles, gamma=param)
#     elif lr_decay == 'poly':
#         lr_update = lambda epoch: (1.001 - epoch / miles) ** param
#         return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_update)
#     elif lr_decay == 'cosine':
#         return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=miles, eta_min=lr*0.001)
#     elif lr_decay == 'exp':
#         lr_update = lambda epoch: param ** (epoch / miles)
#         return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_update)
#     elif lr_decay == 'wavy':
#         lr_update = lambda x: param ** (x - 1)
#         return torch.optim.lr_scheduler.CyclicLR(optimizer, lr/2, lr,
#              step_size_up=miles, scale_fn=lr_update, scale_mode='cycle')
#     else:
#         raise KeyError('NEBULAE ERROR ⨷ %s decay is not supported or defined.' % lr_decay)
#
# def _selectParam(self, global_var, update_scope, ignore_name):
#     update_var = []
#     if ignore_name is None:
#         ignore_name = []
#     for l, p in global_var.items():
#         if (update_scope is None or l.startswith(update_scope)) and (l not in ignore_name):
#             update_var.extend(list(p))
#     return update_var
#
# def momentum(self, **kwargs):
#     recargs = {k: v for k, v in kwargs.items() if k not in ('name', 'inputs', 'media')}
#     recordConfig(os.path.join(os.getcwd(), 'nebulae_temp_config.json'),
#                  recargs, overwrite=False)
#
#     message = kwargs.get('inputs', '_INPUTS')
#     return Pod(partial(self._momentum, **kwargs), ['inputs', 'media'], kwargs['name'], [message, '_MEDIA'], {'type': 'OPTZ'})
# def _momentum(self, name, lr, inputs=None, mmnt=0.9, update_scope=None, ignore_name=None,
#               lr_decay=None, lr_miles=None, decay_param=None, grad_limit=None, media=None):
#     # N.B. media is network and wd factor
#     mod_params, wd = media
#     wd = wd.numpy().tolist()[0]
#     update_params = self._selectParam(mod_params, update_scope, ignore_name)
#     nn.utils.clip_grad_value_(update_params, grad_limit)
#     optz = torch.optim.SGD(update_params, lr=lr, momentum=mmnt, weight_decay=wd)
#     if isinstance(lr_decay, str):
#         optz = self._lrStrategy(optz, lr, lr_decay, lr_miles, decay_param)
#     return optz
#
# def nesterov(self, **kwargs):
#     recargs = {k: v for k, v in kwargs.items() if k not in ('name', 'inputs', 'media')}
#     recordConfig(os.path.join(os.getcwd(), 'nebulae_temp_config.json'),
#                  recargs, overwrite=False)
#
#     message = kwargs.get('inputs', '_INPUTS')
#     return Pod(partial(self._nesterov, **kwargs), ['inputs', 'media'], kwargs['name'], [message, '_MEDIA'], {'type': 'OPTZ'})
# def _nesterov(self, name, lr, inputs=None, mmnt=0.9, update_scope=None, ignore_name=None,
#               lr_decay=None, lr_miles=None, decay_param=None, grad_limit=None, media=None):
#     # N.B. media is network  and wd factor
#     mod_params, wd = media
#     wd = wd.numpy().tolist()[0]
#     update_params = self._selectParam(mod_params, update_scope, ignore_name)
#     nn.utils.clip_grad_value_(update_params, grad_limit)
#     optz = torch.optim.SGD(update_params, lr=lr, momentum=mmnt, weight_decay=wd, nesterov=True)
#     if isinstance(lr_decay, str):
#         optz = self._lrStrategy(optz, lr, lr_decay, lr_miles, decay_param)
#     return optz
#
# def adam(self, **kwargs):
#     recargs = {k:v for k,v in kwargs.items() if k not in ('name', 'inputs', 'media')}
#     recordConfig(os.path.join(os.getcwd(), 'nebulae_temp_config.json'),
#                  recargs, overwrite=False)
#     message = kwargs.get('inputs', '_INPUTS')
#     return Pod(partial(self._adam, **kwargs), ['inputs', 'media'], kwargs['name'], [message, '_MEDIA'], {'type': 'OPTZ'})
# def _adam(self, name, lr, inputs=None, mmnt1=0.9, mmnt2=0.999, update_scope=None, ignore_name=None,
#               lr_decay=None, lr_miles=None, decay_param=None, grad_limit=None, media=None):
#     # N.B. media is network and wd factor
#     mod_params, wd = media
#     wd = wd.numpy().tolist()[0]
#     update_params = self._selectParam(mod_params, update_scope, ignore_name)
#     nn.utils.clip_grad_value_(update_params, grad_limit)
#     optz = torch.optim.Adam(update_params, lr=lr, betas=(mmnt1, mmnt2), weight_decay=wd)
#     if isinstance(lr_decay, str):
#         optz = self._lrStrategy(optz, lr, lr_decay, lr_miles, decay_param)
#     return optz
#
# # ------------------------------------ Metric ------------------------------------ #
#
