from operator import concat
import numpy as np
import torch
import torch.nn as nn
from torch.nn import init
from torch.autograd import Variable
from scipy.special import factorial2
from torch.nn import functional as F
from typing import Optional, Any
from torch import Tensor
# from torch import linalg as LA
import math


def create_variables(shape):
    '''
    :param name: A string. The name of the new variable
    :param shape: A list of dimensions
    :param initializer: User Xavier as default.
    :return: The created variable
    '''

    regularizer = None
    var = torch.empty(shape)
    nn.init.normal_(var)
    return var
    new_variables = torch.nn.parameter.Parameter(var)
    return new_variables


def bias_variable(shape):
    var = torch.empty(shape)
    nn.init.constant_(var, 0.1)
    new_variables = torch.nn.Parameter(var)
    return new_variables


def update_op_bn(X, pop_mean, pop_var, batch_mean, batch_var, decay, training):
    if training:
        pop_mean = batch_mean*(1-decay) + pop_mean*(decay)
        pop_var = batch_var*(1-decay) + pop_var*(decay)

    Xin = (X - batch_mean)
    Xin /= torch.sqrt(batch_var + 1e-8)
    return Xin


class Taylor(nn.Module):
    def __init__(
        self,
        k: int,
        planes: int,
        scale=None,
        decay=0.99,
        reuse=False,
        track_running_stats: bool = True,
        device=None,
        dtype=None
    ) -> None:
        self.h = []
        self.k = k
        self.scale = scale
        self.decay = decay
        self.reuse = reuse
        self.planes = planes
        factory_kwargs = {'device': device, 'dtype': dtype}
        self.track_running_stats = track_running_stats
        super().__init__()  # init the base class
        self.register_buffer(
            'pop_mean', torch.zeros([1, self.planes, 1, 1, self.k], **factory_kwargs))
        self.register_buffer(
            'pop_var', torch.ones([1, self.planes, 1, 1, self.k], **factory_kwargs))
        self.attention = torch.nn.Parameter(
            torch.sigmoid(create_variables(shape=[1,1,1,1, k])))
        self.kernel = torch.nn.Parameter(create_variables(shape=[1,planes, 1,1,k]))
        self.bias = bias_variable(shape=[1,planes,1,1])
        self.reset_running_stats()

    def reset_running_stats(self) -> None:
        self.pop_mean.zero_()
        self.pop_var.fill_(1)

    def forward(self, input):
        l = []
        shape = input.shape
        rank = len(shape)
        for i in range(self.k):
            temp = input ** (i+1)
            l.append(temp)
        X = torch.stack(l, dim=rank)
        # torch saves tensors as (batch, C, H, W)
        batch_var = torch.var(X, [0, 2, 3], keepdim=True)
        batch_mean = torch.mean(X, [0, 2, 3], keepdim=True)
        Xin = update_op_bn(X, self.pop_mean, self.pop_var,
                           batch_mean, batch_var, self.decay, self.training)
        # conv = torch.sum(self.attention * self.kernel * Xin, dim=-1) + self.bias
        conv = torch.sum(self.kernel * Xin, dim=-1) + self.bias
        return conv
