# -*- coding: utf-8 -*-
"""
Created on 2020/9/14
@project: SPAIC
@filename: Torch_Backend
@author: Hong Chaofei
@contact: hongchf@gmail.com
@description:
"""
from .Backend import Backend, backends
import torch
import numpy as np
# from torch import fx
#from torch.nn import Module, Parameter
import torch.nn.functional as fn
import torch.nn as nn
from typing import Tuple, Dict, Callable

class Torch_Engine(torch.nn.Module):
    def __init__(self, graph_operations):
        super(Torch_Engine, self).__init__()
        self._graph_operations = graph_operations

    def forward(self, variables: Dict[str,torch.Tensor]):
        temp_dict = dict()
        update_dict = dict()
        reduce_dict = dict()

        for op in self._graph_operations:
            # for inputs
            inputs = []
            for var in op.input:
                if var[0] == 'variables_dict':
                    inputs.append(variables[var[1]])
                elif var[0] == 'temp_dict':
                    inputs.append(temp_dict[var[1]])
                elif var[0] == 'update_dict':
                    inputs.append(update_dict[var[1]])
                elif var[0] == 'reduce_dict':
                    inputs.append(reduce_dict[var[1]])
            # compute the operation
            result = op.func(*inputs)
            if len(op.output) == 1: result = [result]
            # assign the result variables
            for ind, var in enumerate(op.output):
                if var[0] == 'temp_dict':
                    temp_dict[var[1]] = result[ind]
                elif var[0] == 'update_dict':
                    update_dict[var[1]] = result[ind]
                elif var[0] == 'reduce_dict':
                    if var[1] in reduce_dict:
                        reduce_dict[var[1]].append(result[ind])
                    else:
                        reduce_dict[var[1]] = [result[ind]]

        return update_dict



