import torch
from torch import nn,einsum
import torch.nn.functional as F

from einops import rearrange, repeat
from einops.layers.torch import Rearrange

def pair(t):
    return t if isinstance(t, tuple) else (t, t)

class PreNorm(nn.Module):
    def __init__(self, dim: int, fn):  # dim: input demonsion, fn: FeedForward or MSA
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn

    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)
    
    def fix_parameters(self):
        for i in self.parameters():
            i.requires_grad = False

    def release_parameters(self):
        for i in self.parameters():
            i.requires_grad = True


class FeedForward(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, dropout: float = 0.):  # dim: input dimension
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.net(x)


class Attention(nn.Module):
    # dim: dimension of one token
    def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
        super().__init__()
        inner_dim = dim_head * heads
        project_out = not (heads == 1 and dim == dim_head)

        self.heads = heads
        self.scale = dim_head ** -0.5

        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout),
        ) if project_out else nn.Identity()



    def forward(self, x):
        b, n, _, h = *x.shape, self.heads
        # (b, n, dim) * (dim, 3*heads*head_dim) -> (b, n, 3*heads*head_dim) -> (3, b, n, heads*head_dim)
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = [rearrange(feat, 'b n (h d) -> b h n d', h=h)
                   for feat in qkv]  # (b, heads, n, head_dim)

        attn = (q @ k.transpose(-2, -1)) * self.scale  # (b, heads, n, n)
        attn = F.softmax(attn, dim=-1)

        attn = attn @ v  # (b, heads, n, head_dim)

        attn = self.to_out(rearrange(attn, 'b h n d -> b n (h d)'))

        return attn

class PatchEmbedding(nn.Module):
    def __init__(self, num_patches, dim, emb_dropout):
        super().__init__()
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
        self.dropout = nn.Dropout(emb_dropout)

    def forward(self, x):
        b, n, _ = x.shape

        cls_tokens = repeat(self.cls_token, '() n d -> b n d', b=b)
        x = torch.cat((cls_tokens, x), dim=1)
        x += self.pos_embedding[:, :(n + 1)]

        x = self.dropout(x)
        return x
    
class InputLayer(nn.Module):
    def __init__(self,
                 image_size,
                 patch_size,
                 dim,
                 channels=3,
                 emb_dropout=0.):

        super().__init__()

        image_height, image_width = pair(image_size)
        patch_height, patch_width = pair(patch_size)

        num_patches = (image_height // patch_height) * (image_width //
                                                        patch_width)
        patch_dim = channels * patch_height * patch_width
        # print(patch_dim)
        self.params = nn.ModuleList([
            nn.Sequential(
            Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)",
                      p1=patch_height,
                      p2=patch_width), nn.Linear(patch_dim, dim)),
            PatchEmbedding(num_patches, dim, emb_dropout)
        ])

    def forward(self, img):
        # print(img.shape)
        x = self.params[0](img)
        x = self.params[1](x)
        return x

class OutputLayer(nn.Module):
    def __init__(self, num_classes, dim, pool="cls"):
        super().__init__()
        assert pool in {"cls", "mean"}
        self.params = nn.ModuleList()
        self.params.append(
            nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, num_classes))
            )
        self.pool = pool

    def forward(self, x):
        x = x.mean(dim=1) if self.pool == "mean" else x[:, 0]
        return self.params[0](x)
    
class ViTLayer(nn.Module):
    def __init__(self, dim, heads, dim_head, mlp_dim, dropout=0.):
        super().__init__()
        self.params = nn.ModuleList([
            PreNorm(dim,
                Attention(dim,
                          heads=heads,
                          dim_head=dim_head,
                          dropout=dropout)),
            PreNorm(dim, FeedForward(dim, mlp_dim,
                                     dropout=dropout))
        ])

    def forward(self, x):
        attn, ff = self.params
        x = attn(x) + x
        x = ff(x) + x
        return x

class ViT(nn.Module):
    def __init__(self, dim=1024, patch_size=16, input_size=(32, 32), output_size=10):
        super().__init__()
        self.patch_size = patch_size
        self.dim = dim
        self.input_size = input_size
        self.output_size = output_size
        self.layers = [
            InputLayer(self.input_size[1], self.patch_size, self.dim,
                      emb_dropout=0.3)
        ]
        for i in range(5):
            self.layers.append(
                ViTLayer(self.dim, heads=6, dim_head=64, mlp_dim=2048)
            )
        self.layers.append(
            OutputLayer(self.output_size, dim=self.dim)
        )
        self.layers = nn.Sequential(*self.layers)
        
    def forward(self, x):
        return self.layers(x)
        
def build_model(args):
    model = ViT(input_size=args.image_shape, output_size=args.num_classes)
    setattr(model, "losses", lambda x, y: {"cls_loss": nn.CrossEntropyLoss()(x, y)})
    return model