from .config import VITConfig
from transformers import PreTrainedModel
import torch
from torch import nn
import os
from einops import rearrange,repeat
from einops.layers.torch import Rearrange
import numpy as np
import torch.nn.functional as F

class shift_patch(nn.Module):
    def __init__(self,patch_size=16):
        super().__init__()        
        self.patch_size = patch_size // 2
    def forward(self, x):
        bs,_, _, chanel = x.size()
        shifts = ((-self.patch_size, self.patch_size, -self.patch_size, self.patch_size), 
                    (self.patch_size, -self.patch_size, -self.patch_size, self.patch_size),
                    (-self.patch_size, self.patch_size, self.patch_size, -self.patch_size), 
                    (self.patch_size, -self.patch_size, self.patch_size, -self.patch_size))
        shifted_x = list(map(lambda shift: F.pad(x, shift), shifts))
        return shifted_x
    
def posemb_sincos_2d(h, w, dim, temperature: int = 10000, dtype = torch.float32):
    y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
    assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sincos emb"
    omega = torch.arange(dim // 4) / (dim // 4 - 1)
    omega = 1.0 / (temperature ** omega)

    y = y.flatten()[:, None] * omega[None, :]
    x = x.flatten()[:, None] * omega[None, :]
    pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=1)
    return pe.type(dtype)

class patch_embedding(nn.Module):
    def __init__(self,height=224,width=224,n_channels=3,patch_size=16,batch_size=2,dim=512):
        super().__init__()
        
        assert height%patch_size==0 and width%patch_size==0 ,"Height and Width should be multiples of patch size wich is {0}".format(patch_size)
        #self.class_token = nn.Parameter(torch.randn(1, 1, dim))
        self.pos_embedding = posemb_sincos_2d(
            h = height // patch_size,
            w = width // patch_size,
            dim = dim,
        ) 
        self.patch_size = patch_size
        self.n_patchs = height*width//(patch_size**2)
        #self.embedding = torch.nn.Parameter(torch.randn(1,self.n_patchs+1, dim))
        self.projection = nn.Sequential(nn.LayerNorm(patch_size*patch_size*n_channels*5),
                                        nn.Linear(patch_size*patch_size*n_channels*5,dim),
                                        nn.LayerNorm(dim)
                                       )
        self.shift = shift_patch(patch_size=patch_size)
        self.patch = nn.Sequential(
            Rearrange("b c (h p1) (w p2)  -> b (h w)  (p1 p2 c)", p1 = patch_size, p2 = patch_size),
        )
    def forward(self, x):
        #x bs,h,w,c
        embedding = self.pos_embedding.to(x.device)
        left_up,right_up,left_down,right_down = self.shift(x)
        x = torch.cat([x,left_up,right_up,left_down,right_down],dim=1)
        #first we resize the inputs Bs,Num_patchs,*
        x = self.patch(x)
        #projection on the dim of model
        x = self.projection(x)
        outputs = x + embedding
        return outputs
        
class LocallySelfAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super(LocallySelfAttention, self).__init__()
        assert embed_dim % num_heads == 0, "Embedding dimension must be 0 modulo number of heads."
        
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.norm = nn.LayerNorm(embed_dim)
        self.tau = nn.Parameter(torch.sqrt(torch.tensor(self.head_dim)))
        self.qkv_proj = nn.Linear(embed_dim, 3*embed_dim,bias = False)
        self.o_proj = nn.Linear(embed_dim, embed_dim)

        self._reset_parameters()

    def _reset_parameters(self):
        # Original Transformer initialization, see PyTorch documentation
        nn.init.xavier_uniform_(self.qkv_proj.weight)
        nn.init.xavier_uniform_(self.o_proj.weight)
        self.o_proj.bias.data.fill_(0)
        
    def _masked_softmax(self, attention_scores, attention_mask):
        if attention_mask is not None:
            attention_scores = attention_scores.masked_fill(attention_mask, float('-inf'))
        return F.softmax(attention_scores, dim=-1)

    def forward(self,x, need_weights=True, attn_mask=None):
        x = self.norm(x)
        batch_size, seq_length, _ = x.size()
        qkv = self.qkv_proj(x).chunk(3, dim = -1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.num_heads), qkv)
        attention_scores = torch.matmul(q, k.transpose(-1, -2)) / self.tau
        attention_scores = self._masked_softmax(attention_scores, attn_mask)
        attention_output = torch.matmul(attention_scores, v)
        attention_output = rearrange(attention_output, 'b h n d -> b n (h d)')
        return self.o_proj(attention_output)

class feedforward(nn.Module):
    def __init__(self,embed_dim=512,ff_hidden_dim=1024,dropout_rate = 0.1):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.LayerNorm(embed_dim),
            nn.Linear(embed_dim, ff_hidden_dim),
            nn.GELU(),
            nn.Dropout(p=dropout_rate),
            nn.Linear(ff_hidden_dim, embed_dim),
            nn.GELU(),
            nn.Dropout(p=dropout_rate),
        )
    def forward(self, x):
        outputs = self.mlp(x)
        return outputs
        
class Transformer(nn.Module):
    def __init__(self, embed_dim=512, depth=4, heads=2, ff_hidden_dim=1024, dropout_rate = 0.1):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                LocallySelfAttention(embed_dim,heads),
                feedforward(embed_dim,ff_hidden_dim, dropout_rate)
            ]))
    def forward(self, x,mask=None):
        for attn, ff in self.layers:
            x = attn(x,attn_mask = mask) + x
            x = ff(x) + x
        return x
        
class vit(nn.Module):
    def __init__(self,height=224,width=224,n_channels=3,patch_size=16,batch_size=2,dim=512,n_head=2,feed_forward=1024,num_blocks=4,num_classes=1):
        super().__init__()
        self.embedding = patch_embedding(height,width,n_channels,patch_size,batch_size,dim)
        self.n_patchs = height*width//(patch_size**2)
        # Create a diagonal attention mask
        self.diag_attn_mask = torch.eye(self.n_patchs, dtype=torch.bool)
        self.transformer_encoder = Transformer(dim,num_blocks,n_head,feed_forward)
        self.mlp_block =  nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, num_classes),
        )
    def forward(self, inputs):
        #x bs,h,w,c
        device_ = inputs.device
        x = self.embedding(inputs)
        x = self.transformer_encoder(x,mask=self.diag_attn_mask.to(device_))
        outputs = self.mlp_block(x.mean(dim=1))
        return outputs
        
        
class VITModelForImageClassification(PreTrainedModel):
    config_class = VITConfig

    def __init__(self, config : VITConfig):
        super().__init__(config)
        self.model = vit(config.height,config.width,
                         config.n_channels,config.patch_size,
                         config.batch_size,config.dim,
                         config.n_head,config.feed_forward,
                         config.num_blocks,config.num_classes)

    def forward(self, img,labels = None):
        logits = self.model(img)
        if labels is not None:
            loss_fnc = nn.CrossEntropyLoss()
            loss = loss_fnc(logits, labels)
            return {"loss": loss, "logits": logits}
        return {"logits": logits}