# encoding: utf-8
"""
@author:  liaoxingyu
@contact: sherlockliao01@gmail.com
"""

import torch
from torch import nn
from torch.nn.modules.batchnorm import BatchNorm2d
from torch.nn import ReLU, LeakyReLU
from torch.nn.parameter import Parameter


class TLU(nn.Module):
    def __init__(self, num_features):
        """max(y, tau) = max(y - tau, 0) + tau = ReLU(y - tau) + tau"""
        super(TLU, self).__init__()
        self.num_features = num_features
        self.tau = Parameter(torch.Tensor(num_features))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.zeros_(self.tau)

    def extra_repr(self):
        return 'num_features={num_features}'.format(**self.__dict__)

    def forward(self, x):
        return torch.max(x, self.tau.view(1, self.num_features, 1, 1))


class FRN(nn.Module):
    def __init__(self, num_features, eps=1e-6, is_eps_leanable=False):
        """
        weight = gamma, bias = beta
        beta, gamma:
            Variables of shape [1, 1, 1, C]. if TensorFlow
            Variables of shape [1, C, 1, 1]. if PyTorch
        eps: A scalar constant or learnable variable.
        """
        super(FRN, self).__init__()

        self.num_features = num_features
        self.init_eps = eps
        self.is_eps_leanable = is_eps_leanable

        self.weight = Parameter(torch.Tensor(num_features))
        self.bias = Parameter(torch.Tensor(num_features))
        if is_eps_leanable:
            self.eps = Parameter(torch.Tensor(1))
        else:
            self.register_buffer('eps', torch.Tensor([eps]))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.ones_(self.weight)
        nn.init.zeros_(self.bias)
        if self.is_eps_leanable:
            nn.init.constant_(self.eps, self.init_eps)

    def extra_repr(self):
        return 'num_features={num_features}, eps={init_eps}'.format(**self.__dict__)

    def forward(self, x):
        """
        0, 1, 2, 3 -> (B, H, W, C) in TensorFlow
        0, 1, 2, 3 -> (B, C, H, W) in PyTorch
        TensorFlow code
            nu2 = tf.reduce_mean(tf.square(x), axis=[1, 2], keepdims=True)
            x = x * tf.rsqrt(nu2 + tf.abs(eps))
            # This Code include TLU function max(y, tau)
            return tf.maximum(gamma * x + beta, tau)
        """
        # Compute the mean norm of activations per channel.
        nu2 = x.pow(2).mean(dim=[2, 3], keepdim=True)

        # Perform FRN.
        x = x * torch.rsqrt(nu2 + self.eps.abs())

        # Scale and Bias
        x = self.weight.view(1, self.num_features, 1, 1) * x + self.bias.view(1, self.num_features, 1, 1)
        # x = self.weight * x + self.bias
        return x


def bnrelu_to_frn(module):
    """
    Convert 'BatchNorm2d + ReLU' to 'FRN + TLU'
    """
    mod = module
    before_name = None
    before_child = None
    is_before_bn = False

    for name, child in module.named_children():
        if is_before_bn and isinstance(child, (ReLU, LeakyReLU)):
            # Convert BN to FRN
            if isinstance(before_child, BatchNorm2d):
                mod.add_module(
                    before_name, FRN(num_features=before_child.num_features))
            else:
                raise NotImplementedError()

            # Convert ReLU to TLU
            mod.add_module(name, TLU(num_features=before_child.num_features))
        else:
            mod.add_module(name, bnrelu_to_frn(child))

        before_name = name
        before_child = child
        is_before_bn = isinstance(child, BatchNorm2d)
    return mod


def convert(module, flag_name):
    mod = module
    before_ch = None
    for name, child in module.named_children():
        if hasattr(child, flag_name) and getattr(child, flag_name):
            if isinstance(child, BatchNorm2d):
                before_ch = child.num_features
                mod.add_module(name, FRN(num_features=child.num_features))
            # TODO bn is no good...
            if isinstance(child, (ReLU, LeakyReLU)):
                mod.add_module(name, TLU(num_features=before_ch))
        else:
            mod.add_module(name, convert(child, flag_name))
    return mod


def remove_flags(module, flag_name):
    mod = module
    for name, child in module.named_children():
        if hasattr(child, 'is_convert_frn'):
            delattr(child, flag_name)
            mod.add_module(name, remove_flags(child, flag_name))
        else:
            mod.add_module(name, remove_flags(child, flag_name))
    return mod


def bnrelu_to_frn2(model, input_size=(3, 128, 128), batch_size=2, flag_name='is_convert_frn'):
    forard_hooks = list()
    backward_hooks = list()

    is_before_bn = [False]

    def register_forward_hook(module):
        def hook(self, input, output):
            if isinstance(module, (nn.Sequential, nn.ModuleList)) or (module == model):
                is_before_bn.append(False)
                return

            # input and output is required in hook def
            is_converted = is_before_bn[-1] and isinstance(self, (ReLU, LeakyReLU))
            if is_converted:
                setattr(self, flag_name, True)
            is_before_bn.append(isinstance(self, BatchNorm2d))

        forard_hooks.append(module.register_forward_hook(hook))

    is_before_relu = [False]

    def register_backward_hook(module):
        def hook(self, input, output):
            if isinstance(module, (nn.Sequential, nn.ModuleList)) or (module == model):
                is_before_relu.append(False)
                return
            is_converted = is_before_relu[-1] and isinstance(self, BatchNorm2d)
            if is_converted:
                setattr(self, flag_name, True)
            is_before_relu.append(isinstance(self, (ReLU, LeakyReLU)))

        backward_hooks.append(module.register_backward_hook(hook))

    # multiple inputs to the network
    if isinstance(input_size, tuple):
        input_size = [input_size]

    # batch_size of 2 for batchnorm
    x = [torch.rand(batch_size, *in_size) for in_size in input_size]

    # register hook
    model.apply(register_forward_hook)
    model.apply(register_backward_hook)

    # make a forward pass
    output = model(*x)
    output.sum().backward()  # Raw output is not enabled to use backward()

    # remove these hooks
    for h in forard_hooks:
        h.remove()
    for h in backward_hooks:
        h.remove()

    model = convert(model, flag_name=flag_name)
    model = remove_flags(model, flag_name=flag_name)
    return model