class Torch_Backend(Backend):
    backend_name = 'pytorch'

    def __init__(self, device='cpu'):
        super(Torch_Backend, self).__init__()

        self.device = device[0] if isinstance(device,list) else device
        self.data_type = torch.float32
        self.debug_data = []
        self.nograd_decorator = torch.no_grad()
        self.enablegrad_decorator = torch.enable_grad()
        pass

    def build(self):
        from torch import fx
        # self._graph_var_dicts = {'variables_dict': self._variables, 'temp_dict': dict(), 'update_dict': dict(),
        #                          'reduce_dict': dict()}
        # self._graph_var_dicts['temp_dict']['example_temp_dict_pytorch_datatype'] = torch.empty(1)
        # self._graph_var_dicts['update_dict']['example_temp_dict_pytorch_datatype'] = torch.empty(1)
        # self._graph_var_dicts['reduce_dict']['example_temp_dict_pytorch_datatype'] = torch.empty(1)
        #
        #self.update_step = jit.trace(self.update_step)
        self.engine = Torch_Engine(self._graph_operations)
        self.engine = fx.symbolic_trace(self.engine)

        # self.graph_update_step = torch.jit.script(self.engine)
        self.graph_update_step = self.engine
        # print(self.engine.code)

    def build_graph(self):

        for key, value in self._InitVariables_dict.items():
            if isinstance(value, torch.Tensor):
                value = value.to(self.device)
                self._InitVariables_dict[key] = value

        super(Torch_Backend, self).build_graph()

    def to_nograd_func(self, func):
        return self.nograd_decorator(func)

    def to_grad_func(self, func):
        return self.enablegrad_decorator(func)


    # def graph_update_step(self):
    #
    #     for op in self._graph_operations:
    #         inputs = []
    #         for var in op.input:
    #             inputs.append(self._graph_var_dicts[var[0]][var[1]])
    #
    #         if op.output[0] is 'reduce_dict':
    #             self._graph_var_dicts['reduce_dict'][op.output[1]].append(op.func(*inputs))
    #         else:
    #             self._graph_var_dicts[op.output[0]][op.output[1]] = op.func(*inputs)
    #
    #     return tuple(self._graph_var_dicts['variables_dict'].values())

    #  As of now, autograd support floating point Tensor types ( half, float, double and bfloat16) and complex Tensor types (cfloat, cdouble).
    def add_backend_variable(self, module, name, shape, value=None, grad=False, is_sparse=False, init=None, init_param=None):
        '''
        Parameters
        ----------
        name
        shape
        value
        init
        Returns
        -------
        '''

        if init_param is None:
            init_param = dict()
        if value is not None:
            if hasattr(value, "__len__"):
                if tuple(value.shape) == (1,) and isinstance(value, torch.Tensor):
                    self._variables[name] = value.to(self.device) * torch.ones(shape, dtype=self.data_type, device=self.device,
                                                               requires_grad=grad)
                elif tuple(value.shape) != tuple(shape):
                    raise ValueError("Value is not scalar and the shape of Value is not equal to shape")
                # add a sparse matrices with all dimensions greater than 2
                elif is_sparse:
                    i = np.nonzero(value)
                    v = value[i]

                    # Index for sparse matrix
                    sparse_index = name + '_sparse_index'
                    self._variables[sparse_index] = torch.LongTensor(i).to(device=self.device)
                    self._InitVariables_dict[sparse_index] = self._variables[sparse_index]

                    # Value for sparse matrix
                    sparse_value = name + '_sparse_value'
                    if init is not None:
                        # self._variables[sparse_value] = self.init_param(True, init)
                        data = torch.empty(shape, dtype=self.data_type, device=self.device, requires_grad=True)
                        self._variables[sparse_value] = self.param_init_operate[init](data, **init_param)
                    else:
                        self._variables[sparse_value] = torch.tensor(v, dtype=self.data_type, requires_grad=True, device=self.device)
                    self._parameters_dict[sparse_value] = self._variables[sparse_value]

                    # The shape of sparse matrix
                    sparse_shape = name + '_sparse_shape'
                    self._variables[sparse_shape] = torch.Size(shape)
                    self._InitVariables_dict[sparse_shape] = self._variables[sparse_shape]

                    # Sparse matrix
                    self._variables[name] = torch.sparse.FloatTensor(self._variables[sparse_index], self._variables[sparse_value], self._variables[sparse_shape])
                else:
                    # add a non sparse matrices with all dimensions greater than 2
                    if init is not None:
                        data = torch.empty(shape, dtype=self.data_type, device=self.device, requires_grad=grad)
                        init = init.lower()
                        if init in self.param_init_operate.keys():
                            self._variables[name] = self.param_init_operate[init](data, **init_param)
                        else:
                            raise ValueError("No initialize method: %s in param_init_operate" % init)
                    else:
                        if isinstance(value, torch.Tensor):
                            self._variables[name] = value.clone().detach().to(self.device)
                        else:
                            self._variables[name] = torch.tensor(value, dtype=self.data_type, device=self.device,
                                                                   requires_grad=grad)

            elif len(shape) == 0:
                # add constant
                self._variables[name] = torch.tensor(value, dtype=self.data_type, device=self.device, requires_grad=grad)

            else:
                # add a matrix through constant
                if init is not None:
                    # self._variables[name] = self.init_param(grad, init)
                    data = value*torch.ones(shape, dtype=self.data_type, device=self.device, requires_grad=grad)
                    init = init.lower()
                    if init in self.param_init_operate.keys():
                        self._variables[name] = self.param_init_operate[init](data, **init_param)
                    else:
                        raise ValueError("No initialize method: %s in param_init_operate" % init)
                else:
                    self._variables[name] = value*torch.ones(shape, dtype=self.data_type, device=self.device, requires_grad=grad)

        return self._variables[name]

    def set_variable_value(self, name, value, is_parameter):
        if is_parameter:
            assert name in self._parameters_dict
            assert self._parameters_dict[name].shape == value.shape
            if not isinstance(value, torch.Tensor):
                value = torch.tensor(value, dtype=self._parameters_dict[name].dtype, device=self._parameters_dict[name].device)
            with torch.no_grad():
                self._parameters_dict[name].data = value
        else:
            assert name in self._InitVariables_dict
            assert self._backend._InitVariables_dict[name].shape == value.shape
            if not isinstance(value, torch.Tensor):
                value = torch.tensor(value, dtype=self._InitVariables_dict[name].dtype, device=self._InitVariables_dict[name].device)
            with torch.no_grad():
                self._InitVariables_dict[name].data = value


    # def init_param(self, grad, *init):
    #     if init[0] in self.param_init_operate:
    #         init_op = self.param_init_operate[init[0]]
    #     else:
    #         raise ValueError("No init operate %s in param_init_operate" % init[0])
    #     inputs = []
    #     shape = init[1]
    #     data = torch.empty(shape, dtype=self.data_type, device=self.device, requires_grad=grad)
    #     inputs.append(data)
    #
    #     for var in init[2:]:
    #         inputs.append(var)
    #     return init_op(*inputs)

    def sparse_to_dense(self, index_name, value_name, shape_name):
        return torch.sparse.FloatTensor(self._variables[index_name], self._variables[value_name], self._variables[shape_name])

    def get_str(self, level):
        return level*' ' + 'torch_backend'

    def threshold(self, x, v_th):
        return torch.gt(x, v_th).type(self.data_type)

    def reset(self, v, o):
        return o.eq(0)*v

    def cat(self, x, dim=1):
        return torch.cat(x, dim)

    def stack(self, x, dim=1):       # 在指定维度dim上连接（concatenate）若干个张量。
        try:
            return torch.stack(x, dim)
        except:
            # patch for SLIF 2[O]
            for ii in range(len(x)):
                if x[ii].dim() ==2:
                    tmp = torch.zeros_like(x[ii])
                    tmp = torch.stack([x[ii], tmp], dim=1)
                    x[ii] = tmp
            return torch.stack(x, dim)

    def reduce_sum(self, x, *dim):
        if len(dim)==0:
            dim = 1
        return torch.sum(x, dim=dim)

    def index_select(self, x,  indices, dim=1):
        return torch.index_select(x, dim=dim, index=indices)

    def permute(self, x, permute_dim):
        return x.permute(permute_dim)

    def view(self, x, view_dim):

        x = x.contiguous().view(view_dim)
        return x


    def scatter(self, x, indices):
        return torch.scatter(x, dim=0, index=indices)

    def conv1d(self, x, kernel):
        return torch.conv1d(x, kernel)

    def conv_trans1d(self, x, kernel):
        return torch.conv_transpose1d(x, kernel)

    def conv_2d(self, x, kernel, stride, padding, dilation, groups):
        if x.dim() == kernel.dim() + 1:
            xshape = list(x.shape)
            xshape[0] = xshape[0]*xshape[1]
            extend_size = xshape[1]
            xshape.pop(1)
            out = fn.conv2d(x.reshape(xshape), kernel, stride=stride, padding=padding, dilation=dilation, groups=groups)
            outshape = list(out.shape)
            outshape[0] = outshape[0]//extend_size
            outshape.insert(1, extend_size)
            return out.view(outshape)
        else:
            return fn.conv2d(x, kernel, stride=stride, padding=padding, dilation=dilation, groups=groups)

    def conv_2d_complex(self, x, kernel, stride, padding, dilation, groups, beta):
        if x.dtype.is_complex:
            x = beta ** x.imag * (x.real * (0 + 1.0j))
        else:
            x = x*(0+1.0j)
        real = fn.conv2d(x.real, kernel, stride=stride, padding=padding, dilation=dilation, groups=groups)
        imag = fn.conv2d(x.imag, kernel, stride=stride, padding=padding, dilation=dilation, groups=groups)
        return torch.complex(real, imag)



    def conv_max_pool2d(self, x, kernel, pool_kernel, stride, pool_stride, padding, pool_padding, dilation, groups):
        return fn.max_pool2d(fn.conv2d(x, kernel, stride=stride, padding=padding,
                             dilation=dilation, groups=groups), kernel_size=pool_kernel,
                             stride=pool_stride, padding=pool_padding)

        # return fn.conv2d(fn.max_pool2d(x, int(max_kernel[0])), kernel, stride=int(stride), padding=int(padding), dilation=int(dilation), groups=int(groups))

    def conv_avg_pool2d(self, x, kernel, pool_kernel, stride, pool_stride, padding, pool_padding, dilation, groups):
        return fn.avg_pool2d(fn.conv2d(x, kernel, stride=stride, padding=padding,
                             dilation=dilation, groups=groups), kernel_size=pool_kernel,
                             stride=pool_stride, padding=pool_padding)
    def conv_add_bias(self, x, bias):
        bias_t = bias.repeat(x.shape[-2], x.shape[-1], 1).permute(2, 1, 0)
        return x+bias_t
    def max_pool2d(self, x, pool_kernel, pool_stride, pool_padding):
        return fn.max_pool2d(x, kernel_size=pool_kernel, stride=pool_stride, padding=pool_padding)
    def post_max_pool2d_complex(self, x, pool_kernel, pool_stride, pool_padding):
        pool_imag, pool_index = fn.max_pool2d(x.imag, kernel_size=pool_kernel, return_indices=True,
                                              stride=pool_stride, padding=pool_padding)
        x_shape = x.shape
        pool_shape = pool_index.shape
        pool_real = torch.gather(x.real.view(x_shape[0], x_shape[1], -1), dim=-1,
                                 index=pool_index.view(x_shape[0], x_shape[1],-1)).view(pool_shape[0], pool_shape[1],
                                                                                        pool_shape[2], pool_shape[3])


        return torch.complex(real=pool_real, imag=pool_imag)

    def avg_pool2d(self, x, pool_kernel, pool_stride, pool_padding):
        return fn.avg_pool2d(x, kernel_size=pool_kernel, stride=pool_stride, padding=pool_padding)

    def batchnorm2d(self, x, num_features):
        # 该实现方式忽略了running_mean 和 running_var
        # 当 batch_size 较小时，在推理阶段的统计特性就会和全局统计特性有着较大偏差，从而导致糟糕的效果，这种情况下推荐使用
        # SPAIC 的 Module 模块实现 batchnorm。
        device = x.device
        bn_2d = torch.nn.BatchNorm2d(num_features).to(device=device)
        return bn_2d(x)

    def dropout(self, x, p, inplace=False):
        return fn.dropout(x, p=p, inplace=inplace)

    def reshape_mat_mult(self, A, X):

        if A.dim() == 4:
            (batchsize, outchannels, H, W) = A.shape
            A = A.view(batchsize, -1)
        elif A.dim() == 5:
            (batchsize, extend, outchannels, H, W) = A.shape
            A = A.view(batchsize, extend, -1)

        return torch.matmul(A, X.permute(1, 0))

    def im2col_indices(self, x, kh, kw, padding, stride):
        return fn.unfold(x, (kh, kw), padding=padding, stride=stride)

    def conv2d_flatten(self, x):
        return x.view(x.shape[0], x.shape[1], -1)

    def feature_map_flatten(self, x):
        return x.view(x.shape[0], -1)

    def add(self, x, y):

        return x + y

    def minus(self, x, y):
        return x - y

    def div(self, x, y):
        return torch.div(x, y)

    def relu(self, x):
        return torch.relu(x)

    def sigmoid(self, x):
        return torch.sigmoid(x)

    def mat_mult_weight(self, A, X):
        '''
        Parameters
        ----------
        A--->preGroup:input
        X--->postGroup:weight
        Returns
        -------
        '''

        X = X.permute(1, 0)
        return torch.matmul(A, X)

    def mat_mult_weight_complex(self, A, X, beta):
        '''
        Parameters
        ----------
        A--->preGroup:input
        X--->postGroup:weight
        beta---> postGroup:beta_complex
        Returns
        -------
        '''
        if A.dtype.is_complex:
            A = A.unsqueeze(-2)
            beta = beta.unsqueeze(-1)
            real = A.real
            imag = A.imag
            O = beta ** imag * (real * (0 + 1.0j))
            return torch.sum(O * X, dim=-1)
        else:
            X = X.permute(1, 0)
            Out = torch.matmul(A.to(X.dtype), X)
            Out = Out*(0.0+1.0j)
            return Out


    def mat_mult_pre(self, A, X):
        '''
        Parameters
        ----------
        A--->preGroup:input
        X--->postGroup:weight
        Returns
        -------
        '''
        A = A.permute(1, 0)
        return torch.matmul(A, X)

    def mat_mult(self, A, X):
        '''
        Parameters
        ----------
        A--->preGroup:input
        X--->postGroup:weight
        Returns
        -------
        '''
        return torch.matmul(A, X)

    def bmm(self, A, X):
        '''
        Parameters
        ----------
        A---> postGroup
        X---> preGroup
        Returns
        -------
        '''
        return torch.bmm(A, X)

    def ger(self, A, X):
        '''
        Parameters
        ----------
        A---> postGroup
        X---> preGroup
        Returns
        -------
        '''
        return torch.ger(A, X)

    def sparse_mat_mult_weight(self, A, X):
        '''
       Parameters
       ----------
       A--->preGroup:sparseWeight(post, pre)
       X--->postGroup:input(batch, pre)
       Returns
       -------
       '''
        X = X.permute(1, 0)
        result = torch.sparse.mm(A, X)
        result = result.permute(1, 0)
        return result

    def var_mult(self, A, X):
        return A * X

    def mult_sum_weight(self, A, X):
        # X = X.permute(1, 0)
        A = A.permute(0, 2, 1)
        return torch.sum(A * X, dim=-1)

    def mat_linear(self, A, X, b):
        return torch.matmul(A, X) + b

    def var_linear(self, A, X, b):

        return A*X+b

    def unsqueeze(self, X, dim):
        return torch.unsqueeze(X, dim)

    def to_numpy(self, data: torch.Tensor):
        return data.detach().cpu().numpy()

    def to_tensor(self, data):
        if isinstance(data, torch.Tensor):
            return data.to(torch.float).to(self.device)
        else:
            return torch.tensor(data, dtype=torch.float, device=self.device)

    def exp(self, x):
        return torch.exp(x)


    def clamp_(self, data, min, max):
        with torch.no_grad():
            data.clamp_(min, max)


    def clamp_max_(self, data, max):
        with torch.no_grad():
            data.clamp_max_(max)



    def clamp_min_(self, data, min):
        with torch.no_grad():
            data.clamp_min_(min)


    def uniform(self, data, a=-0.0, b=1.0):
        '''
        Args:
            data(tensor): an n-dimensional torch.Tensor
            a(float): the lower bound of the uniform distribution
            b(float): the upper bound of the uniform distribution
        Returns:
            torch.nn.init.uniform_(data, a=0.0, b=1.0)
        '''
        return torch.nn.init.uniform_(data, a, b)


    def normal(self, data, mean=0.0, std=1.0):
        '''
        Args:
            data(tensor): an n-dimensional torch.Tensor
            mean(float): the mean of the normal distribution
            std(float): the standard deviation of the normal distribution
        Returns:
            torch.nn.init.normal_(data, mean=0.0, std=1.0)
        '''
        return torch.nn.init.normal_(data, mean, std)



    def xavier_normal(self, data, gain=1.0):
        '''
        Args:
            data(tensor): an n-dimensional torch.Tensor
            gain: an optional scaling factor
        Returns:
            torch.nn.init.xavier_normal_(data, gain=1.0)
        '''
        return torch.nn.init.xavier_normal_(data, gain)


    def xavier_uniform(self, data, gain=1.0):
        '''
        Args:
            data(tensor): an n-dimensional torch.Tensor
            gain: an optional scaling factor
        Returns:
            torch.nn.init.xavier_uniform_(data, gain=1.0)
        '''
        return torch.nn.init.xavier_uniform_(data, gain)


    def kaiming_normal(self, data, a=0, mode='fan_in', nonlinearity='leaky_relu'):
        '''
        Args:
            data(tensor): an n-dimensional torch.Tensor
            a: the negative slope of the rectifier used after this layer (only used with 'leaky_relu')
            mode: either 'fan_in' (default) or 'fan_out'. Choosing 'fan_in' preserves the magnitude of the variance of the weights in the forward pass. Choosing 'fan_out' preserves the magnitudes in the backwards pass.
            nonlinearity: the non-linear function (nn.functional name), recommended to use only with 'relu' or 'leaky_relu' (default).
        Returns:
            torch.nn.init.kaiming_normal_(data, a=0, mode='fan_in', nonlinearity='leaky_relu')
        '''
        return torch.nn.init.kaiming_normal_(data, a, mode, nonlinearity)


    def kaiming_uniform(self, data, a=0, mode='fan_in', nonlinearity='leaky_relu'):
        '''
        Args:
            data(tensor): an n-dimensional torch.Tensor
            a: the negative slope of the rectifier used after this layer (only used with 'leaky_relu')
            mode: either 'fan_in' (default) or 'fan_out'. Choosing 'fan_in' preserves the magnitude of the variance of the weights in the forward pass. Choosing 'fan_out' preserves the magnitudes in the backwards pass.
            nonlinearity: the non-linear function (nn.functional name), recommended to use only with 'relu' or 'leaky_relu' (default).
        Returns:
            torch.nn.init.kaiming_uniform_(data, a=0, mode='fan_in', nonlinearity='leaky_relu')
        '''
        return torch.nn.init.kaiming_uniform_(data, a, mode, nonlinearity)

    def constant(self, data, constant_value=0.0):
        '''
        Args:
            data(tensor): an n-dimensional torch.Tensor
            constant_value(float): the value to fill the tensor with
        Returns:
            torch.nn.init.constant_(data, constant_value)
        '''
        return torch.nn.init.constant_(data, constant_value)

    def sparse(self, data, sparsity=0.1, std=0.01):
        '''
        Args:
            data(tensor): an n-dimensional `torch.Tensor`
            sparsity(float): The fraction of elements in each column to be set to zero
            std(float): the standard deviation of the normal distribution used to generate
            the non-zero values
        Returns:
            torch.nn.init.sparse_(data, sparsity, std)
        '''
        return torch.nn.init.sparse_(data, sparsity, std)

    def weight_norm(self, weight, amp):
        w_norm = torch.norm(weight, p=2, dim=1, keepdim=True)
        # print(amp.item(), w_norm.item())
        return weight*amp/w_norm


    def sin(self, x):
        return torch.sin(x)

    def cos(self, x):
        return torch.cos(x)

    def tan(self, x):
        return torch.tan(x)

    def log(self, x):
        return torch.log(x)

    def log2(self, x):
        return torch.log2(x)

    def log10(self, x):
        return torch.log10(x)

    # def reset(self, x, v_reset, u_reset, spike):
    #
    #     # if hasattr(x, "__len__"):
    #     #     if x.shape != spike.shape:
    #     #         raise ValueError("%s and %s do not match" % (x.shape, spike.shape))
    #     mask = torch.eq(spike, 1)
    #     x[mask] = v_reset
    #     x[mask] += u_reset
    #     return x

    # def izh_v(self, v, u, psp):
    #     v = v+self.dt*(0.04*v*v+5*v+140-u+psp)
    #     return v
    #
    # def izh_u(self, a, b, v, u):
    #     u = u+self.dt*a*(b*v-u)
    #     return u


backends[Torch_Backend.backend_name] = Torch_Backend





# test = Torch_Backend()
# th = test.basic_operate['threshold']
# print(th(-1.0))