"""Implements the binary convolution modules.
"""
import math

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


class BinaryQuantize(Function):
    @staticmethod
    def forward(ctx, input, k, t):
        ctx.save_for_backward(input, k, t)
        out = torch.sign(input)
        return out

    @staticmethod
    def backward(ctx, grad_output):
        input, k, t = ctx.saved_tensors
        grad_input = k * t * (1 - torch.pow(torch.tanh(input * t), 2)) * grad_output
        return grad_input, None, None


class BinaryConv2d(nn.Conv2d):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
        super().__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.k = torch.tensor([10]).float().to(self.device)
        self.t = torch.tensor([0.1]).float().to(self.device)

    def forward(self, a):
        w = self.weight
        bw = w - w.view(w.size(0), -1).mean(-1).view(w.size(0), 1, 1, 1)
        bw = bw / bw.view(bw.size(0), -1).std(-1).view(bw.size(0), 1, 1, 1)

        bw = BinaryQuantize().apply(bw, self.k, self.t)
        ba = BinaryQuantize().apply(a, self.k, self.t)
        output = F.conv2d(ba, bw, self.bias, self.stride, self.padding, self.dilation, self.groups)
        return output
