import torch
from torch import nn
from einops import rearrange
from torch import nn, einsum
from einops import rearrange
from mmseg.models.builder import MODELS
import math
import torch
from torch import nn as nn
from mmseg.models.builder import MODELS
from timm.layers import DropPath, trunc_normal_
from typing import List
from timm.layers import create_act_layer
from functools import partial
import torch.nn.functional as F


import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from timm.layers import CondConv2d, get_condconv_initializer, create_conv2d, DropPath, get_norm_act_layer


class LoRaMLP(nn.Module):
    def __init__(self, in_dim, out_dim, rank_dim=8):
        super().__init__()
        self.loramlp = nn.Sequential(
            nn.Linear(in_dim, rank_dim, bias=False),
            nn.Linear(rank_dim, out_dim, bias=False),
        )

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


class CrossAttention(nn.Module):
    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, rank_dim=None):
        super().__init__()
        inner_dim = dim_head * heads  # 512
        context_dim = query_dim if context_dim is None else context_dim

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

        if not rank_dim:
            self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
            self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
            self.to_v = nn.Linear(context_dim, inner_dim, bias=False)

            self.to_out = nn.Linear(inner_dim, query_dim, bias=False)
        else:
            self.to_q = LoRaMLP(query_dim, inner_dim, rank_dim=rank_dim)
            self.to_k = LoRaMLP(context_dim, inner_dim, rank_dim=rank_dim)
            self.to_v = LoRaMLP(context_dim, inner_dim, rank_dim=rank_dim)

            self.to_out = LoRaMLP(inner_dim, query_dim, rank_dim=rank_dim)

    def forward(self, x, context):
        h = self.heads

        q = self.to_q(x)
        k = self.to_k(context)
        v = self.to_v(context)

        q, k, v = map(lambda t: rearrange(
            t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))

        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale

        attn = sim.softmax(dim=-1)

        out = einsum('b i j, b j d -> b i d', attn, v)
        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)

        return self.to_out(out)


def num_groups(group_size, channels):
    if not group_size:
        return 1
    else:
        assert channels % group_size == 0
        return channels // group_size


def _init_weight_goog(m, n='', fix_group_fanout=True):
    if isinstance(m, CondConv2d):
        fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
        if fix_group_fanout:
            fan_out //= m.groups
        init_weight_fn = get_condconv_initializer(
            lambda w: nn.init.normal_(w, 0, math.sqrt(2.0 / fan_out)), m.num_experts, m.weight_shape)
        init_weight_fn(m.weight)
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Conv2d):
        fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
        if fix_group_fanout:
            fan_out //= m.groups
        nn.init.normal_(m.weight, 0, math.sqrt(2.0 / fan_out))
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.BatchNorm2d):
        nn.init.ones_(m.weight)
        nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Linear):
        fan_out = m.weight.size(0)
        fan_in = 0
        if 'routing_fn' in n:
            fan_in = m.weight.size(1)
        init_range = 1.0 / math.sqrt(fan_in + fan_out)
        nn.init.uniform_(m.weight, -init_range, init_range)
        if m.bias is not None:
            nn.init.zeros_(m.bias)


class DepthwiseSeparableConv(nn.Module):
    def __init__(
            self, in_chs, out_chs, dw_kernel_size=3, stride=1, dilation=1, group_size=1, pad_type='',
            noskip=False, pw_kernel_size=1, pw_act=False, act_layer=nn.ReLU, norm_layer=nn.BatchNorm2d,
            se_layer=None, drop_path_rate=0.):
        super(DepthwiseSeparableConv, self).__init__()
        norm_act_layer = get_norm_act_layer(norm_layer)
        groups = num_groups(group_size, in_chs)
        self.has_skip = (stride == 1 and in_chs == out_chs) and not noskip
        self.has_pw_act = pw_act

        self.conv_dw = create_conv2d(
            in_chs, in_chs, dw_kernel_size, stride=stride, dilation=dilation, padding=pad_type, groups=groups)
        self.bn1 = norm_act_layer(in_chs, inplace=True)

        self.se = se_layer(
            in_chs, act_layer=act_layer) if se_layer else nn.Identity()

        self.conv_pw = create_conv2d(
            in_chs, out_chs, pw_kernel_size, padding=pad_type)
        self.bn2 = norm_act_layer(
            out_chs, inplace=True, apply_act=self.has_pw_act)
        self.drop_path = DropPath(
            drop_path_rate) if drop_path_rate else nn.Identity()

    def feature_info(self, location):
        if location == 'expansion':
            return dict(module='conv_pw', hook_type='forward_pre', num_chs=self.conv_pw.in_channels)
        else:
            return dict(module='', hook_type='', num_chs=self.conv_pw.out_channels)

    def forward(self, x):
        shortcut = x
        x = self.conv_dw(x)
        x = self.bn1(x)
        x = self.se(x)
        x = self.conv_pw(x)
        x = self.bn2(x)
        if self.has_skip:
            x = self.drop_path(x) + shortcut
        return x


