from functools import partial
import torch
from torch import nn, einsum

from einops import rearrange
from einops.layers.torch import Rearrange, Reduce
import torch
import math
import warnings

from torch.nn.init import _calculate_fan_in_and_fan_out
import torch.nn.functional as F

# helpers
NUM_VARIANTS = 18

def cast_tuple(val, depth):
    return val if isinstance(val, tuple) else ((val,) * depth)

# LayerNorm = partial(nn.InstanceNorm2d, affine = True)

# classes

def drop_path(x, drop_prob: float = 0., training: bool = False):
    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
    'survival rate' as the argument.

    """
    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()  # binarize
    output = x.div(keep_prob) * random_tensor
    return output


class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
    """
    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)



def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    # Cut & paste from PyTorch official master until it's in a few official releases - RW
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1. + math.erf(x / math.sqrt(2.))) / 2.

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
                      "The distribution of values may be incorrect.",
                      stacklevel=2)

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        l = norm_cdf((a - mean) / std)
        u = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [l, u], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * l - 1, 2 * u - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor


def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
    # type: (Tensor, float, float, float, float) -> Tensor
    r"""Fills the input Tensor with values drawn from a truncated
    normal distribution. The values are effectively drawn from the
    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
    with values outside :math:`[a, b]` redrawn until they are within
    the bounds. The method used for generating the random values works
    best when :math:`a \leq \text{mean} \leq b`.
    Args:
        tensor: an n-dimensional `torch.Tensor`
        mean: the mean of the normal distribution
        std: the standard deviation of the normal distribution
        a: the minimum cutoff value
        b: the maximum cutoff value
    Examples:
        >>> w = torch.empty(3, 5)
        >>> nn.init.trunc_normal_(w)
    """
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)


def variance_scaling_(tensor, scale=1.0, mode='fan_in', distribution='normal'):
    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
    if mode == 'fan_in':
        denom = fan_in
    elif mode == 'fan_out':
        denom = fan_out
    elif mode == 'fan_avg':
        denom = (fan_in + fan_out) / 2

    variance = scale / denom

    if distribution == "truncated_normal":
        # constant is stddev of standard normal truncated to (-2, 2)
        trunc_normal_(tensor, std=math.sqrt(variance) / .87962566103423978)
    elif distribution == "normal":
        tensor.normal_(std=math.sqrt(variance))
    elif distribution == "uniform":
        bound = math.sqrt(3 * variance)
        tensor.uniform_(-bound, bound)
    else:
        raise ValueError(f"invalid distribution {distribution}")


def lecun_normal_(tensor):
    variance_scaling_(tensor, mode='fan_in', distribution='truncated_normal')







