import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class Attention(nn.Module):
    """
    Compute 'Scaled Dot Product Attention
    """

    def __init__(self, p=0.1):
        super(Attention, self).__init__()
        self.dropout = nn.Dropout(p=p)

    def forward(self, query, key, value):
        scores = torch.matmul(query, key.transpose(-2, -1)
                              ) / math.sqrt(query.size(-1))
        p_attn = F.softmax(scores, dim=-1)
        p_attn = self.dropout(p_attn)
        p_val = torch.matmul(p_attn, value)
        return p_val, p_attn


class SWMHSA_depthGlobalWindowConcatLN_qkFlow_reweightFlow(nn.Module):
    def __init__(self, token_size, window_size, kernel_size, d_model, flow_dModel, head, p=0.1):
        super(SWMHSA_depthGlobalWindowConcatLN_qkFlow_reweightFlow, self).__init__()
        self.h, self.w = token_size
        self.head = head
        self.window_size = window_size
        self.d_model = d_model
        self.flow_dModel = flow_dModel
        in_channels = d_model + flow_dModel
        self.query_embedding = nn.Linear(in_channels, d_model)
        self.key_embedding = nn.Linear(in_channels, d_model)
        self.value_embedding = nn.Linear(d_model, d_model)
        self.output_linear = nn.Linear(d_model, d_model)
        self.attention = Attention(p)
        self.pad_l = self.pad_t = 0
        self.pad_r = (self.window_size - self.w % self.window_size) % self.window_size
        self.pad_b = (self.window_size - self.h % self.window_size) % self.window_size
        self.new_h, self.new_w = self.h + self.pad_b, self.w + self.pad_r
        self.group_h, self.group_w = self.new_h // self.window_size, self.new_w // self.window_size
        self.global_extract_v = nn.Conv2d(d_model, d_model, kernel_size=kernel_size, stride=kernel_size, padding=0,
                                          groups=d_model)
        self.global_extract_k = nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size, stride=kernel_size,
                                          padding=0,
                                          groups=in_channels)
        self.q_norm = nn.LayerNorm(d_model + flow_dModel)
        self.k_norm = nn.LayerNorm(d_model + flow_dModel)
        self.v_norm = nn.LayerNorm(d_model)
        self.reweightFlow = nn.Sequential(
            nn.Linear(in_channels, flow_dModel),
            nn.Sigmoid()
        )

    def inference(self, x, f, h, w):
        pad_r = (self.window_size - w % self.window_size) % self.window_size
        pad_b = (self.window_size - h % self.window_size) % self.window_size
        new_h, new_w = h + pad_b, w + pad_r
        group_h, group_w = new_h // self.window_size, new_w // self.window_size
        bt, n, c = x.shape
        cf = f.shape[2]
        x = x.view(bt, h, w, c)
        f = f.view(bt, h, w, cf)
        if pad_r > 0 or pad_b > 0:
            x = F.pad(x, (0, 0, self.pad_l, pad_r, self.pad_t, pad_b))
            f = F.pad(f, (0, 0, self.pad_l, pad_r, self.pad_t, pad_b))
        y = x.permute(0, 3, 1, 2)
        xf = torch.cat((x, f), dim=-1)
        flow_weights = self.reweightFlow(xf)
        f = f * flow_weights
        qk = torch.cat((x, f), dim=-1)  # [b, h, w, c]
        qk_c = qk.shape[-1]
        # generate q
        q = qk.reshape(bt, group_h, self.window_size, group_w, self.window_size, qk_c).transpose(2, 3)
        q = q.reshape(bt, group_h * group_w, self.window_size * self.window_size, qk_c)
        # generate k
        ky = qk.permute(0, 3, 1, 2)  # [b, c, h, w]
        k_global = self.global_extract_k(ky)
        k_global = k_global.permute(0, 2, 3, 1).reshape(bt, -1, qk_c).unsqueeze(1).repeat(1, group_h * group_w, 1, 1)
        k = torch.cat((q, k_global), dim=2)
        # norm q and k
        q = self.q_norm(q)
        k = self.k_norm(k)
        # generate v
        global_tokens = self.global_extract_v(y)  # [bt, c, h', w']
        global_tokens = global_tokens.permute(0, 2, 3, 1).reshape(bt, -1, c).unsqueeze(1).repeat(1,
                                                                                                 group_h * group_w,
                                                                                                 1,
                                                                                                 1)  # [bt, gh * gw, h'*w', c]
        x = x.reshape(bt, group_h, self.window_size, group_w, self.window_size, c).transpose(2,
                                                                                             3)  # [bt, gh, gw, ws, ws, c]
        x = x.reshape(bt, group_h * group_w, self.window_size * self.window_size, c)  # [bt, gh * gw, ws^2, c]
        v = torch.cat((x, global_tokens), dim=2)
        v = self.v_norm(v)
        query = self.query_embedding(q)  # [bt, self.group_h, self.group_w, self.window_size, self.window_size, c]
        key = self.key_embedding(k)
        value = self.value_embedding(v)
        query = query.reshape(bt, group_h * group_w, self.window_size * self.window_size, self.head,
                              c // self.head).permute(0, 1, 3, 2, 4)
        key = key.reshape(bt, group_h * group_w, -1, self.head,
                          c // self.head).permute(0, 1, 3, 2, 4)
        value = value.reshape(bt, group_h * group_w, -1, self.head,
                              c // self.head).permute(0, 1, 3, 2, 4)
        attn, _ = self.attention(query, key, value)
        x = attn.transpose(2, 3).reshape(bt, group_h, group_w, self.window_size, self.window_size, c)
        x = x.transpose(2, 3).reshape(bt, group_h * self.window_size, group_w * self.window_size, c)
        if pad_r > 0 or pad_b > 0:
            x = x[:, :h, :w, :].contiguous()
        x = x.reshape(bt, n, c)
        output = self.output_linear(x)
        return output

    def forward(self, x, f, t, h=0, w=0):
        if h != 0 or w != 0:
            return self.inference(x, f, h, w)
        bt, n, c = x.shape
        cf = f.shape[2]
        x = x.view(bt, self.h, self.w, c)
        f = f.view(bt, self.h, self.w, cf)
        if self.pad_r > 0 or self.pad_b > 0:
            x = F.pad(x, (0, 0, self.pad_l, self.pad_r, self.pad_t, self.pad_b))
            f = F.pad(f, (0, 0, self.pad_l, self.pad_r, self.pad_t, self.pad_b))  # [bt, cf, h, w]
        y = x.permute(0, 3, 1, 2)
        xf = torch.cat((x, f), dim=-1)
        weights = self.reweightFlow(xf)
        f = f * weights
        qk = torch.cat((x, f), dim=-1)  # [b, h, w, c]
        qk_c = qk.shape[-1]
        # generate q
        q = qk.reshape(bt, self.group_h, self.window_size, self.group_w, self.window_size, qk_c).transpose(2, 3)
        q = q.reshape(bt, self.group_h * self.group_w, self.window_size * self.window_size, qk_c)
        # generate k
        ky = qk.permute(0, 3, 1, 2)  # [b, c, h, w]
        k_global = self.global_extract_k(ky)  # [b, qk_c, h, w]
        k_global = k_global.permute(0, 2, 3, 1).reshape(bt, -1, qk_c).unsqueeze(1).repeat(1,
                                                                                          self.group_h * self.group_w,
                                                                                          1, 1)
        k = torch.cat((q, k_global), dim=2)
        # norm q and k
        q = self.q_norm(q)
        k = self.k_norm(k)
        # generate v
        global_tokens = self.global_extract_v(y)  # [bt, c, h', w']
        global_tokens = global_tokens.permute(0, 2, 3, 1).reshape(bt, -1, c).unsqueeze(1).repeat(1,
                                                                                                 self.group_h * self.group_w,
                                                                                                 1,
                                                                                                 1)  # [bt, gh * gw, h'*w', c]
        x = x.reshape(bt, self.group_h, self.window_size, self.group_w, self.window_size, c).transpose(2,
                                                                                                       3)  # [bt, gh, gw, ws, ws, c]
        x = x.reshape(bt, self.group_h * self.group_w, self.window_size * self.window_size, c)  # [bt, gh * gw, ws^2, c]
        v = torch.cat((x, global_tokens), dim=2)
        v = self.v_norm(v)
        query = self.query_embedding(q)  # [bt, self.group_h, self.group_w, self.window_size, self.window_size, c]
        key = self.key_embedding(k)
        value = self.value_embedding(v)
        query = query.reshape(bt, self.group_h * self.group_w, self.window_size * self.window_size, self.head,
                              c // self.head).permute(0, 1, 3, 2, 4)
        key = key.reshape(bt, self.group_h * self.group_w, -1, self.head,
                          c // self.head).permute(0, 1, 3, 2, 4)
        value = value.reshape(bt, self.group_h * self.group_w, -1, self.head,
                              c // self.head).permute(0, 1, 3, 2, 4)
        attn, _ = self.attention(query, key, value)
        x = attn.transpose(2, 3).reshape(bt, self.group_h, self.group_w, self.window_size, self.window_size, c)
        x = x.transpose(2, 3).reshape(bt, self.group_h * self.window_size, self.group_w * self.window_size, c)
        if self.pad_r > 0 or self.pad_b > 0:
            x = x[:, :self.h, :self.w, :].contiguous()
        x = x.reshape(bt, n, c)
        output = self.output_linear(x)
        return output