class PMAAConvBlock(nn.Module):
    def __init__(self, in_channels=3, hidden_channels=256, depth=4, norm=nn.BatchNorm2d, act=nn.ReLU, return_multi_feats=False, return_last_feature=True, has_stem=True, has_block=True):
        super().__init__()
        self.return_last_feature = return_last_feature
        self.depth = depth
        self.has_stem = has_stem
        self.return_multi_feats = return_multi_feats

        self.proj_1x1 = DepthwiseSeparableConv(
            in_channels, hidden_channels, dw_kernel_size=1, norm_layer=norm, act_layer=act)

        self.spp_dw = nn.ModuleList()

        if has_stem:
            self.spp_dw.append(
                DepthwiseSeparableConv(hidden_channels, hidden_channels, dw_kernel_size=3,
                                       stride=1, group_size=hidden_channels, pad_type="same")
            )
        else:
            self.spp_dw.append(nn.Identity())

        if has_block:
            for _ in range(self.depth):
                self.spp_dw.append(
                    DepthwiseSeparableConv(
                        hidden_channels, hidden_channels, dw_kernel_size=3, stride=2, group_size=hidden_channels
                    )
                )
        else:
            for _ in range(self.depth):
                self.spp_dw.append(
                    nn.MaxPool2d(kernel_size=2, stride=2)
                )
        self._init_weights()

    def forward(self, x):
        B, C, H, W = x.shape
        output1 = self.proj_1x1(x)
        output = [self.spp_dw[0](output1)]

        for k in range(1, self.depth+1):
            out_k = self.spp_dw[k](output[-1])
            output.append(out_k)

        if self.return_multi_feats:
            return output[1:]
        else:
            if self.return_last_feature:
                return output[-1]
            global_f = torch.zeros(
                output[-1].shape, requires_grad=True, device=output1.device)
            for fea in output:
                global_f = global_f + F.adaptive_avg_pool2d(
                    fea, output_size=output[-1].shape[-2:]
                )
            return global_f

    def _init_weights(self):
        init_fn = _init_weight_goog
        for n, m in self.named_modules():
            init_fn(m, n)


class ConvnextInteractiveModule(nn.Module):
    def __init__(self, emd_dim=1024, context_dim=256, rank_dim=None):
        super().__init__()
        self.attn = CrossAttention(emd_dim, context_dim, rank_dim=rank_dim)

    def forward(self, x, cache, index):
        # x: 1024 2 1024
        if isinstance(cache, list) or isinstance(cache, tuple):
            # len(cache) 4 cache[4]-23
            # 0-5->0 6-11 -> 1 12-17->2 18-23->3
            cache = cache[index]
        cache = F.interpolate(
            cache, (int(math.sqrt(x.shape[0])), int(math.sqrt(x.shape[0]))), mode="bilinear", align_corners=False
        )
        cache = cache.flatten(2)  # B C N
        cache = cache.permute(2, 0, 1)  # N B C

        # Reshape: batch first
        x = x.permute(1, 0, 2)  # B N C
        cache = cache.permute(1, 0, 2)  # B N C
        return (x + self.attn(x, cache)).permute(1, 0, 2)


class PMAAInteractiveModule(nn.Module):
    def __init__(self,
                 emd_dim=1024,
                 context_dim=64,
                 kernel: int = 1,
                 norm=nn.BatchNorm2d,
                 local_groups=32,
                 global_groups=2,
                 return_multi_feats=False,
                 ):
        super().__init__()
        self.return_multi_feats = return_multi_feats
        self.local_embedding = nn.Sequential(
            nn.Conv2d(emd_dim, emd_dim, kernel, groups=local_groups,
                      padding=int((kernel - 1) / 2), bias=False),
            norm(emd_dim)
        )
        self.global_embedding = nn.Sequential(
            nn.Conv2d(context_dim, emd_dim, kernel, groups=global_groups,
                      padding=int((kernel - 1) / 2), bias=False),
            norm(emd_dim)
        )
        self.global_act = nn.Sequential(
            nn.Conv2d(context_dim, emd_dim, kernel, groups=global_groups,
                      padding=int((kernel - 1) / 2), bias=False),
            norm(emd_dim)
        )
        self.act = nn.Sigmoid()
        self._init_weights()

    def _init_weights(self):
        init_fn = _init_weight_goog
        for n, m in self.named_modules():
            init_fn(m, n)

    def forward(self, x, cache, index):
        if isinstance(cache, list) or isinstance(cache, tuple):
            cache = cache[index]
        N, B, C = x.shape
        H = W = int(math.sqrt(N))
        # reshape x -> B, C, H, W
        x = x.permute(1, 2, 0).reshape(B, C, H, W)
        local_feat = self.local_embedding(x)  # 32
        global_act = self.global_act(cache)
        sig_act = F.interpolate(self.act(global_act), size=(H, W))  # 32

        global_feat = self.global_embedding(cache)
        global_feat = F.interpolate(global_feat, size=(H, W))  # 32

        out = local_feat * sig_act + global_feat

        return out.permute(2, 3, 0, 1).reshape(N, B, C)


