# Copyright (c) 2015-present, Facebook, Inc.
# All rights reserved.

# Modified from
# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
# Copyright 2020 Ross Wightman, Apache-2.0 License

import torch
import itertools
import utils

from timm.models.vision_transformer import trunc_normal_
from timm.models.registry import register_model

specification = {
    'LeViT_128S': {
        'C': '128_256_384',
        'D': 16,
        'N': '4_6_8',
        'X': '2_3_4',
        'drop_path': 0,
        'weights':
        'https://dl.fbaipublicfiles.com/LeViT/LeViT-128S-96703c44.pth'
    },
    'LeViT_128': {
        'C': '128_256_384',
        'D': 16,
        'N': '4_8_12',
        'X': '4_4_4',
        'drop_path': 0,
        'weights':
        'https://dl.fbaipublicfiles.com/LeViT/LeViT-128-b88c2750.pth'
    },
    'LeViT_192': {
        'C': '192_288_384',
        'D': 32,
        'N': '3_5_6',
        'X': '4_4_4',
        'drop_path': 0,
        'weights':
        'https://dl.fbaipublicfiles.com/LeViT/LeViT-192-92712e41.pth'
    },
    'LeViT_256': {
        'C': '256_384_512',
        'D': 32,
        'N': '4_6_8',
        'X': '4_4_4',
        'drop_path': 0,
        'weights':
        'https://dl.fbaipublicfiles.com/LeViT/LeViT-256-13b5763e.pth'
    },
    'LeViT_384': {
        'C': '384_512_768',
        'D': 32,
        'N': '6_9_12',
        'X': '4_4_4',
        'drop_path': 0.1,
        'weights':
        'https://dl.fbaipublicfiles.com/LeViT/LeViT-384-9bdaf2e2.pth'
    },
}

__all__ = [specification.keys()]


@register_model
def LeViT_128S(num_classes=1000,
               distillation=True,
               pretrained=False,
               fuse=False):
    return model_factory(**specification['LeViT_128S'],
                         num_classes=num_classes,
                         distillation=distillation,
                         pretrained=pretrained,
                         fuse=fuse)


@register_model
def LeViT_128(num_classes=1000,
              distillation=True,
              pretrained=False,
              fuse=False):
    return model_factory(**specification['LeViT_128'],
                         num_classes=num_classes,
                         distillation=distillation,
                         pretrained=pretrained,
                         fuse=fuse)


@register_model
def LeViT_192(num_classes=1000,
              distillation=True,
              pretrained=False,
              fuse=False):
    return model_factory(**specification['LeViT_192'],
                         num_classes=num_classes,
                         distillation=distillation,
                         pretrained=pretrained,
                         fuse=fuse)


@register_model
def LeViT_256(num_classes=1000,
              distillation=True,
              pretrained=False,
              fuse=False):
    return model_factory(**specification['LeViT_256'],
                         num_classes=num_classes,
                         distillation=distillation,
                         pretrained=pretrained,
                         fuse=fuse)


@register_model
def LeViT_384(num_classes=1000,
              distillation=True,
              pretrained=False,
              fuse=False):
    return model_factory(**specification['LeViT_384'],
                         num_classes=num_classes,
                         distillation=distillation,
                         pretrained=pretrained,
                         fuse=fuse)


FLOPS_COUNTER = 0


