import timm
import torch
import torch.nn as nn
import numpy as np

from .utils import activations, get_activation, Transpose


def forward_levit(pretrained, x):
    pretrained.model.forward_features(x)

    layer_1 = pretrained.activations["1"]
    layer_2 = pretrained.activations["2"]
    layer_3 = pretrained.activations["3"]

    layer_1 = pretrained.act_postprocess1(layer_1)
    layer_2 = pretrained.act_postprocess2(layer_2)
    layer_3 = pretrained.act_postprocess3(layer_3)

    return layer_1, layer_2, layer_3


def _make_levit_backbone(
        model,
        hooks=[3, 11, 21],
        patch_grid=[14, 14]
):
    pretrained = nn.Module()

    pretrained.model = model
    pretrained.model.blocks[hooks[0]].register_forward_hook(get_activation("1"))
    pretrained.model.blocks[hooks[1]].register_forward_hook(get_activation("2"))
    pretrained.model.blocks[hooks[2]].register_forward_hook(get_activation("3"))

    pretrained.activations = activations

    patch_grid_size = np.array(patch_grid, dtype=int)

    pretrained.act_postprocess1 = nn.Sequential(
        Transpose(1, 2),
        nn.Unflatten(2, torch.Size(patch_grid_size.tolist()))
    )
    pretrained.act_postprocess2 = nn.Sequential(
        Transpose(1, 2),
        nn.Unflatten(2, torch.Size((np.ceil(patch_grid_size / 2).astype(int)).tolist()))
    )
    pretrained.act_postprocess3 = nn.Sequential(
        Transpose(1, 2),
        nn.Unflatten(2, torch.Size((np.ceil(patch_grid_size / 4).astype(int)).tolist()))
    )

    return pretrained


class ConvTransposeNorm(nn.Sequential):
    """
    Modification of
        https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/levit.py: ConvNorm
    such that ConvTranspose2d is used instead of Conv2d.
    """

    def __init__(
            self, in_chs, out_chs, kernel_size=1, stride=1, pad=0, dilation=1,
            groups=1, bn_weight_init=1):
        super().__init__()
        self.add_module('c',
                        nn.ConvTranspose2d(in_chs, out_chs, kernel_size, stride, pad, dilation, groups, bias=False))
        self.add_module('bn', nn.BatchNorm2d(out_chs))

        nn.init.constant_(self.bn.weight, bn_weight_init)

    @torch.no_grad()
    def fuse(self):
        c, bn = self._modules.values()
        w = bn.weight / (bn.running_var + bn.eps) ** 0.5
        w = c.weight * w[:, None, None, None]
        b = bn.bias - bn.running_mean * bn.weight / (bn.running_var + bn.eps) ** 0.5
        m = nn.ConvTranspose2d(
            w.size(1), w.size(0), w.shape[2:], stride=self.c.stride,
            padding=self.c.padding, dilation=self.c.dilation, groups=self.c.groups)
        m.weight.data.copy_(w)
        m.bias.data.copy_(b)
        return m


def stem_b4_transpose(in_chs, out_chs, activation):
    """
    Modification of
        https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/levit.py: stem_b16
    such that ConvTranspose2d is used instead of Conv2d and stem is also reduced to the half.
    """
    return nn.Sequential(
        ConvTransposeNorm(in_chs, out_chs, 3, 2, 1),
        activation(),
        ConvTransposeNorm(out_chs, out_chs // 2, 3, 2, 1),
        activation())


def _make_pretrained_levit_384(pretrained, hooks=None):
    model = timm.create_model("levit_384", pretrained=pretrained)

    hooks = [3, 11, 21] if hooks == None else hooks
    return _make_levit_backbone(
        model,
        hooks=hooks
    )