class Mlp(nn.Module):
    """ MLP as used in Vision Transformer, MLP-Mixer and related networks
    """
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class Attention(nn.Module):
    def __init__(self, num_variants, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
        super().__init__()

        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5
        self.num_variants = num_variants
        self.dim = dim

        self.Wkv_odd = nn.Linear(dim, 2 * head_dim, bias=qkv_bias)
        self.Wq_odd = nn.Linear(dim, dim, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x, num_layer):
        B, N, C = x.shape  # torch.Size([num_variants,batch_size,  num_patches, embedding_dim])
        q = self.Wq_odd(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)  # torch.Size([batch, #heads, #patches, embedding])
        kv = self.Wkv_odd(x).reshape(B, N, 2, C // self.num_heads).permute(2, 0, 1, 3)
        k, v = kv[0], kv[1]  # torch.Size([batch, #patches, embedding])

        attn = einsum('b h i d, b  j d -> b  h i j', q, k) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        x = einsum('b h i j, b j d -> b h i d', attn, v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

    def flops(self, N):
        # N - num tokens
        flops = 0
        flops += N * self.dim * self.dim
        flops += 2  * N * self.dim * self.dim
        flops += self.num_heads * N * (self.dim // self.num_heads) * N
        flops += self.num_heads * N * N * (self.dim // self.num_heads)
        flops += N * self.dim * self.dim

        return flops

class Attention_variants(nn.Module):
    def __init__(self, num_variants, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
        super().__init__()

        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5
        self.num_variants = num_variants
        self.dim = dim
        self.Wkv_even = nn.Linear(dim,  2* head_dim, bias=qkv_bias)
        self.Wq_even = nn.Linear(dim, dim, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x, variants_patches, num_layer):
        num_variants, B, N, C = variants_patches.shape  # torch.Size([num_variants,batch_size,  num_patches, embedding_dim])
        assert num_variants == self.num_variants, f
        "num variants ({num_variants}) doesn't match model ({self.num_variants})."

        variants_patches = variants_patches.permute(1, 0, 2, 3)  # torch.Size([batch, #varaints , #patches, embedding])
        q = self.Wq_even(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3).unsqueeze(dim=-2) # torch.Size([batch, #heads, #patches,1, embedding])
        kv = self.Wkv_even(variants_patches).reshape(B, num_variants, N, 2, C // self.num_heads).permute(3, 0, 2, 1, 4).unsqueeze(dim=2)

        k, v = kv[0], kv[1]  # k:  torch.Size([batch,  #patches,  #varaints, embedding])            v:  torch.Size([batch, 1 ,   #patches,  #varaints, embedding])

        attn = (q @ k.transpose(-2, -1)) * self.scale  # attn:  torch.Size([batch, #heads, #patches, 1, #varaints ])

        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        x = (attn @ v).squeeze()
        x = x.transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)

        return x

    def flops(self, N):
        # N - num tokens
        flops = 0
        flops += N * self.dim * self.dim
        flops += 2 * self.num_variants * N * self.dim * self.dim
        flops += self.num_heads * N * (self.dim // self.num_heads) * self.num_variants
        flops += self.num_heads * N * self.num_variants * (self.dim // self.num_heads)
        flops += N * self.dim * self.dim
        return flops


class Reduce_image_size(nn.Module):
    def __init__(self, dim, dim_out, seq_len):
        super().__init__()
        self.conv =  nn.Conv2d(dim, dim_out, 3, padding = 1)
        self.norm =  partial(nn.LayerNorm, eps=1e-6)(dim_out)
        self.pool = nn.MaxPool2d(3, stride = 2, padding = 1)
        self.img_size = seq_len ** 0.5
        self.dim_out = dim_out
        self.dim = dim
    def forward(self, x):
        B, embedding_dim, h, w = x.shape
        x = self.conv(x)
        x= x.permute(0,2, 3,1) # 44*128, 32, 32, 192
        x = self.norm(x) #TODO: layernorm on embedding dim?
        x = x.permute(0,3,1,2) # 44*128, 192, 32, 32
        x = self.pool(x)
        _, embedding_dim, h_new, w_new = x.shape

        x = x.reshape( B, embedding_dim, h_new, w_new)

        return x
    def flops(self):
        Ho, Wo = self.img_size, self.img_size
        flops = Ho * Wo * self.dim_out * self.dim * 3 * 3
        flops += Ho * Wo * self.dim_out
        flops += Ho * Wo * self.dim_out
        return flops


class Transformer_first(nn.Module):
    def __init__(self, num_variants, num_patches, depth, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, proj_drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=partial(nn.LayerNorm, eps=1e-6)):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.num_patches = num_patches
        self.dim = dim

        self.layers = nn.ModuleList([])
        self.pos_emb = nn.Parameter(torch.zeros(num_variants, num_patches, dim))
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.num_variants = num_variants
        self.embedding_dim = dim
        self.mlp_hidden_dim = mlp_hidden_dim

        for i in range(depth):
            self.layers.append(nn.ModuleList([
                 Attention(num_variants, dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=proj_drop) if not i == 0
                                             else Attention_variants(num_variants, dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=proj_drop),
                Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=proj_drop)
            ]))

        trunc_normal_(self.pos_emb, std=.02)

    def forward(self, patches):
        num_variants, B, num_patches, embedding_dim = patches.shape
        assert num_variants == self.num_variants
        assert embedding_dim == self.embedding_dim
        patches = patches.transpose(0, 1)  # 128, 44, 64, 16, 192
        pos_emb = self.pos_emb.unsqueeze(dim=0)
        patches = patches + pos_emb
        patches = patches.transpose(0, 1)
        x = patches[0]
        num_layer = 0
        for attn, mlp in self.layers:
            if num_layer == 0:
                x = x + self.drop_path(attn(self.norm1(x), self.norm1(patches), num_layer=num_layer))
            else:
                x = x + self.drop_path(attn(self.norm1(x), num_layer=num_layer))

            x = x + self.drop_path(mlp(self.norm2(x)))
            num_layer += 1
        #x = x.unsqueeze(dim=0)
        #patches = torch.cat((x, patches[1:]))
        return x

    def flops(self):
        flops = 0
        print("transfomer flops")
        for attn, mlp in self.layers:
            flops += attn.flops(self.num_patches)
            flops += self.num_patches * self.dim
            flops += 2* self.num_patches * self.dim * self.mlp_hidden_dim
            flops += self.num_patches * self.dim
        return flops


class Transformer(nn.Module):
    def __init__(self, num_variants, num_patches, depth, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, proj_drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=partial(nn.LayerNorm, eps=1e-6)):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.num_patches = num_patches
        self.dim = dim

        self.layers = nn.ModuleList([])
        self.pos_emb = nn.Parameter(torch.zeros(num_patches, dim))
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.num_variants = num_variants
        self.embedding_dim = dim
        self.mlp_hidden_dim = mlp_hidden_dim

        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                Attention(num_variants, dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=proj_drop),
                Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=proj_drop)
            ]))

        trunc_normal_(self.pos_emb, std=.02)

    def forward(self, x):
        B, num_patches,  embedding_dim = x.shape

        # assert  num_blocks == self.num_blocks
        assert embedding_dim == self.embedding_dim
        pos_emb = self.pos_emb.unsqueeze(dim=0)
        x = x + pos_emb
        num_layer = 0
        for attn, mlp in self.layers:
            x = x + self.drop_path(attn(self.norm1(x), num_layer=num_layer))
            x = x + self.drop_path(mlp(self.norm2(x)))
            num_layer +=1
        return x

    def flops(self):
        flops = 0
        print("transfomer flops")
        for attn, mlp in self.layers:
            # attn
            print("num patches: ", self.num_patches)

            flops += attn.flops(self.num_patches)
            print("attention flops: ", flops/ 1e9)
            # norm
            flops += self.num_patches * self.dim
            print("norm flops: ", flops/ 1e9)

            # mlp
            flops += 2* self.num_patches * self.dim * self.mlp_hidden_dim
            print("mlp flops: ", flops/ 1e9)

            # norm
            flops += self.num_patches * self.dim
            print("norm flops: ", flops/ 1e9)

        return flops
class PatchEmbed(nn.Module):
    """ 2D Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, stride_size=1, padding_size=1, in_chans=3, embed_dim=768, norm_layer=None):
        super().__init__()
        img_size = (img_size,img_size)
        patch_size = (patch_size,patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
        self.num_patches = self.grid_size[0] * self.grid_size[1]
        self.patch_size = patch_size
        self.embed_dim = embed_dim

        # original x patches
        # self.proj1 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj2 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj3 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj4 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj5 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj6 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj7 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj8 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj9 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj10 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj11 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj12 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj13 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj14 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj15 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj16 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj17 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj18 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)

        self.projections = nn.ModuleList([])

        for _ in range(NUM_VARIANTS):
            self.projections.append(
                nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
            )

        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

    def forward(self, img):
        B, C, H, W = img.shape
        assert H == self.img_size[0] and W == self.img_size[1], \
            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."


        patches = []
        img_new = img
        new_patch = self.norm(self.projections[0](img_new))
        new_patch = new_patch.flatten(2).transpose(1, 2)
        patches.append(new_patch)

        paddings = [(1,0), (0,1), (-1,0), (0,-1), (2,0), (0,2), (-2,0), (0,-2),
                     (1,1), (1,2), (1, -1), (-1,1), (-1,2), (-1,-1), (2,1),
                     (2,2), (2,-1)]
        for i in range(len(paddings)):
            cur_padding = paddings[i]
            img_new = torch.roll(img, shifts=cur_padding, dims=(2, 3))
            new_patch = self.norm(self.projections[i+1](img_new))
            new_patch = new_patch.flatten(2).transpose(1, 2)
            patches.append(new_patch)


        patches = torch.stack(patches)

        return  patches

    def flops(self):
        Ho, Wo = self.grid_size
        for i in range(10):
            flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
            flops += Ho * Wo * self.embed_dim
        return flops


class ShiftingTransformer(nn.Module):
    def __init__(self,*, scaling_factor, output_dir, num_variants, image_size,patch_size,num_classes,embedding_dim,heads,num_hierarchies,num_layers_per_block, mlp_mult = 4,channels = 3, dim_head = 64, qkv_bias=True,attn_drop=0.0,
                 proj_drop=0.0,stochastic_depth_drop = 0.1 , init_patch_embed_size =1, kernel_size=3, stride_size=1, padding_size=1):
        super().__init__()
        assert (image_size % patch_size) == 0, 'Image dimensions must be divisible by the patch size.'
        fmap_size = image_size // patch_size
        len_pyramid = len(num_layers_per_block)
        input_size_after_patch = image_size // init_patch_embed_size
        initial_num_blocks = fmap_size * fmap_size

        assert input_size_after_patch % patch_size == 0
        down_sample_ratio = input_size_after_patch // patch_size
        # assert len_pyramid == int(torch.log( torch.tensor([down_sample_ratio]).to(torch.float32))
        #                           / torch.log(torch.tensor([2]).to(torch.float32)) + 1)

        self.patch_size = patch_size
        self.kernel_size = kernel_size
        self.stride_size = stride_size
        self.padding_size= padding_size
        self.num_hierarchies = num_hierarchies
        hierarchies = list(reversed(range(num_hierarchies)))
        layer_heads = heads
        # layer_dims = list(map(lambda t: t * dim, mults))

        layer_dims = embedding_dim
        self.start_embedding = embedding_dim[0][0]
        self.end_embedding = embedding_dim[-1][-1]
        # dim_pairs = zip(layer_dims[:-1], layer_dims[1:])
        num_blocks = (initial_num_blocks, initial_num_blocks, initial_num_blocks , initial_num_blocks)
        if num_hierarchies == 4:
            seq_len = (initial_num_blocks, initial_num_blocks//(scaling_factor), initial_num_blocks//(scaling_factor**2), initial_num_blocks//(scaling_factor**3))
        elif num_hierarchies == 3:
            seq_len = (initial_num_blocks, initial_num_blocks//(scaling_factor), initial_num_blocks//(scaling_factor**2))
        self.end_num_patches = seq_len[-1]


        self.to_patch_embedding = PatchEmbed(img_size=image_size, patch_size=self.kernel_size, stride_size=self.stride_size, padding_size=padding_size, in_chans=3,
                                             embed_dim=self.start_embedding)
        block_repeats = cast_tuple(num_layers_per_block, num_hierarchies)
        layer_heads = cast_tuple(layer_heads, num_hierarchies)
        num_blocks = cast_tuple(num_blocks, num_hierarchies)
        seq_lens = cast_tuple(seq_len, num_hierarchies)

        dim_pairs = cast_tuple(layer_dims,num_hierarchies)
        self.layers = nn.ModuleList([])

        # print("build pyramid: ")
        # print("levels:", hierarchies, "heds: ", layer_heads, "dim paors: ", dim_pairs, "block repeats: ", block_repeats, "num blocks: ", num_blocks)
        for level, heads, (dim_in, dim_out), block_repeat, seq_len in zip(hierarchies, layer_heads, dim_pairs, block_repeats,seq_lens):
            print("level: ", level, "heads: ", heads, "dim in dim out: ", (dim_in, dim_out),
                  "block repeat: ", block_repeat,"seq len: ", seq_len)
            with open(output_dir + "/pyramid_parameters.txt", "a+") as text_file:
                text_file.write(" level: ")
                text_file.write(str(level))
                text_file.write(" heads: ")
                text_file.write(str(heads))
                text_file.write(" dim in: ")
                text_file.write(str(dim_in))
                text_file.write(" dim out: ")
                text_file.write(str(dim_out))
                text_file.write(" block repeat: ")
                text_file.write(str(block_repeat))
                text_file.write(" seq len: ")
                text_file.write(str(seq_len))
                text_file.write("\n")

            is_last = level == 0
            depth = block_repeat
            is_first = level == (num_hierarchies-1)
            print("is first: ", is_first, "is last: ", is_last)
            self.layers.append(nn.ModuleList([
                Transformer(num_variants, seq_len, depth, dim_in, heads, mlp_mult, qkv_bias=qkv_bias, attn_drop=attn_drop,
                            drop_path=stochastic_depth_drop, proj_drop=proj_drop) if not is_first else
                Transformer_first(num_variants, seq_len, depth, dim_in, heads, mlp_mult, qkv_bias=qkv_bias, attn_drop=attn_drop,
                            drop_path=stochastic_depth_drop, proj_drop=proj_drop),
                Reduce_image_size(dim_in, dim_out, seq_len) if not is_last else nn.Identity()
            ]))


        self.norm = partial(nn.LayerNorm, eps=1e-6)(self.end_embedding)
        self.mlp_head = nn.Linear(self.end_embedding, num_classes)
        self.apply(_init_vit_weights)
        text_file.close()


    def forward(self, img):
        # input shape:  128, 3, 32 , 32
        # print("input size: ", img.shape)
        patches = self.to_patch_embedding(img) # 44, 128, 256,192
        # print("x after embedding: ", patches.shape)
        num_hierarchies = len(self.layers)
        for level, (transformer, reduce_image_size) in zip(reversed(range(num_hierarchies)), self.layers):
            patches = transformer(patches)

            if level > 0:
                grid_size = (int(patches.shape[1]**0.5), int(patches.shape[1]**0.5))
                patches = to_image_plane(patches, grid_size, self.patch_size)
                patches = reduce_image_size(patches)
                patches = to_patches_plane(patches, self.patch_size)
        patches = self.norm(patches)
        patches_pool = torch.mean(patches, dim=(1))
        return self.mlp_head(patches_pool)


    def flops(self):
        flops = 0
        flops += self.to_patch_embedding.flops()
        print(flops/ 1e9)
        for level, (transformer, reduce_image_size) in zip(reversed(range(len(self.layers))), self.layers):
            flops += transformer.flops()
            print(flops/ 1e9)
            if level > 0:
                flops += reduce_image_size.flops()
                print(flops/ 1e9)
        # last norm
        flops += self.end_embedding * self.end_num_patches
        print(flops/ 1e9)
        # MLP
        flops += self.end_embedding * self.num_classes
        print(flops/ 1e9)
        return flops


def to_patches_plane(x, patch_size):
    patch_size = (patch_size, patch_size)
    batch, depth, height, width = x.shape # ([128, 192, 8, 8])
    x = x.reshape( batch, depth, height*width) #  128, 192, 64
    x = x.permute(0,2,1) #  128, 64, 192
    return x

def to_image_plane(x, grid_size, patch_size):
    batch, num_patches, depth = x.shape #  128, 256, 192
    x = x.permute(0,2,1)  # 128, 192, 256
    x = x.reshape(batch, depth, grid_size[0], grid_size[1])

    return x


def _init_vit_weights(m, n: str = '', head_bias: float = 0., jax_impl: bool = False):
    """ ViT weight initialization
    * When called without n, head_bias, jax_impl args it will behave exactly the same
      as my original init for compatibility with prev hparam / downstream use cases (ie DeiT).
    * When called w/ valid n (module name) and jax_impl=True, will (hopefully) match JAX impl
    """
    if isinstance(m, nn.Linear):
        trunc_normal_(m.weight, std=.02)
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Conv2d):
        # NOTE conv was left to pytorch default in my original init
        trunc_normal_(m.weight, std=.02)
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.LayerNorm):

        nn.init.zeros_(m.bias)
        nn.init.ones_(m.weight)

#
# if __name__ == '__main__':
#     image_size = 32
#     x = torch.rand(1, 3, image_size, image_size).cuda()
#     embedding_dim = 192
#     embedding_dim = ((embedding_dim, embedding_dim), (embedding_dim, embedding_dim), (embedding_dim, embedding_dim),
#                      (embedding_dim, embedding_dim))
#     heads = 3
#     num_hierarchies = 4 # number of hierarchies
#     num_layers_per_block = (2, 4, 4, 4) #the number of transformer blocks at each heirarchy, starting from the bottom
#     model = ShiftingTransformer(
#         scaling_factor=4,
#         output_dir=".",
#         num_variants=18,
#         image_size=32,
#         patch_size=1,
#         embedding_dim=embedding_dim,
#         heads=heads,
#         num_hierarchies=num_hierarchies,  # number of hierarchies
#         num_layers_per_block=num_layers_per_block,  # the number of transformer blocks at each heirarchy, starting from the bottom
#         num_classes=10,
#         init_patch_embed_size =1,
#         kernel_size = 3,
#         stride_size = 1,
#         padding_size = 1
#     ).cuda()
#     model.eval()
#     flops = model.flops()
#     print(f"number of GFLOPs: {flops / 1e9}")
#
#
#     # througput
#     repetitions = 1000
#     total_time = 0
#     optimal_batch_size = 128 #TODO
#     dummy_input = torch.randn(optimal_batch_size, 3, 32, 32, dtype=torch.float).cuda()
#     with torch.no_grad():
#         for rep in range(repetitions):
#             starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
#             starter.record()
#             _ = model(dummy_input)
#             torch.cuda.synchronize()
#             curr_time = starter.elapsed_time(ender)/1000
#             total_time += curr_time
#     Throughput = (repetitions*optimal_batch_size)/total_time
#     print("final throughput:", Throughput)
#

from functools import partial
import torch
from torch import nn, einsum

from einops import rearrange
from einops.layers.torch import Rearrange, Reduce
import torch
import math
import warnings

from torch.nn.init import _calculate_fan_in_and_fan_out
import torch.nn.functional as F

# helpers
NUM_VARIANTS = 18

def cast_tuple(val, depth):
    return val if isinstance(val, tuple) else ((val,) * depth)

# LayerNorm = partial(nn.InstanceNorm2d, affine = True)

# classes

def drop_path(x, drop_prob: float = 0., training: bool = False):
    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
    'survival rate' as the argument.

    """
    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()  # binarize
    output = x.div(keep_prob) * random_tensor
    return output


class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
    """
    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)



def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    # Cut & paste from PyTorch official master until it's in a few official releases - RW
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1. + math.erf(x / math.sqrt(2.))) / 2.

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
                      "The distribution of values may be incorrect.",
                      stacklevel=2)

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        l = norm_cdf((a - mean) / std)
        u = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [l, u], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * l - 1, 2 * u - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor


def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
    # type: (Tensor, float, float, float, float) -> Tensor
    r"""Fills the input Tensor with values drawn from a truncated
    normal distribution. The values are effectively drawn from the
    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
    with values outside :math:`[a, b]` redrawn until they are within
    the bounds. The method used for generating the random values works
    best when :math:`a \leq \text{mean} \leq b`.
    Args:
        tensor: an n-dimensional `torch.Tensor`
        mean: the mean of the normal distribution
        std: the standard deviation of the normal distribution
        a: the minimum cutoff value
        b: the maximum cutoff value
    Examples:
        >>> w = torch.empty(3, 5)
        >>> nn.init.trunc_normal_(w)
    """
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)


def variance_scaling_(tensor, scale=1.0, mode='fan_in', distribution='normal'):
    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
    if mode == 'fan_in':
        denom = fan_in
    elif mode == 'fan_out':
        denom = fan_out
    elif mode == 'fan_avg':
        denom = (fan_in + fan_out) / 2

    variance = scale / denom

    if distribution == "truncated_normal":
        # constant is stddev of standard normal truncated to (-2, 2)
        trunc_normal_(tensor, std=math.sqrt(variance) / .87962566103423978)
    elif distribution == "normal":
        tensor.normal_(std=math.sqrt(variance))
    elif distribution == "uniform":
        bound = math.sqrt(3 * variance)
        tensor.uniform_(-bound, bound)
    else:
        raise ValueError(f"invalid distribution {distribution}")


def lecun_normal_(tensor):
    variance_scaling_(tensor, mode='fan_in', distribution='truncated_normal')







class Mlp(nn.Module):
    """ MLP as used in Vision Transformer, MLP-Mixer and related networks
    """
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class Attention(nn.Module):
    def __init__(self, num_variants, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
        super().__init__()

        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5
        self.num_variants = num_variants
        self.dim = dim

        self.Wkv_odd = nn.Linear(dim, 2 * head_dim, bias=qkv_bias)
        self.Wq_odd = nn.Linear(dim, dim, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x, num_layer):
        B, N, C = x.shape  # torch.Size([num_variants,batch_size,  num_patches, embedding_dim])
        q = self.Wq_odd(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)  # torch.Size([batch, #heads, #patches, embedding])
        kv = self.Wkv_odd(x).reshape(B, N, 2, C // self.num_heads).permute(2, 0, 1, 3)
        k, v = kv[0], kv[1]  # torch.Size([batch, #patches, embedding])

        attn = einsum('b h i d, b  j d -> b  h i j', q, k) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        x = einsum('b h i j, b j d -> b h i d', attn, v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

    def flops(self, N):
        # N - num tokens
        flops = 0
        flops += N * self.dim * self.dim
        flops += 2  * N * self.dim * self.dim
        flops += self.num_heads * N * (self.dim // self.num_heads) * N
        flops += self.num_heads * N * N * (self.dim // self.num_heads)
        flops += N * self.dim * self.dim

        return flops

class Attention_variants(nn.Module):
    def __init__(self, num_variants, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
        super().__init__()

        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5
        self.num_variants = num_variants
        self.dim = dim
        self.Wkv_even = nn.Linear(dim,  2* head_dim, bias=qkv_bias)
        self.Wq_even = nn.Linear(dim, dim, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x, variants_patches, num_layer):
        num_variants, B, N, C = variants_patches.shape  # torch.Size([num_variants,batch_size,  num_patches, embedding_dim])
        assert num_variants == self.num_variants, f
        "num variants ({num_variants}) doesn't match model ({self.num_variants})."

        variants_patches = variants_patches.permute(1, 0, 2, 3)  # torch.Size([batch, #varaints , #patches, embedding])
        q = self.Wq_even(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3).unsqueeze(dim=-2) # torch.Size([batch, #heads, #patches,1, embedding])
        kv = self.Wkv_even(variants_patches).reshape(B, num_variants, N, 2, C // self.num_heads).permute(3, 0, 2, 1, 4).unsqueeze(dim=2)

        k, v = kv[0], kv[1]  # k:  torch.Size([batch,  #patches,  #varaints, embedding])            v:  torch.Size([batch, 1 ,   #patches,  #varaints, embedding])

        attn = (q @ k.transpose(-2, -1)) * self.scale  # attn:  torch.Size([batch, #heads, #patches, 1, #varaints ])

        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        x = (attn @ v).squeeze()
        x = x.transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)

        return x

    def flops(self, N):
        # N - num tokens
        flops = 0
        flops += N * self.dim * self.dim
        flops += 2 * self.num_variants * N * self.dim * self.dim
        flops += self.num_heads * N * (self.dim // self.num_heads) * self.num_variants
        flops += self.num_heads * N * self.num_variants * (self.dim // self.num_heads)
        flops += N * self.dim * self.dim
        return flops


class Reduce_image_size(nn.Module):
    def __init__(self, dim, dim_out, seq_len):
        super().__init__()
        self.conv =  nn.Conv2d(dim, dim_out, 3, padding = 1)
        self.norm =  partial(nn.LayerNorm, eps=1e-6)(dim_out)
        self.pool = nn.MaxPool2d(3, stride = 2, padding = 1)
        self.img_size = seq_len ** 0.5
        self.dim_out = dim_out
        self.dim = dim
    def forward(self, x):
        B, embedding_dim, h, w = x.shape
        x = self.conv(x)
        x= x.permute(0,2, 3,1) # 44*128, 32, 32, 192
        x = self.norm(x) #TODO: layernorm on embedding dim?
        x = x.permute(0,3,1,2) # 44*128, 192, 32, 32
        x = self.pool(x)
        _, embedding_dim, h_new, w_new = x.shape

        x = x.reshape( B, embedding_dim, h_new, w_new)

        return x
    def flops(self):
        Ho, Wo = self.img_size, self.img_size
        flops = Ho * Wo * self.dim_out * self.dim * 3 * 3
        flops += Ho * Wo * self.dim_out
        flops += Ho * Wo * self.dim_out
        return flops


class Transformer_first(nn.Module):
    def __init__(self, num_variants, num_patches, depth, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, proj_drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=partial(nn.LayerNorm, eps=1e-6)):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.num_patches = num_patches
        self.dim = dim

        self.layers = nn.ModuleList([])
        self.pos_emb = nn.Parameter(torch.zeros(num_variants, num_patches, dim))
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.num_variants = num_variants
        self.embedding_dim = dim
        self.mlp_hidden_dim = mlp_hidden_dim

        for i in range(depth):
            self.layers.append(nn.ModuleList([
                 Attention(num_variants, dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=proj_drop) if not i == 0
                                             else Attention_variants(num_variants, dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=proj_drop),
                Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=proj_drop)
            ]))

        trunc_normal_(self.pos_emb, std=.02)

    def forward(self, patches):
        num_variants, B, num_patches, embedding_dim = patches.shape
        assert num_variants == self.num_variants
        assert embedding_dim == self.embedding_dim
        patches = patches.transpose(0, 1)  # 128, 44, 64, 16, 192
        pos_emb = self.pos_emb.unsqueeze(dim=0)
        patches = patches + pos_emb
        patches = patches.transpose(0, 1)
        x = patches[0]
        num_layer = 0
        for attn, mlp in self.layers:
            if num_layer == 0:
                x = x + self.drop_path(attn(self.norm1(x), self.norm1(patches), num_layer=num_layer))
            else:
                x = x + self.drop_path(attn(self.norm1(x), num_layer=num_layer))

            x = x + self.drop_path(mlp(self.norm2(x)))
            num_layer += 1
        #x = x.unsqueeze(dim=0)
        #patches = torch.cat((x, patches[1:]))
        return x

    def flops(self):
        flops = 0
        print("transfomer flops")
        for attn, mlp in self.layers:
            flops += attn.flops(self.num_patches)
            flops += self.num_patches * self.dim
            flops += 2* self.num_patches * self.dim * self.mlp_hidden_dim
            flops += self.num_patches * self.dim
        return flops


class Transformer(nn.Module):
    def __init__(self, num_variants, num_patches, depth, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, proj_drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=partial(nn.LayerNorm, eps=1e-6)):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.num_patches = num_patches
        self.dim = dim

        self.layers = nn.ModuleList([])
        self.pos_emb = nn.Parameter(torch.zeros(num_patches, dim))
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.num_variants = num_variants
        self.embedding_dim = dim
        self.mlp_hidden_dim = mlp_hidden_dim

        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                Attention(num_variants, dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=proj_drop),
                Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=proj_drop)
            ]))

        trunc_normal_(self.pos_emb, std=.02)

    def forward(self, x):
        B, num_patches,  embedding_dim = x.shape

        # assert  num_blocks == self.num_blocks
        assert embedding_dim == self.embedding_dim
        pos_emb = self.pos_emb.unsqueeze(dim=0)
        x = x + pos_emb
        num_layer = 0
        for attn, mlp in self.layers:
            x = x + self.drop_path(attn(self.norm1(x), num_layer=num_layer))
            x = x + self.drop_path(mlp(self.norm2(x)))
            num_layer +=1
        return x

    def flops(self):
        flops = 0
        print("transfomer flops")
        for attn, mlp in self.layers:
            # attn
            print("num patches: ", self.num_patches)

            flops += attn.flops(self.num_patches)
            print("attention flops: ", flops/ 1e9)
            # norm
            flops += self.num_patches * self.dim
            print("norm flops: ", flops/ 1e9)

            # mlp
            flops += 2* self.num_patches * self.dim * self.mlp_hidden_dim
            print("mlp flops: ", flops/ 1e9)

            # norm
            flops += self.num_patches * self.dim
            print("norm flops: ", flops/ 1e9)

        return flops
class PatchEmbed(nn.Module):
    """ 2D Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, stride_size=1, padding_size=1, in_chans=3, embed_dim=768, norm_layer=None):
        super().__init__()
        img_size = (img_size,img_size)
        patch_size = (patch_size,patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
        self.num_patches = self.grid_size[0] * self.grid_size[1]
        self.patch_size = patch_size
        self.embed_dim = embed_dim

        # original x patches
        # self.proj1 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj2 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj3 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj4 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj5 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj6 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj7 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj8 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj9 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj10 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj11 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj12 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj13 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj14 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj15 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj16 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj17 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
        # self.proj18 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)

        self.projections = nn.ModuleList([])

        for _ in range(NUM_VARIANTS):
            self.projections.append(
                nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size)
            )

        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

    def forward(self, img):
        B, C, H, W = img.shape
        assert H == self.img_size[0] and W == self.img_size[1], \
            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."


        patches = []
        img_new = img
        new_patch = self.norm(self.projections[0](img_new))
        new_patch = new_patch.flatten(2).transpose(1, 2)
        patches.append(new_patch)

        paddings = [(1,0), (0,1), (-1,0), (0,-1), (2,0), (0,2), (-2,0), (0,-2),
                     (1,1), (1,2), (1, -1), (-1,1), (-1,2), (-1,-1), (2,1),
                     (2,2), (2,-1)]
        for i in range(len(paddings)):
            cur_padding = paddings[i]
            img_new = torch.roll(img, shifts=cur_padding, dims=(2, 3))
            new_patch = self.norm(self.projections[i+1](img_new))
            new_patch = new_patch.flatten(2).transpose(1, 2)
            patches.append(new_patch)


        patches = torch.stack(patches)

        return  patches

    def flops(self):
        Ho, Wo = self.grid_size
        for i in range(10):
            flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
            flops += Ho * Wo * self.embed_dim
        return flops


class ShiftingTransformer(nn.Module):
    def __init__(self,*, scaling_factor, output_dir, num_variants, image_size,patch_size,num_classes,embedding_dim,heads,num_hierarchies,num_layers_per_block, mlp_mult = 4,channels = 3, dim_head = 64, qkv_bias=True,attn_drop=0.0,
                 proj_drop=0.0,stochastic_depth_drop = 0.1 , init_patch_embed_size =1, kernel_size=3, stride_size=1, padding_size=1):
        super().__init__()
        assert (image_size % patch_size) == 0, 'Image dimensions must be divisible by the patch size.'
        fmap_size = image_size // patch_size
        len_pyramid = len(num_layers_per_block)
        input_size_after_patch = image_size // init_patch_embed_size
        initial_num_blocks = fmap_size * fmap_size

        assert input_size_after_patch % patch_size == 0
        down_sample_ratio = input_size_after_patch // patch_size
        # assert len_pyramid == int(torch.log( torch.tensor([down_sample_ratio]).to(torch.float32))
        #                           / torch.log(torch.tensor([2]).to(torch.float32)) + 1)

        self.patch_size = patch_size
        self.kernel_size = kernel_size
        self.stride_size = stride_size
        self.padding_size= padding_size
        self.num_hierarchies = num_hierarchies
        hierarchies = list(reversed(range(num_hierarchies)))
        layer_heads = heads
        # layer_dims = list(map(lambda t: t * dim, mults))

        layer_dims = embedding_dim
        self.start_embedding = embedding_dim[0][0]
        self.end_embedding = embedding_dim[-1][-1]
        # dim_pairs = zip(layer_dims[:-1], layer_dims[1:])
        num_blocks = (initial_num_blocks, initial_num_blocks, initial_num_blocks , initial_num_blocks)
        if num_hierarchies == 4:
            seq_len = (initial_num_blocks, initial_num_blocks//(scaling_factor), initial_num_blocks//(scaling_factor**2), initial_num_blocks//(scaling_factor**3))
        elif num_hierarchies == 3:
            seq_len = (initial_num_blocks, initial_num_blocks//(scaling_factor), initial_num_blocks//(scaling_factor**2))
        self.end_num_patches = seq_len[-1]


        self.to_patch_embedding = PatchEmbed(img_size=image_size, patch_size=self.kernel_size, stride_size=self.stride_size, padding_size=padding_size, in_chans=3,
                                             embed_dim=self.start_embedding)
        block_repeats = cast_tuple(num_layers_per_block, num_hierarchies)
        layer_heads = cast_tuple(layer_heads, num_hierarchies)
        num_blocks = cast_tuple(num_blocks, num_hierarchies)
        seq_lens = cast_tuple(seq_len, num_hierarchies)

        dim_pairs = cast_tuple(layer_dims,num_hierarchies)
        self.layers = nn.ModuleList([])

        # print("build pyramid: ")
        # print("levels:", hierarchies, "heds: ", layer_heads, "dim paors: ", dim_pairs, "block repeats: ", block_repeats, "num blocks: ", num_blocks)
        for level, heads, (dim_in, dim_out), block_repeat, seq_len in zip(hierarchies, layer_heads, dim_pairs, block_repeats,seq_lens):
            print("level: ", level, "heads: ", heads, "dim in dim out: ", (dim_in, dim_out),
                  "block repeat: ", block_repeat,"seq len: ", seq_len)
            with open(output_dir + "/pyramid_parameters.txt", "a+") as text_file:
                text_file.write(" level: ")
                text_file.write(str(level))
                text_file.write(" heads: ")
                text_file.write(str(heads))
                text_file.write(" dim in: ")
                text_file.write(str(dim_in))
                text_file.write(" dim out: ")
                text_file.write(str(dim_out))
                text_file.write(" block repeat: ")
                text_file.write(str(block_repeat))
                text_file.write(" seq len: ")
                text_file.write(str(seq_len))
                text_file.write("\n")

            is_last = level == 0
            depth = block_repeat
            is_first = level == (num_hierarchies-1)
            print("is first: ", is_first, "is last: ", is_last)
            self.layers.append(nn.ModuleList([
                Transformer(num_variants, seq_len, depth, dim_in, heads, mlp_mult, qkv_bias=qkv_bias, attn_drop=attn_drop,
                            drop_path=stochastic_depth_drop, proj_drop=proj_drop) if not is_first else
                Transformer_first(num_variants, seq_len, depth, dim_in, heads, mlp_mult, qkv_bias=qkv_bias, attn_drop=attn_drop,
                            drop_path=stochastic_depth_drop, proj_drop=proj_drop),
                Reduce_image_size(dim_in, dim_out, seq_len) if not is_last else nn.Identity()
            ]))


        self.norm = partial(nn.LayerNorm, eps=1e-6)(self.end_embedding)
        self.mlp_head = nn.Linear(self.end_embedding, num_classes)
        self.apply(_init_vit_weights)
        text_file.close()


    def forward(self, img):
        # input shape:  128, 3, 32 , 32
        # print("input size: ", img.shape)
        patches = self.to_patch_embedding(img) # 44, 128, 256,192
        # print("x after embedding: ", patches.shape)
        num_hierarchies = len(self.layers)
        for level, (transformer, reduce_image_size) in zip(reversed(range(num_hierarchies)), self.layers):
            patches = transformer(patches)

            if level > 0:
                grid_size = (int(patches.shape[1]**0.5), int(patches.shape[1]**0.5))
                patches = to_image_plane(patches, grid_size, self.patch_size)
                patches = reduce_image_size(patches)
                patches = to_patches_plane(patches, self.patch_size)
        patches = self.norm(patches)
        patches_pool = torch.mean(patches, dim=(1))
        return self.mlp_head(patches_pool)


    def flops(self):
        flops = 0
        flops += self.to_patch_embedding.flops()
        print(flops/ 1e9)
        for level, (transformer, reduce_image_size) in zip(reversed(range(len(self.layers))), self.layers):
            flops += transformer.flops()
            print(flops/ 1e9)
            if level > 0:
                flops += reduce_image_size.flops()
                print(flops/ 1e9)
        # last norm
        flops += self.end_embedding * self.end_num_patches
        print(flops/ 1e9)
        # MLP
        flops += self.end_embedding * self.num_classes
        print(flops/ 1e9)
        return flops


def to_patches_plane(x, patch_size):
    patch_size = (patch_size, patch_size)
    batch, depth, height, width = x.shape # ([128, 192, 8, 8])
    x = x.reshape( batch, depth, height*width) #  128, 192, 64
    x = x.permute(0,2,1) #  128, 64, 192
    return x

def to_image_plane(x, grid_size, patch_size):
    batch, num_patches, depth = x.shape #  128, 256, 192
    x = x.permute(0,2,1)  # 128, 192, 256
    x = x.reshape(batch, depth, grid_size[0], grid_size[1])

    return x


def _init_vit_weights(m, n: str = '', head_bias: float = 0., jax_impl: bool = False):
    """ ViT weight initialization
    * When called without n, head_bias, jax_impl args it will behave exactly the same
      as my original init for compatibility with prev hparam / downstream use cases (ie DeiT).
    * When called w/ valid n (module name) and jax_impl=True, will (hopefully) match JAX impl
    """
    if isinstance(m, nn.Linear):
        trunc_normal_(m.weight, std=.02)
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Conv2d):
        # NOTE conv was left to pytorch default in my original init
        trunc_normal_(m.weight, std=.02)
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.LayerNorm):

        nn.init.zeros_(m.bias)
        nn.init.ones_(m.weight)

#
# if __name__ == '__main__':
#     image_size = 32
#     x = torch.rand(1, 3, image_size, image_size).cuda()
#     embedding_dim = 192
#     embedding_dim = ((embedding_dim, embedding_dim), (embedding_dim, embedding_dim), (embedding_dim, embedding_dim),
#                      (embedding_dim, embedding_dim))
#     heads = 3
#     num_hierarchies = 4 # number of hierarchies
#     num_layers_per_block = (2, 4, 4, 4) #the number of transformer blocks at each heirarchy, starting from the bottom
#     model = ShiftingTransformer(
#         scaling_factor=4,
#         output_dir=".",
#         num_variants=18,
#         image_size=32,
#         patch_size=1,
#         embedding_dim=embedding_dim,
#         heads=heads,
#         num_hierarchies=num_hierarchies,  # number of hierarchies
#         num_layers_per_block=num_layers_per_block,  # the number of transformer blocks at each heirarchy, starting from the bottom
#         num_classes=10,
#         init_patch_embed_size =1,
#         kernel_size = 3,
#         stride_size = 1,
#         padding_size = 1
#     ).cuda()
#     model.eval()
#     flops = model.flops()
#     print(f"number of GFLOPs: {flops / 1e9}")
#
#
#     # througput
#     repetitions = 1000
#     total_time = 0
#     optimal_batch_size = 128 #TODO
#     dummy_input = torch.randn(optimal_batch_size, 3, 32, 32, dtype=torch.float).cuda()
#     with torch.no_grad():
#         for rep in range(repetitions):
#             starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
#             starter.record()
#             _ = model(dummy_input)
#             torch.cuda.synchronize()
#             curr_time = starter.elapsed_time(ender)/1000
#             total_time += curr_time
#     Throughput = (repetitions*optimal_batch_size)/total_time
#     print("final throughput:", Throughput)
#

# helpers # LayerNorm = partial(nn.InstanceNorm2d, affine = True) # classes Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).      This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,     the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...     See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for     changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use     'survival rate' as the argument. # work with diff dim tensors, not just 2D ConvNets # binarize Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks). # Cut & paste from PyTorch official master until it's in a few official releases - RW # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf # Computes standard normal cumulative distribution function # Values are generated by using a truncated uniform distribution and # then using the inverse CDF for the normal distribution. # Get upper and lower cdf values # Uniformly fill tensor with values from [l, u], then translate to # [2l-1, 2u-1]. # Use inverse cdf transform for normal distribution to get truncated # standard normal # Transform to proper mean, std # Clamp to ensure it's in the proper range # type: (Tensor, float, float, float, float) -> Tensor Fills the input Tensor with values drawn from a truncated     normal distribution. The values are effectively drawn from the     normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`     with values outside :math:`[a, b]` redrawn until they are within     the bounds. The method used for generating the random values works     best when :math:`a \leq \text{mean} \leq b`.     Args:         tensor: an n-dimensional `torch.Tensor`         mean: the mean of the normal distribution         std: the standard deviation of the normal distribution         a: the minimum cutoff value         b: the maximum cutoff value     Examples:         >>> w = torch.empty(3, 5)         >>> nn.init.trunc_normal_(w) # constant is stddev of standard normal truncated to (-2, 2) MLP as used in Vision Transformer, MLP-Mixer and related networks # torch.Size([num_variants,batch_size,  num_patches, embedding_dim]) # torch.Size([batch, #heads, #patches, embedding]) # torch.Size([batch, #patches, embedding]) # N - num tokens # torch.Size([num_variants,batch_size,  num_patches, embedding_dim]) # torch.Size([batch, #varaints , #patches, embedding]) # torch.Size([batch, #heads, #patches,1, embedding]) # k:  torch.Size([batch,  #patches,  #varaints, embedding])            v:  torch.Size([batch, 1 ,   #patches,  #varaints, embedding]) # attn:  torch.Size([batch, #heads, #patches, 1, #varaints ]) # N - num tokens # 44*128, 32, 32, 192 #TODO: layernorm on embedding dim? # 44*128, 192, 32, 32 # 128, 44, 64, 16, 192 #x = x.unsqueeze(dim=0) #patches = torch.cat((x, patches[1:])) # assert  num_blocks == self.num_blocks # attn # norm # mlp # norm 2D Image to Patch Embedding # original x patches # self.proj1 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj2 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj3 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj4 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj5 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj6 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj7 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj8 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj9 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj10 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj11 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj12 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj13 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj14 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj15 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj16 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj17 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # self.proj18 = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride_size, padding=padding_size) # assert len_pyramid == int(torch.log( torch.tensor([down_sample_ratio]).to(torch.float32)) #                           / torch.log(torch.tensor([2]).to(torch.float32)) + 1) # layer_dims = list(map(lambda t: t * dim, mults)) # dim_pairs = zip(layer_dims[:-1], layer_dims[1:]) # print("build pyramid: ") # print("levels:", hierarchies, "heds: ", layer_heads, "dim paors: ", dim_pairs, "block repeats: ", block_repeats, "num blocks: ", num_blocks) # input shape:  128, 3, 32 , 32 # print("input size: ", img.shape) # 44, 128, 256,192 # print("x after embedding: ", patches.shape) # last norm # MLP # ([128, 192, 8, 8]) #  128, 192, 64 #  128, 64, 192 #  128, 256, 192 # 128, 192, 256 ViT weight initialization     * When called without n, head_bias, jax_impl args it will behave exactly the same       as my original init for compatibility with prev hparam / downstream use cases (ie DeiT).     * When called w/ valid n (module name) and jax_impl=True, will (hopefully) match JAX impl # NOTE conv was left to pytorch default in my original init # # if __name__ == '__main__': #     image_size = 32 #     x = torch.rand(1, 3, image_size, image_size).cuda() #     embedding_dim = 192 #     embedding_dim = ((embedding_dim, embedding_dim), (embedding_dim, embedding_dim), (embedding_dim, embedding_dim), #                      (embedding_dim, embedding_dim)) #     heads = 3 #     num_hierarchies = 4 # number of hierarchies #     num_layers_per_block = (2, 4, 4, 4) #the number of transformer blocks at each heirarchy, starting from the bottom #     model = ShiftingTransformer( #         scaling_factor=4, #         output_dir=".", #         num_variants=18, #         image_size=32, #         patch_size=1, #         embedding_dim=embedding_dim, #         heads=heads, #         num_hierarchies=num_hierarchies,  # number of hierarchies #         num_layers_per_block=num_layers_per_block,  # the number of transformer blocks at each heirarchy, starting from the bottom #         num_classes=10, #         init_patch_embed_size =1, #         kernel_size = 3, #         stride_size = 1, #         padding_size = 1 #     ).cuda() #     model.eval() #     flops = model.flops() #     print(f"number of GFLOPs: {flops / 1e9}") # # #     # througput #     repetitions = 1000 #     total_time = 0 #     optimal_batch_size = 128 #TODO #     dummy_input = torch.randn(optimal_batch_size, 3, 32, 32, dtype=torch.float).cuda() #     with torch.no_grad(): #         for rep in range(repetitions): #             starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) #             starter.record() #             _ = model(dummy_input) #             torch.cuda.synchronize() #             curr_time = starter.elapsed_time(ender)/1000 #             total_time += curr_time #     Throughput = (repetitions*optimal_batch_size)/total_time #     print("final throughput:", Throughput) #