class Conv2d_BN(torch.nn.Sequential):
    def __init__(self,
                 a,
                 b,
                 ks=1,
                 stride=1,
                 pad=0,
                 dilation=1,
                 groups=1,
                 bn_weight_init=1,
                 resolution=-10000):
        super().__init__()
        self.add_module(
            'c',
            torch.nn.Conv2d(a,
                            b,
                            ks,
                            stride,
                            pad,
                            dilation,
                            groups,
                            bias=False))
        bn = torch.nn.BatchNorm2d(b)
        torch.nn.init.constant_(bn.weight, bn_weight_init)
        torch.nn.init.constant_(bn.bias, 0)
        self.add_module('bn', bn)

        global FLOPS_COUNTER
        output_points = ((resolution + 2 * pad - dilation *
                          (ks - 1) - 1) // stride + 1)**2
        FLOPS_COUNTER += a * b * output_points * (ks**2) // groups

    @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 = torch.nn.Conv2d(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


class Linear_BN(torch.nn.Sequential):
    def __init__(self, a, b, bn_weight_init=1, resolution=-100000):
        # print(f'a, b in Linear BN : {a,b}')
        super().__init__()
        self.add_module('c', torch.nn.Linear(a, b, bias=False))
        bn = torch.nn.BatchNorm1d(b)
        torch.nn.init.constant_(bn.weight, bn_weight_init)
        torch.nn.init.constant_(bn.bias, 0)
        self.add_module('bn', bn)

        global FLOPS_COUNTER
        output_points = resolution**2
        FLOPS_COUNTER += a * b * output_points

    @torch.no_grad()
    def fuse(self):
        l, bn = self._modules.values()
        w = bn.weight / (bn.running_var + bn.eps)**0.5
        w = l.weight * w[:, None]
        b = bn.bias - bn.running_mean * bn.weight / \
            (bn.running_var + bn.eps)**0.5
        m = torch.nn.Linear(w.size(1), w.size(0))
        m.weight.data.copy_(w)
        m.bias.data.copy_(b)
        return m

    def forward(self, x):
        l, bn = self._modules.values()
        x = l(x)
        return bn(x.flatten(0, 1)).reshape_as(x)


class BN_Linear(torch.nn.Sequential):
    def __init__(self, a, b, bias=True, std=0.02):
        super().__init__()
        self.add_module('bn', torch.nn.BatchNorm1d(a))
        l = torch.nn.Linear(a, b, bias=bias)
        trunc_normal_(l.weight, std=std)
        if bias:
            torch.nn.init.constant_(l.bias, 0)
        self.add_module('l', l)
        global FLOPS_COUNTER
        FLOPS_COUNTER += a * b

    @torch.no_grad()
    def fuse(self):
        bn, l = self._modules.values()
        w = bn.weight / (bn.running_var + bn.eps)**0.5
        b = bn.bias - self.bn.running_mean * \
            self.bn.weight / (bn.running_var + bn.eps)**0.5
        w = l.weight * w[None, :]
        if l.bias is None:
            b = b @ self.l.weight.T
        else:
            b = (l.weight @ b[:, None]).view(-1) + self.l.bias
        m = torch.nn.Linear(w.size(1), w.size(0))
        m.weight.data.copy_(w)
        m.bias.data.copy_(b)
        return m


def b16(n, activation, resolution=224):
    return torch.nn.Sequential(
        Conv2d_BN(3, n // 8, 3, 2, 1, resolution=resolution), activation(),
        Conv2d_BN(n // 8, n // 4, 3, 2, 1, resolution=resolution // 2),
        activation(),
        Conv2d_BN(n // 4, n // 2, 3, 2, 1, resolution=resolution // 4),
        activation(), Conv2d_BN(n // 2, n, 3, 2, 1,
                                resolution=resolution // 8))


class Residual(torch.nn.Module):
    def __init__(self, m, drop):
        super().__init__()
        self.m = m
        self.drop = drop

    def forward(self, x):
        if self.training and self.drop > 0:
            return x + self.m(x) * torch.rand(
                x.size(0), 1, 1, device=x.device).ge_(
                    self.drop).div(1 - self.drop).detach()
        else:
            return x + self.m(x)


class Attention(torch.nn.Module):
    def __init__(self,
                 dim,
                 key_dim,
                 num_heads=8,
                 attn_ratio=4,
                 activation=None,
                 resolution=14):
        super().__init__()
        # print(f'num_heads : {num_heads}')
        self.num_heads = num_heads
        self.scale = key_dim**-0.5
        self.key_dim = key_dim
        ## Smaller keys => Key Dim(32)
        self.nh_kd = nh_kd = key_dim * num_heads
        self.d = int(attn_ratio * key_dim)
        self.dh = int(attn_ratio * key_dim) * num_heads
        self.attn_ratio = attn_ratio
        print(f'self.dh : {self.dh} nh_kd :{nh_kd}')
        h = self.dh + nh_kd * 2
        self.qkv = Linear_BN(dim, h, resolution=resolution)
        self.proj = torch.nn.Sequential(
            activation(),
            Linear_BN(self.dh, dim, bn_weight_init=0, resolution=resolution))

        points = list(itertools.product(range(resolution), range(resolution)))
        N = len(points)
        attention_offsets = {}
        idxs = []
        for p1 in points:
            for p2 in points:
                offset = (abs(p1[0] - p2[0]), abs(p1[1] - p2[1]))
                if offset not in attention_offsets:
                    attention_offsets[offset] = len(attention_offsets)
                idxs.append(attention_offsets[offset])
        # print(F'attention_offsets : {len(attention_offsets)}')
        self.attention_biases = torch.nn.Parameter(
            torch.zeros(num_heads, len(attention_offsets)))
        self.register_buffer('attention_bias_idxs',
                             torch.LongTensor(idxs).view(N, N))

        global FLOPS_COUNTER
        #queries * keys
        FLOPS_COUNTER += num_heads * (resolution**4) * key_dim
        # softmax
        FLOPS_COUNTER += num_heads * (resolution**4)
        #attention * v
        FLOPS_COUNTER += num_heads * self.d * (resolution**4)

    @torch.no_grad()
    def train(self, mode=True):
        super().train(mode)
        if mode and hasattr(self, 'ab'):
            del self.ab
        else:
            self.ab = self.attention_biases[:, self.attention_bias_idxs]

    def forward(self, x):
        print(f'#' * 10)
        print(f'attention Input : {x.shape}')
        # x (B,N,C)
        B, N, _ = x.shape
        # (6,196,256)
        qkv = self.qkv(x)
        # print(f'qkv : {qkv.shape}')
        # (6, 196, 256x2)
        ## B 6, N 196 num_heads 4, 128
        ## Split [32, 32, 64] at Dimension 3
        q, k, v = qkv.view(B, N, self.num_heads,
                           -1).split([self.key_dim, self.key_dim, self.d],
                                     dim=3)
        # print(f'q {q.shape} k {k.shape} v {v.shape}')
        # q torch.Size([6, 196, 4, 32]) k torch.Size([6, 196, 4, 32]) v torch.Size([6, 196, 4, 64])
        # B, N,
        q = q.permute(0, 2, 1, 3)  # Batch, Head, Patch Num, Key Dim
        k = k.permute(0, 2, 1, 3)
        v = v.permute(0, 2, 1, 3)
        # matrix multiple
        # q torch.Size([6, 196, 4, 32]) k torch.Size([6, 196, 4, 32])
        # torch.matmul https://pytorch.org/docs/stable/generated/torch.matmul.html
        # q @ k.transpose(-2, -1) => q @ k.transpoes(-2, -1) : torch.Size([6, 4, 196, 196])
        # (6, 4, 196, 32) @ (6, 4, 196, 32)) => torch.Size([6, 4, 196, 196])
        # [196, 32] x [32, 196] = [196, 196]
        # print(f' self.scale : {self.scale}')
        # print(F'self.attention_bias_idxs] : {self.attention_bias_idxs}')
        # print(F'self.attention_bias_idxs[0] : {self.attention_bias_idxs[0]}')
        # print(f'self.attention_biases :{self.attention_biases.shape}')
        attn = ((q @ k.transpose(-2, -1)) * self.scale +
                (self.attention_biases[:, self.attention_bias_idxs]
                 if self.training else self.ab))
        attn = attn.softmax(dim=-1)
        # print(f' after Q K = ', attn.shape
        #       )  # attn.shape = [Batch, num_heads, Patch Num, Patch Num]
        # print(f'num_heads in Attention : {self.num_heads}')
        # torch.Size([6, 4, 196, 196])
        # After dividing Q KT / sqrt(Dh), Attn = A
        # print(f' attn : {attn.shape}')  # B,H, N, N => B, H, N, N
        # print(F' attn : {torch.sum(attn, dim=-1)}')
        # printx : torch.Size([6, 196, 256])(
        #     F'v : {v.shape} key_dim : {self.key_dim} attn_ratio : {self.attn_ratio}'
        # )  # [B, num_heads, Patch Num, ]
        # (attn @ v) => [6,4,196,196] @ [6,4,196,64] = [6,4,196,64]
        # print(f'attn @ v : {(attn@v).shape}')
        # Multihead の結果の統合
        x = (attn @ v).transpose(1, 2).reshape(B, N, self.dh)
        # x : [6, 196, 256]
        # print(f'x : {x.shape}')
        x = self.proj(x)
        print(f'attention Output : {x.shape}')
        return x


class Subsample(torch.nn.Module):
    def __init__(self, stride, resolution):
        super().__init__()
        self.stride = stride
        self.resolution = resolution

    def forward(self, x):
        print(f'subsample Module Input : {x.shape}')
        B, N, C = x.shape
        print(
            f'self.resolution : {self.resolution}  self.stride {self.stride}')
        # x =  [6, 196, 256]
        x = x.view(B, self.resolution, self.resolution,
                   C)[:, ::self.stride, ::self.stride].reshape(B, -1, C)
        # x.view(...) = [6, 14, 14, 256]
        # x.view(...)[...] = [6, 7, 7, 256]
        # x.view(...)[...].reshape() = [6, 49, 256]
        return x


class AttentionSubsample(torch.nn.Module):
    def __init__(self,
                 in_dim,
                 out_dim,
                 key_dim,
                 num_heads=8,
                 attn_ratio=2,
                 activation=None,
                 stride=2,
                 resolution=14,
                 resolution_=7):
        super().__init__()
        self.num_heads = num_heads
        self.scale = key_dim**-0.5
        self.key_dim = key_dim
        self.nh_kd = nh_kd = key_dim * num_heads
        self.d = int(attn_ratio * key_dim)
        self.dh = int(attn_ratio * key_dim) * self.num_heads
        self.attn_ratio = attn_ratio
        self.resolution_ = resolution_
        self.resolution_2 = resolution_**2
        h = self.dh + nh_kd
        self.kv = Linear_BN(in_dim, h, resolution=resolution)

        self.q = torch.nn.Sequential(
            Subsample(stride, resolution),
            Linear_BN(in_dim, nh_kd, resolution=resolution_))
        self.proj = torch.nn.Sequential(
            activation(), Linear_BN(self.dh, out_dim, resolution=resolution_))

        self.stride = stride
        self.resolution = resolution
        points = list(itertools.product(range(resolution), range(resolution)))
        points_ = list(
            itertools.product(range(resolution_), range(resolution_)))
        N = len(points)
        N_ = len(points_)
        attention_offsets = {}
        idxs = []
        for p1 in points_:
            for p2 in points:
                size = 1
                offset = (abs(p1[0] * stride - p2[0] + (size - 1) / 2),
                          abs(p1[1] * stride - p2[1] + (size - 1) / 2))
                if offset not in attention_offsets:
                    attention_offsets[offset] = len(attention_offsets)
                idxs.append(attention_offsets[offset])
        # print(f'len(attention_offsets): {len(attention_offsets)}') 196 or 49
        # B[0][0] = q, k is eqaul
        self.attention_biases = torch.nn.Parameter(
            torch.zeros(num_heads, len(attention_offsets)))
        self.register_buffer('attention_bias_idxs',
                             torch.LongTensor(idxs).view(N_, N))

        global FLOPS_COUNTER
        #queries * keys
        FLOPS_COUNTER += num_heads * \
            (resolution**2) * (resolution_**2) * key_dim
        # softmax
        FLOPS_COUNTER += num_heads * (resolution**2) * (resolution_**2)
        #attention * v
        FLOPS_COUNTER += num_heads * \
            (resolution**2) * (resolution_**2) * self.d

    @torch.no_grad()
    def train(self, mode=True):
        super().train(mode)
        if mode and hasattr(self, 'ab'):
            del self.ab
        else:
            self.ab = self.attention_biases[:, self.attention_bias_idxs]

    def forward(self, x):
        # print(f'attention Input ; {x.shape}')
        # Attention Input ; torch.Size([6, 196, 256])
        B, N, _ = x.shape
        # self.kv(x) => [6, 196, 768]
        # self.kv(x).view(B, N, self.num_heads, -1) => [6, 196, 8, **]
        # self.kv(x).view(B, N, self.num_heads, -1).split([]) => [6, 196, 8, **]
        k, v = self.kv(x).view(B, N, self.num_heads,
                               -1).split([self.key_dim, self.d], dim=3)
        k = k.permute(0, 2, 1, 3)  # BHNC
        v = v.permute(0, 2, 1, 3)  # BHNC
        q = self.q(x).view(B, self.resolution_2, self.num_heads,
                           self.key_dim).permute(0, 2, 1, 3)
        # Bh|x−x|,|y−y0|
        # our goal of attention bias is to provide positional information within each attention block
        # self.attention_bias_idxs
        attn = (q @ k.transpose(-2, -1)) * self.scale + \
            (self.attention_biases[:, self.attention_bias_idxs]
             if self.training else self.ab)
        attn = attn.softmax(dim=-1)

        x = (attn @ v).transpose(1, 2).reshape(B, -1, self.dh)
        x = self.proj(x)
        return x


class LeViT(torch.nn.Module):
    """ Vision Transformer with support for patch or hybrid CNN input stage
    """
    def __init__(self,
                 img_size=224,
                 patch_size=16,
                 in_chans=3,
                 num_classes=1000,
                 embed_dim=[192],
                 key_dim=[64],
                 depth=[12],
                 num_heads=[3],
                 attn_ratio=[2],
                 mlp_ratio=[2],
                 hybrid_backbone=None,
                 down_ops=[],
                 attention_activation=torch.nn.Hardswish,
                 mlp_activation=torch.nn.Hardswish,
                 distillation=True,
                 drop_path=0):
        super().__init__()
        global FLOPS_COUNTER

        self.num_classes = num_classes
        print(f'embed_dim[-1]: {embed_dim[-1]} {embed_dim}')
        print(f'key_dim : {key_dim}')
        self.num_features = embed_dim[-1]
        self.embed_dim = embed_dim
        self.distillation = distillation

        self.patch_embed = hybrid_backbone
        # print(self.patch_embed) # 3x3 conv
        self.blocks = []
        down_ops.append([''])
        resolution = img_size // patch_size
        for i, (ed, kd, dpth, nh, ar, mr, do) in enumerate(
                zip(embed_dim, key_dim, depth, num_heads, attn_ratio,
                    mlp_ratio, down_ops)):
            print(
                f'({i}) embed_dim {ed}, key_dim{kd}, depth{dpth}, num_heads{nh}, attn_ratio{ar}, mlp_ratio{mr}, down_ops {do}'
            )
            for _ in range(dpth):
                self.blocks.append(
                    Residual(
                        Attention(
                            ed,
                            kd,
                            nh,
                            attn_ratio=ar,
                            activation=attention_activation,
                            resolution=resolution,
                        ), drop_path))
                if mr > 0:
                    h = int(ed * mr)
                    self.blocks.append(
                        Residual(
                            torch.nn.Sequential(
                                Linear_BN(ed, h, resolution=resolution),
                                mlp_activation(),
                                Linear_BN(h,
                                          ed,
                                          bn_weight_init=0,
                                          resolution=resolution),
                            ), drop_path))
            if do[0] == 'Subsample':
                #('Subsample',key_dim, num_heads, attn_ratio, mlp_ratio, stride)
                resolution_ = (resolution - 1) // do[5] + 1
                self.blocks.append(
                    AttentionSubsample(*embed_dim[i:i + 2],
                                       key_dim=do[1],
                                       num_heads=do[2],
                                       attn_ratio=do[3],
                                       activation=attention_activation,
                                       stride=do[5],
                                       resolution=resolution,
                                       resolution_=resolution_))
                resolution = resolution_
                if do[4] > 0:  # mlp_ratio
                    h = int(embed_dim[i + 1] * do[4])
                    self.blocks.append(
                        Residual(
                            torch.nn.Sequential(
                                Linear_BN(embed_dim[i + 1],
                                          h,
                                          resolution=resolution),
                                mlp_activation(),
                                Linear_BN(h,
                                          embed_dim[i + 1],
                                          bn_weight_init=0,
                                          resolution=resolution),
                            ), drop_path))
        self.blocks = torch.nn.Sequential(*self.blocks)

        # Classifier head
        self.head = BN_Linear(
            embed_dim[-1],
            num_classes) if num_classes > 0 else torch.nn.Identity()
        if distillation:
            self.head_dist = BN_Linear(
                embed_dim[-1],
                num_classes) if num_classes > 0 else torch.nn.Identity()

        self.FLOPS = FLOPS_COUNTER
        FLOPS_COUNTER = 0

    @torch.jit.ignore
    def no_weight_decay(self):
        return {x for x in self.state_dict().keys() if 'attention_biases' in x}

    def forward(self, x):
        print(f'input of LeViT : {x.shape}')
        x = self.patch_embed(x)
        ## X : (B, C, H, W)
        print(F'after patch embedding : {x.shape}')
        # b, c, _, _ = x.shape
        # y = x.view(
        #     b,
        #     c,
        #     -1,
        # ).transpose(1, 2)
        # print(f'y : {y.shape}')
        x = x.flatten(2).transpose(1, 2)
        ## B,C,H,W -> B,C,N -> B,N,C
        print(f'After Patching : {x.shape}')
        ## X : (B, N, C)
        # Encoder = self.blocks
        x = self.blocks(x)
        print(f'x before average pooling : {x.shape}')
        x = x.mean(1)
        print(f'x after average pooling : {x.shape}')
        if self.distillation:
            x = self.head(x), self.head_dist(x)
            if not self.training:
                x = (x[0] + x[1]) / 2
        else:
            x = self.head(x)
        return x


def model_factory(C, D, X, N, drop_path, weights, num_classes, distillation,
                  pretrained, fuse):
    embed_dim = [int(x) for x in C.split('_')]
    num_heads = [int(x) for x in N.split('_')]
    depth = [int(x) for x in X.split('_')]
    act = torch.nn.Hardswish
    model = LeViT(
        patch_size=16,
        embed_dim=embed_dim,
        num_heads=num_heads,
        key_dim=[D] * 3,
        depth=depth,
        attn_ratio=[2, 2, 2],
        mlp_ratio=[2, 2, 2],
        down_ops=[
            #('Subsample',key_dim, num_heads, attn_ratio, mlp_ratio, stride)
            ['Subsample', D, embed_dim[0] // D, 4, 2, 2],
            ['Subsample', D, embed_dim[1] // D, 4, 2, 2],
        ],
        attention_activation=act,
        mlp_activation=act,
        hybrid_backbone=b16(embed_dim[0], activation=act),
        num_classes=num_classes,
        drop_path=drop_path,
        distillation=distillation)
    if pretrained:
        checkpoint = torch.hub.load_state_dict_from_url(weights,
                                                        map_location='cpu')
        model.load_state_dict(checkpoint['model'])
    if fuse:
        utils.replace_batchnorm(model)

    return model


if __name__ == '__main__':
    for name in specification:
        net = globals()[name](fuse=True, pretrained=True)
        net.eval()
        net(torch.randn(4, 3, 224, 224))
        print(name, net.FLOPS, 'FLOPs',
              sum(p.numel() for p in net.parameters() if p.requires_grad),
              'parameters')
