import torch
import torch.utils.data
from torch import nn, optim
from torch.nn import functional as F
from torchvision import datasets, transforms
from torchvision.utils import save_image
from torch.autograd import Function


# ###copy
# class LowerBound(Function):
#     def forward(ctx, inputs, bound):
#         b = torch.ones_like(inputs) * bound
#         ctx.save_for_backward(inputs, b)
#         return torch.max(inputs, b)

#     def backward(ctx, grad_output):
#         inputs, b = ctx.saved_tensors

#         pass_through_1 = inputs >= b
#         pass_through_2 = grad_output < 0

#         pass_through = pass_through_1 | pass_through_2
#         return pass_through.type(grad_output.dtype) * grad_output, None


# class GDN(nn.Module):
#     """Generalized divisive normalization layer.
#     y[i] = x[i] / sqrt(beta[i] + sum_j(gamma[j, i] * x[j]))
#     """

#     def __init__(self,
#                  ch,
#                  inverse=False,
#                  beta_min=1e-6,
#                  gamma_init=.1,
#                  reparam_offset=2 ** -18):
#         super(GDN, self).__init__()
#         self.inverse = inverse
#         self.beta_min = beta_min
#         self.gamma_init = gamma_init
#         self.reparam_offset = torch.FloatTensor([reparam_offset])

#         self.build(ch)

#     def build(self, ch):
#         self.pedestal = self.reparam_offset ** 2
#         self.beta_bound = ((self.beta_min + self.reparam_offset ** 2) ** .5)
#         self.gamma_bound = self.reparam_offset

#         # Create beta param
#         beta = torch.sqrt(torch.ones(ch) + self.pedestal)
#         self.beta = nn.Parameter(beta.cuda())

#         # Create gamma param
#         eye = torch.eye(ch)
#         g = self.gamma_init * eye
#         g = g + self.pedestal
#         gamma = torch.sqrt(g)

#         self.gamma = nn.Parameter(gamma.cuda())
#         self.gamma_bound = self.gamma_bound.cuda()
#         self.beta_bound = self.beta_bound.cuda()
#         self.pedestal = self.pedestal.cuda()

#     def forward(self, inputs):
#         unfold = False
#         if inputs.dim() == 5:
#             unfold = True
#             bs, ch, d, w, h = inputs.size()
#             inputs = inputs.view(bs, ch, d * w, h)

#         _, ch, _, _ = inputs.size()

#         # Beta bound and reparam

#         beta = LowerBound()(self.beta, self.beta_bound.cuda())
#         beta = beta ** 2 - self.pedestal

#         # Gamma bound and reparam

#         gamma = LowerBound()(self.gamma, self.gamma_bound.cuda())
#         gamma = gamma ** 2 - self.pedestal
#         gamma = gamma.view(ch, ch, 1, 1)

#         # Norm pool calc
#         norm_ = nn.functional.conv2d(inputs ** 2, gamma, beta)
#         norm_ = torch.sqrt(norm_)

#         # Apply norm
#         if self.inverse:
#             outputs = inputs * norm_
#         else:
#             outputs = inputs / norm_

#         if unfold:
#             outputs = outputs.view(bs, ch, d, w, h)
#         return outputs


# check

# Ones_tensor = [{} for i in range(8)]
class LowerBound(Function):
    @staticmethod
    def forward(ctx, inputs, bound):
        b = torch.ones_like(inputs) * bound
        ctx.save_for_backward(inputs, b)
        return torch.max(inputs, b)
        
    @staticmethod
    def backward(ctx, grad_output):
        inputs, b = ctx.saved_tensors
        pass_through_1 = inputs >= b
        pass_through_2 = grad_output < 0

        pass_through = pass_through_1 | pass_through_2
        return pass_through.type(grad_output.dtype) * grad_output, None


class GDN(nn.Module):
    """Generalized divisive normalization layer.
    y[i] = x[i] / sqrt(beta[i] + sum_j(gamma[j, i] * x[j]))
    """

    def __init__(self,
                 ch,
                 inverse=False,
                 beta_min=1e-6,
                 gamma_init=0.1,
                 reparam_offset=2**-18):
        super(GDN, self).__init__()
        self.inverse = inverse
        self.beta_min = beta_min
        self.gamma_init = gamma_init
        self.reparam_offset = reparam_offset

        self.build(ch)

    def build(self, ch):
        self.pedestal = self.reparam_offset**2
        self.beta_bound = ((self.beta_min + self.reparam_offset**2)**0.5)
        self.gamma_bound = self.reparam_offset

        # Create beta param
        beta = torch.sqrt(torch.ones(ch)+self.pedestal)
        self.beta = nn.Parameter(beta)

        # Create gamma param
        eye = torch.eye(ch)
        g = self.gamma_init*eye
        g = g + self.pedestal
        gamma = torch.sqrt(g)

        self.gamma = nn.Parameter(gamma)
        self.pedestal = self.pedestal

    def forward(self, inputs):
        unfold = False
        if inputs.dim() == 5:
            unfold = True
            bs, ch, d, w, h = inputs.size() 
            inputs = inputs.view(bs, ch, d*w, h)

        _, ch, _, _ = inputs.size()

        # Beta bound and reparam
        beta = LowerBound.apply(self.beta, self.beta_bound)
        beta = beta**2 - self.pedestal

        # Gamma bound and reparam
        gamma = LowerBound.apply(self.gamma, self.gamma_bound)
        gamma = gamma**2 - self.pedestal
        gamma = gamma.view(ch, ch, 1, 1)

        # Norm pool calc
        norm_ = nn.functional.conv2d(inputs**2, gamma, beta)
        norm_ = torch.sqrt(norm_)

        # Apply norm
        if self.inverse:
            outputs = inputs * norm_
        else:
            outputs = inputs / norm_

        if unfold:
            outputs = outputs.view(bs, ch, d, w, h)
        return outputs
