import torch
import torch.nn as nn
from timm.models.layers import DropPath
from .models.modules import CrossAttention


class MultilevelInteractionBlock(nn.Module):
    r""" Multilevel Interaction Block. 
    
    Args:
        dim (int): Number of low-level feature channels.
        dim1, dim2 (int): Number of high-level feature channels.
        embed_dim (int): Dimension for attention.
        num_heads (int): Number of attention heads.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
    """

    def __init__(self, dim, dim1, dim2=None, embed_dim=384, num_heads=6, mlp_ratio=3., qkv_bias=False, qk_scale=None,
                 drop=0., attn_drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
        super(MultilevelInteractionBlock, self).__init__()
        self.interact1 = CrossAttention(dim1=dim, dim2=dim1, dim=embed_dim, num_heads=num_heads, qkv_bias=qkv_bias,
                                        qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
        self.norm0 = norm_layer(dim)
        self.norm1 = norm_layer(dim1)
        self.dim = dim
        self.dim2 = dim2
        self.mlp_ratio = mlp_ratio
        if self.dim2:
            self.interact2 = CrossAttention(dim1=dim, dim2=dim2, dim=embed_dim, num_heads=num_heads, qkv_bias=qkv_bias,
                                            qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
            self.norm2 = norm_layer(dim2)

        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

        self.norm = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, dim * mlp_ratio),
            act_layer(),
            nn.Linear(dim * mlp_ratio, dim),
        )

    def forward(self, fea, fea_1, fea_2=None):
        fea = self.norm0(fea)
        fea_1 = self.norm1(fea_1)
        fea_1 = self.interact1(fea, fea_1)
        if self.dim2:
            fea_2 = self.norm2(fea_2)
            fea_2 = self.interact2(fea, fea_2)
        fea = fea + fea_1
        if self.dim2:
            fea = fea + fea_2
        fea = fea + self.drop_path(self.mlp(self.norm(fea)))
        return fea

    def flops(self, N1, N2, N3=None):
        flops = 0
        flops += self.interact1.flops(N1, N2)
        if N3:
            flops += self.interact2.flops(N1, N3)
        flops += self.dim * N1
        flops += 2 * N1 * self.dim * self.dim * self.mlp_ratio
        return flops


if __name__ == '__main__':
    # Test
    model = MultilevelInteractionBlock(dim1=96, dim2=192, dim3=384)
    model.cuda()
    f = []
    f.append(torch.randn((1, 3136, 96)).cuda())
    f.append(torch.randn((1, 784, 192)).cuda())
    f.append(torch.randn((1, 196, 384)).cuda())
    y = model(f[0], f[1], f[2])
    print(y.shape)