class LayerNorm(nn.Module):
    r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. 
    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with 
    shape (batch_size, height, width, channels) while channels_first corresponds to inputs 
    with shape (batch_size, channels, height, width).
    """

    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))
        self.eps = eps
        self.data_format = data_format
        if self.data_format not in ["channels_last", "channels_first"]:
            raise NotImplementedError
        self.normalized_shape = (normalized_shape, )

    def forward(self, x):
        if self.data_format == "channels_last":
            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
        elif self.data_format == "channels_first":
            u = x.mean(1, keepdim=True)
            s = (x - u).pow(2).mean(1, keepdim=True)
            x = (x - u) / torch.sqrt(s + self.eps)
            x = self.weight[:, None, None] * x + self.bias[:, None, None]
            return x


class Block(nn.Module):
    r""" ConvNeXt Block. There are two equivalent implementations:
    (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
    (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
    We use (2) as we find it slightly faster in PyTorch

    Args:
        dim (int): Number of input channels.
        drop_path (float): Stochastic depth rate. Default: 0.0
        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
    """

    def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6):
        super().__init__()
        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7,
                                padding=3, groups=dim)  # depthwise conv
        self.norm = LayerNorm(dim, eps=1e-6)
        # pointwise/1x1 convs, implemented with linear layers
        self.pwconv1 = nn.Linear(dim, 4 * dim)
        self.act = nn.GELU()
        self.pwconv2 = nn.Linear(4 * dim, dim)
        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)),
                                  requires_grad=True) if layer_scale_init_value > 0 else None
        self.drop_path = DropPath(
            drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x):
        input = x
        x = self.dwconv(x)
        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        if self.gamma is not None:
            x = self.gamma * x
        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)

        x = input + self.drop_path(x)
        return x


class ConvNeXt(nn.Module):
    r""" ConvNeXt
        A PyTorch impl of : `A ConvNet for the 2020s`  -
          https://arxiv.org/pdf/2201.03545.pdf

    Args:
        in_chans (int): Number of input image channels. Default: 3
        num_classes (int): Number of classes for classification head. Default: 1000
        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
        drop_path_rate (float): Stochastic depth rate. Default: 0.
        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
    """

    def __init__(self, in_chans=3, depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
                 drop_path_rate=0., layer_scale_init_value=1e-6, out_indices=[0, 1, 2, 3],
                 return_multi_feats=False,
                 return_last_feature=True
                 ):
        super().__init__()
        self.return_last_feature = return_last_feature
        self.return_multi_feats = return_multi_feats

        # stem and 3 intermediate downsampling conv layers
        self.downsample_layers = nn.ModuleList()
        stem = nn.Sequential(
            nn.Conv2d(in_chans, dims[0], kernel_size=2, stride=2),
            LayerNorm(dims[0], eps=1e-6, data_format="channels_first")
        )
        self.downsample_layers.append(stem)
        for i in range(3):
            downsample_layer = nn.Sequential(
                LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
                nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2),
            )
            self.downsample_layers.append(downsample_layer)

        # 4 feature resolution stages, each consisting of multiple residual blocks
        self.stages = nn.ModuleList()
        dp_rates = [x.item()
                    for x in torch.linspace(0, drop_path_rate, sum(depths))]
        cur = 0
        for i in range(4):
            stage = nn.Sequential(
                *[Block(dim=dims[i], drop_path=dp_rates[cur + j],
                        layer_scale_init_value=layer_scale_init_value) for j in range(depths[i])]
            )
            self.stages.append(stage)
            cur += depths[i]

        self.out_indices = out_indices

        norm_layer = partial(LayerNorm, eps=1e-6, data_format="channels_first")
        for i_layer in range(4):
            layer = norm_layer(dims[i_layer])
            layer_name = f'norm{i_layer}'
            self.add_module(layer_name, layer)

        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, (nn.Conv2d, nn.Linear)):
            trunc_normal_(m.weight, std=.02)
            nn.init.constant_(m.bias, 0)

    def init_weights(self, pretrained=None):
        """Initialize the weights in backbone.
        Args:
            pretrained (str, optional): Path to pre-trained weights.
                Defaults to None.
        """

        def _init_weights(m):
            if isinstance(m, nn.Linear):
                trunc_normal_(m.weight, std=.02)
                if isinstance(m, nn.Linear) and m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.LayerNorm):
                nn.init.constant_(m.bias, 0)
                nn.init.constant_(m.weight, 1.0)

        if isinstance(pretrained, str):
            self.apply(_init_weights)
            # logger = get_root_logger()
            # load_checkpoint(self, pretrained, strict=False, logger=logger)
        elif pretrained is None:
            self.apply(_init_weights)
        else:
            raise TypeError('pretrained must be a str or None')

    def forward_features(self, x):
        outs = []
        for i in range(4):
            x = self.downsample_layers[i](x)
            x = self.stages[i](x)
            if i in self.out_indices:
                norm_layer = getattr(self, f'norm{i}')
                x_out = norm_layer(x)
                outs.append(x_out)
        if self.return_multi_feats:
            return tuple(outs)
        if self.return_last_feature:
            return outs[-1]
        global_f = torch.zeros(
            outs[-1].shape, requires_grad=True, device=outs[-1].device)
        for fea in outs:
            global_f = global_f + F.adaptive_avg_pool2d(
                fea, output_size=outs[-1].shape[-2:]
            )
        return global_f

    def forward(self, x):
        x = self.forward_features(x)
        return x


class NoAdaptingModule(nn.Identity):
    def __init__(self):
        super().__init__()

    def forward(self, x, cache, layer):
        return x


@MODELS.register_module()
class CloudAdapter(nn.Module):
    def __init__(self,
                 cnn_type="convnext",  # convnext or mobilenet
                 int_type="convnext",  # cross_attention or
                 # 共同的参数 start
                 emd_dim=1024,
                 num_layers=24,

                 # 先判断是否返回多特征，之后再判断是否进行特征融合
                 return_multi_feats=True,
                 return_last_feature=False,

                 # 共同的参数 end

                 # pmaa 提取单个特征 or 多尺寸特征 start
                 hidden_channels=256,
                 depth=4,
                 norm=nn.BatchNorm2d,
                 act=nn.ReLU,
                 # pmaa 提取单个特征 or 多尺寸特征 end

                 # pmaa net start
                 local_groups=1,
                 global_groups=1,
                 # pmaa net end

                 # convnext 提取单个特征 or 多尺寸特征 start
                 context_dim=256,
                 rank_dim=None,
                 # convnext 提取单个特征 or 多尺寸特征 end,
                 has_stem=True,
                 has_block=True,
                 ):
        super().__init__()
        self.cnn = nn.Identity()
        self.net = nn.Identity()
        if cnn_type == "pmaa":
            self.cnn = PMAAConvBlock(
                hidden_channels=hidden_channels,
                depth=depth,
                norm=norm,
                act=act,
                return_multi_feats=return_multi_feats,
                return_last_feature=return_last_feature,
                has_stem=has_stem,
                has_block=has_block
            )
        elif cnn_type == "convnext":
            self.cnn = ConvNeXt(depths=[1]*4,
                                dims=[context_dim]*4,
                                return_multi_feats=return_multi_feats,
                                return_last_feature=return_last_feature
                                )

        else:
            raise ValueError(
                f"cnn_type must in ['convnext','pmaa'],but got {cnn_type}")

        if int_type == "convnext":
            self.net = nn.ModuleList(
                ConvnextInteractiveModule(emd_dim, context_dim, rank_dim)
                for _ in range(num_layers)
            )
        elif int_type == "pmaa":
            self.net = nn.ModuleList(
                PMAAInteractiveModule(
                    emd_dim, context_dim, local_groups=local_groups, global_groups=global_groups)
                for _ in range(num_layers)
            )

        elif int_type == "no_adapting":
            self.net = nn.ModuleList(
                NoAdaptingModule() for _ in range(num_layers)
            )
        else:
            raise ValueError(
                f"int_type must in ['convnext','pmaa'],but got {int_type}")

    def forward(self, feats, layer, batch_first=True, has_cls_token=True, cache=None):
        if batch_first:
            feats = feats.permute(1, 0, 2)  # 1025 2 1024
        if has_cls_token:
            cls_token, feats = torch.tensor_split(feats, [1], dim=0)
        # 24 // 1
        # feat: 1024 2 1024
        feats = self.net[layer].forward(
            feats, cache, layer//(len(self.net) // 4))

        if has_cls_token:
            feats = torch.cat([cls_token, feats], dim=0)
        if batch_first:
            feats = feats.permute(1, 0, 2)
        return feats