from collections import OrderedDict
from typing import Tuple, Union

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
class LayerNorm(nn.LayerNorm):
    """Subclass torch's LayerNorm to handle fp16."""

    def forward(self, x: torch.Tensor):
        orig_type = x.dtype
        ret = super().forward(x.type(torch.float32))
        return ret.type(orig_type)
class QuickGELU(nn.Module):
    def forward(self, x: torch.Tensor):
        return x * torch.sigmoid(1.702 * x)
class LightTransformer(nn.Module):
    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_head)
        self.ln_1 = LayerNorm(d_model)
        self.mlp = nn.Sequential(OrderedDict([
            ("c_fc", nn.Linear(d_model, d_model * 4)),
            ("gelu", QuickGELU()),
            ("c_proj", nn.Linear(d_model * 4, d_model))
        ]))
        self.ln_2 = LayerNorm(d_model)
        self.attn_mask = attn_mask

    def attention(self, x: torch.Tensor):
        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]

    def forward(self, x: torch.Tensor):
        x = x + self.attention(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x
class SPA_prompt_injection(nn.Module):
    def __init__(self, prompt_size,in_img_size):
        super().__init__()
        self.pro_in = nn.Linear(prompt_size, in_img_size*in_img_size)
        self.pro_out = nn.Linear(in_img_size*in_img_size, prompt_size)
    def forward(self, x, prompt):
        prompt = torch.mean(prompt, dim=0).repeat(x.size(0),1,1)
        x = self.pro_in(x)
        # prompt = torch.softmax(prompt, dim=1)
        x = torch.bmm(x,prompt)
        x = self.pro_out(x)
        return x
class SPE_prompt_injection(nn.Module):
    def __init__(self, prompt_size,in_chans):
        super().__init__()
        self.pro_in = nn.Linear(prompt_size, in_chans)
        self.pro_out = nn.Linear(in_chans, prompt_size)
    def forward(self, x, prompt):
        prompt = torch.mean(prompt, dim=0).repeat(x.size(0),1,1)
        x = self.pro_in(x)
        # prompt = torch.softmax(prompt, dim=1)
        x = torch.bmm(x, prompt)
        x = self.pro_out(x)
        return x