import torch.nn as nn
from torch.nn import functional as F
from torch.autograd import Variable

class LRN(nn.Module):
    def __init__(self, alpha=1.0, beta=0.75, ACROSS_CHANNELS=False, local_size=1):
        super(LRN, self).__init__()
        self.local_size = local_size
        self.ACROSS_CHANNELS = ACROSS_CHANNELS
        if self.ACROSS_CHANNELS:
            self.average = nn.AvgPool3d(kernel_size=(local_size, 1,1),
                                        stride=1,
                                        padding=(int((local_size-1.0)/2),0,0))
        else:
          self.average = nn.AvgPool2d(kernel_size=local_size, stride=1, padding=int((local_size-1.0)/2))
        self.alpha = alpha
        self.beta = beta

    def forward(self, x):
        if self.ACROSS_CHANNELS:
            div = x.pow(2).unsqueeze(1)
            div = self.average(div).squeeze(1)
            div = div.mul(self.alpha).add(1.0).pow(self.beta)
        else:
            div = x.pow(2)
            div = self.average(div)
            div = div.mul(self.alpha).add(1.0).pow(self.beta)
        x = x.div(div)
        return x