from __future__ import division
import math
import torch
from torch import nn
import torch.nn.functional as F
from torch.autograd import Function
from torch.nn.modules.utils import _pair
from torch.nn.modules.conv import _ConvNd
from torch.autograd.function import once_differentiable

from gcn import _C
# from gcn import _C

class GOF_Function(Function):
    @staticmethod
    def forward(ctx, weight, gaborFilterBank):
        ctx.save_for_backward(weight, gaborFilterBank)
        output = _C.gof_forward(weight, gaborFilterBank)
        return output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        weight, gaborFilterBank = ctx.saved_tensors
        grad_weight = _C.gof_backward(grad_output, gaborFilterBank)
        return grad_weight, None 

class MConv(_ConvNd):
    '''
    Baee layer class for modulated convolution
    '''
    def __init__(self, in_channels, out_channels, kernel_size, M=4, nScale=3, stride=1,
                    padding=0, dilation=1, groups=1, bias=True, expand=False, padding_mode='zeros'):
        if groups != 1:
            raise ValueError('Group-conv not supported!')
        kernel_size = (M,) + _pair(kernel_size)
        stride = _pair(stride)
        padding = _pair(padding)
        dilation = _pair(dilation)
        super(MConv, self).__init__(
            in_channels, out_channels, kernel_size, stride, padding, dilation,
            False, _pair(0), groups, bias, padding_mode)
        self.expand = expand
        self.M = M
        self.need_bias = bias
        self.generate_MFilters(nScale, kernel_size)
        self.GOF_Function = GOF_Function.apply

    def generate_MFilters(self, nScale, kernel_size):
        raise NotImplementedError

    def forward(self, x):
        if self.expand:
            x = self.do_expanding(x)
        new_weight = self.GOF_Function(self.weight, self.MFilters)
        new_bias = self.expand_bias(self.bias) if self.need_bias else self.bias
        if self.padding_mode == 'circular':
            expanded_padding = ((self.padding[1] + 1) // 2, self.padding[1] // 2,
                                (self.padding[0] + 1) // 2, self.padding[0] // 2)
            return F.conv2d(F.pad(input, expanded_padding, mode='circular'),
                            self.weight, self.bias, self.stride,
                            _pair(0), self.dilation, self.groups)
        return F.conv2d(x, new_weight, new_bias, self.stride,
                self.padding, self.dilation, self.groups)

    def do_expanding(self, x):
        index = []
        for i in range(x.size(1)):
            for _ in range(self.M):
                index.append(i)
        index = torch.LongTensor(index).cuda() if x.is_cuda else torch.LongTensor(index)
        return x.index_select(1, index)
    
    def expand_bias(self, bias):
        index = []
        for i in range(bias.size()):
            for _ in range(self.M):
                index.append(i)
        index = torch.LongTensor(index).cuda() if bias.is_cuda else torch.LongTensor(index)
        return bias.index_select(0, index)

class GConv(MConv):
    '''
    Gabor Convolutional Operation Layer
    '''
    def __init__(self, in_channels, out_channels, kernel_size, M=4, nScale=3, stride=1,
                    padding=0, dilation=1, groups=1, bias=True, expand=False, padding_mode='zeros'):
        super(GConv, self).__init__(in_channels, out_channels, kernel_size, M, nScale, stride,
                    padding, dilation, groups, bias, expand, padding_mode)

    def generate_MFilters(self, nScale, kernel_size):
        # To generate Gabor Filters
        self.register_buffer('MFilters', getGaborFilterBank(nScale, *kernel_size))

def getGaborFilterBank(nScale, M, h, w):
    Kmax = math.pi / 2
    f = math.sqrt(2)
    sigma = math.pi
    sqsigma = sigma ** 2
    postmean = math.exp(-sqsigma / 2)
    if h != 1:
        gfilter_real = torch.zeros(M, h, w)
        for i in range(M):
            theta = i / M * math.pi
            k = Kmax / f ** (nScale - 1)
            xymax = -1e309
            xymin = 1e309
            for y in range(h):
                for x in range(w):
                    y1 = y + 1 - ((h + 1) / 2)
                    x1 = x + 1 - ((w + 1) / 2)
                    tmp1 = math.exp(-(k * k * (x1 * x1 + y1 * y1) / (2 * sqsigma)))
                    tmp2 = math.cos(k * math.cos(theta) * x1 + k * math.sin(theta) * y1) - postmean # For real part
                    # tmp3 = math.sin(k*math.cos(theta)*x1+k*math.sin(theta)*y1) # For imaginary part
                    gfilter_real[i][y][x] = k * k * tmp1 * tmp2 / sqsigma			
                    xymax = max(xymax, gfilter_real[i][y][x])
                    xymin = min(xymin, gfilter_real[i][y][x])
            gfilter_real[i] = (gfilter_real[i] - xymin) / (xymax - xymin)
    else:
        gfilter_real = torch.ones(M, h, w)
    return gfilter_real