import torch
import torch.nn as nn
import torch.nn.functional as F
from .functions.fakequantn import fakequantn
from .fakequantn_linear import FakeQuantNLinear

class FakeQuantNBatchNormLinear(FakeQuantNLinear):
    def __init__(self, in_features, out_features, nbits, bn_stats_freeze=False, momentum=0.1, eps=1e-5, **kwargs):
        """
        :param list args: args for :class:`antgine.modules.quantization.qconv2d.QConv2d`.
        :param int nbits: number of quantization bits.
        :param bool bn_stats_freeze: Freeze BatchNorm stats or continue computing moving averages.
        :param float momentum: Momentum for moving averages.
        :param float eps: Epsilon value preventing from division by zero.
        :param dict[str, any] kwargs: kwargs for :class:`antgine.modules.quantization.qconv2d.QConv2d`.
        """
        super().__init__(in_features, out_features, nbits=nbits, **kwargs)
        self._bn_stats_freeze = bn_stats_freeze
        self._momentum = momentum
        self._eps = eps
        self.gamma = nn.Parameter(torch.ones(out_features))
        self.beta = nn.Parameter(torch.zeros(out_features))
        self.register_buffer('running_mean', torch.zeros(out_features))
        self.register_buffer('running_var', torch.ones(out_features))
        

    @property
    def bn_stats_freeze(self):
        return self._bn_stats_freeze

    @bn_stats_freeze.setter
    def bn_stats_freeze(self, val):
        self._bn_stats_freeze = val

    def _ema(self, tensor, val):
        tensor.copy_(tensor * (1.0 - self._momentum) + val * self._momentum)

    def forward(self, inputs):
        if self.training and not self._bn_stats_freeze:
            out = F.linear(inputs, self.weight, self.bias)
            mean = out.mean(dim=0)
            var = out.var(dim=0)
            self._ema(self.running_mean, mean.data)
            self._ema(self.running_var, var.data)
        else:
            mean = self.running_mean
            var = self.running_var
        std = torch.sqrt(var + self._eps)
        f = (lambda x: x) if not self.quantize else self._qfunc
        return F.linear(inputs, f((self.weight * self.gamma.view(-1, 1)) / std.view(-1, 1)),
                        self.beta + (self.gamma * ((self.bias if self.bias is not None else 0) - mean)) / std)
                
            
