import copy
from typing import Optional

import paddle
import paddle.nn as nn
import paddle.nn.functional as F

from .extractor import BasicEncoder
from .position_encoding import build_position_encoding
from .weight_init import weight_init_


class attnLayer(nn.Layer):
    def __init__(
            self,
            d_model,
            nhead=8,
            dim_feedforward=2048,
            dropout=0.1,
            activation="relu",
            normalize_before=False,
    ):
        super().__init__()

        self.self_attn = nn.MultiHeadAttention(d_model, nhead, dropout=dropout)
        self.multihead_attn_list = nn.LayerList(
            [
                copy.deepcopy(nn.MultiHeadAttention(d_model, nhead, dropout=dropout))
                for i in range(2)
            ]
        )

        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(p=dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2_list = nn.LayerList(
            [copy.deepcopy(nn.LayerNorm(d_model)) for i in range(2)]
        )

        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(p=dropout)
        self.dropout2_list = nn.LayerList(
            [copy.deepcopy(nn.Dropout(p=dropout)) for i in range(2)]
        )
        self.dropout3 = nn.Dropout(p=dropout)

        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

    def with_pos_embed(self, tensor, pos: Optional[paddle.Tensor]):
        return tensor if pos is None else tensor + pos

    def forward_post(
            self,
            tgt,
            memory_list,
            tgt_mask=None,
            memory_mask=None,
            pos=None,
            memory_pos=None,
    ):
        q = k = self.with_pos_embed(tgt, pos)
        tgt2 = self.self_attn(
            q.transpose((1, 0, 2)),
            k.transpose((1, 0, 2)),
            value=tgt.transpose((1, 0, 2)),
            attn_mask=tgt_mask,
        )
        tgt2 = tgt2.transpose((1, 0, 2))
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)

        for memory, multihead_attn, norm2, dropout2, m_pos in zip(
                memory_list,
                self.multihead_attn_list,
                self.norm2_list,
                self.dropout2_list,
                memory_pos,
        ):
            tgt2 = multihead_attn(
                query=self.with_pos_embed(tgt, pos).transpose((1, 0, 2)),
                key=self.with_pos_embed(memory, m_pos).transpose((1, 0, 2)),
                value=memory.transpose((1, 0, 2)),
                attn_mask=memory_mask,
            ).transpose((1, 0, 2))

            tgt = tgt + dropout2(tgt2)
            tgt = norm2(tgt)

        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)

        return tgt

    def forward_pre(
            self,
            tgt,
            memory,
            tgt_mask=None,
            memory_mask=None,
            pos=None,
            memory_pos=None,
    ):
        tgt2 = self.norm1(tgt)

        q = k = self.with_pos_embed(tgt2, pos)
        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask)
        tgt = tgt + self.dropout1(tgt2)
        tgt2 = self.norm2(tgt)

        tgt2 = self.multihead_attn(
            query=self.with_pos_embed(tgt2, pos),
            key=self.with_pos_embed(memory, memory_pos),
            value=memory,
            attn_mask=memory_mask,
        )
        tgt = tgt + self.dropout2(tgt2)
        tgt2 = self.norm3(tgt)

        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
        tgt = tgt + self.dropout3(tgt2)

        return tgt

    def forward(
            self,
            tgt,
            memory_list,
            tgt_mask=None,
            memory_mask=None,
            pos=None,
            memory_pos=None,
    ):
        if self.normalize_before:
            return self.forward_pre(
                tgt,
                memory_list,
                tgt_mask,
                memory_mask,
                pos,
                memory_pos,
            )
        return self.forward_post(
            tgt,
            memory_list,
            tgt_mask,
            memory_mask,
            pos,
            memory_pos,
        )


def _get_clones(module, N):
    return nn.LayerList([copy.deepcopy(module) for i in range(N)])


def _get_activation_fn(activation):
    """Return an activation function given a string"""
    if activation == "relu":
        return F.relu
    if activation == "gelu":
        return F.gelu
    if activation == "glu":
        return F.glu
    raise RuntimeError(f"activation should be relu/gelu, not {activation}.")


class TransDecoder(nn.Layer):
    def __init__(self, num_attn_layers: int, hidden_dim: int = 128):
        super(TransDecoder, self).__init__()

        attn_layer = attnLayer(hidden_dim)
        self.layers = _get_clones(attn_layer, num_attn_layers)
        self.position_embedding = build_position_encoding(hidden_dim)

    def forward(self, image: paddle.Tensor, query_embed: paddle.Tensor):
        pos = self.position_embedding(
            paddle.ones([image.shape[0], image.shape[2], image.shape[3]], dtype="bool")
        )

        b, c, h, w = image.shape

        image = image.flatten(2).transpose(perm=[2, 0, 1])
        pos = pos.flatten(2).transpose(perm=[2, 0, 1])

        for layer in self.layers:
            query_embed = layer(query_embed, [image], pos=pos, memory_pos=[pos, pos])

        query_embed = query_embed.transpose(perm=[1, 2, 0]).reshape([b, c, h, w])

        return query_embed


class TransEncoder(nn.Layer):
    def __init__(self, num_attn_layers: int, hidden_dim: int = 128):
        super(TransEncoder, self).__init__()

        attn_layer = attnLayer(hidden_dim)
        self.layers = _get_clones(attn_layer, num_attn_layers)
        self.position_embedding = build_position_encoding(hidden_dim)

    def forward(self, image: paddle.Tensor):
        pos = self.position_embedding(
            paddle.ones([image.shape[0], image.shape[2], image.shape[3]], dtype="bool")
        )

        b, c, h, w = image.shape

        image = image.flatten(2).transpose(perm=[2, 0, 1])
        pos = pos.flatten(2).transpose(perm=[2, 0, 1])

        for layer in self.layers:
            image = layer(image, [image], pos=pos, memory_pos=[pos, pos])

        image = image.transpose(perm=[1, 2, 0]).reshape([b, c, h, w])

        return image


class FlowHead(nn.Layer):
    def __init__(self, input_dim=128, hidden_dim=256):
        super(FlowHead, self).__init__()

        self.conv1 = nn.Conv2D(input_dim, hidden_dim, 3, padding=1)
        self.conv2 = nn.Conv2D(hidden_dim, 2, 3, padding=1)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.conv2(self.relu(self.conv1(x)))


class UpdateBlock(nn.Layer):
    def __init__(self, hidden_dim: int = 128):
        super(UpdateBlock, self).__init__()

        self.flow_head = FlowHead(hidden_dim, hidden_dim=256)

        self.mask = nn.Sequential(
            nn.Conv2D(hidden_dim, 256, 3, padding=1),
            nn.ReLU(),
            nn.Conv2D(256, 64 * 9, 1, padding=0),
        )

    def forward(self, image, coords):
        mask = 0.25 * self.mask(image)
        dflow = self.flow_head(image)
        coords = coords + dflow
        return mask, coords


def coords_grid(batch, ht, wd):
    coords = paddle.meshgrid(paddle.arange(end=ht), paddle.arange(end=wd))
    coords = paddle.stack(coords[::-1], axis=0).astype(dtype="float32")
    return coords[None].tile([batch, 1, 1, 1])


def upflow8(flow, mode="bilinear"):
    new_size = 8 * flow.shape[2], 8 * flow.shape[3]
    return 8 * F.interpolate(flow, size=new_size, mode=mode, align_corners=True)


class OverlapPatchEmbed(nn.Layer):
    """Image to Patch Embedding"""

    def __init__(self, img_size=224, patch_size=7, stride=4, in_chans=3, embed_dim=768):
        super().__init__()

        img_size = img_size if isinstance(img_size, tuple) else (img_size, img_size)
        patch_size = (
            patch_size if isinstance(patch_size, tuple) else (patch_size, patch_size)
        )

        self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]
        self.num_patches = self.H * self.W

        self.proj = nn.Conv2D(
            in_chans,
            embed_dim,
            patch_size,
            stride=stride,
            padding=(patch_size[0] // 2, patch_size[1] // 2),
        )
        self.norm = nn.LayerNorm(embed_dim)

        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            weight_init_(m, "trunc_normal_", std=0.02)
        elif isinstance(m, nn.LayerNorm):
            weight_init_(m, "Constant", value=1.0)
        elif isinstance(m, nn.Conv2D):
            weight_init_(
                m.weight, "kaiming_normal_", mode="fan_out", nonlinearity="relu"
            )

    def forward(self, x):
        x = self.proj(x)
        _, _, H, W = x.shape
        x = x.flatten(2)

        perm = list(range(x.ndim))
        perm[1] = 2
        perm[2] = 1
        x = x.transpose(perm=perm)

        x = self.norm(x)

        return x, H, W


class GeoTr(nn.Layer):
    def __init__(self):
        super(GeoTr, self).__init__()

        self.hidden_dim = hdim = 256

        self.fnet = BasicEncoder(output_dim=hdim, norm_fn="instance")

        self.encoder_block = [("encoder_block" + str(i)) for i in range(3)]
        for i in self.encoder_block:
            self.__setattr__(i, TransEncoder(2, hidden_dim=hdim))

        self.down_layer = [("down_layer" + str(i)) for i in range(2)]
        for i in self.down_layer:
            self.__setattr__(i, nn.Conv2D(256, 256, 3, stride=2, padding=1))

        self.decoder_block = [("decoder_block" + str(i)) for i in range(3)]
        for i in self.decoder_block:
            self.__setattr__(i, TransDecoder(2, hidden_dim=hdim))

        self.up_layer = [("up_layer" + str(i)) for i in range(2)]
        for i in self.up_layer:
            self.__setattr__(
                i, nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)
            )

        self.query_embed = nn.Embedding(81, self.hidden_dim)

        self.update_block = UpdateBlock(self.hidden_dim)

    def initialize_flow(self, img):
        N, _, H, W = img.shape
        coodslar = coords_grid(N, H, W)
        coords0 = coords_grid(N, H // 8, W // 8)
        coords1 = coords_grid(N, H // 8, W // 8)
        return coodslar, coords0, coords1

    def upsample_flow(self, flow, mask):
        N, _, H, W = flow.shape

        mask = mask.reshape([N, 1, 9, 8, 8, H, W])
        mask = F.softmax(mask, axis=2)

        up_flow = F.unfold(8 * flow, [3, 3], paddings=1)
        up_flow = up_flow.reshape([N, 2, 9, 1, 1, H, W])

        up_flow = paddle.sum(mask * up_flow, axis=2)
        up_flow = up_flow.transpose(perm=[0, 1, 4, 2, 5, 3])

        return up_flow.reshape([N, 2, 8 * H, 8 * W])

    def forward(self, image):
        fmap = self.fnet(image)
        fmap = F.relu(fmap)

        fmap1 = self.__getattr__(self.encoder_block[0])(fmap)
        fmap1d = self.__getattr__(self.down_layer[0])(fmap1)
        fmap2 = self.__getattr__(self.encoder_block[1])(fmap1d)
        fmap2d = self.__getattr__(self.down_layer[1])(fmap2)
        fmap3 = self.__getattr__(self.encoder_block[2])(fmap2d)

        query_embed0 = self.query_embed.weight.unsqueeze(1).tile([1, fmap3.shape[0], 1])

        fmap3d_ = self.__getattr__(self.decoder_block[0])(fmap3, query_embed0)
        fmap3du_ = (
            self.__getattr__(self.up_layer[0])(fmap3d_)
            .flatten(2)
            .transpose(perm=[2, 0, 1])
        )
        fmap2d_ = self.__getattr__(self.decoder_block[1])(fmap2, fmap3du_)
        fmap2du_ = (
            self.__getattr__(self.up_layer[1])(fmap2d_)
            .flatten(2)
            .transpose(perm=[2, 0, 1])
        )
        fmap_out = self.__getattr__(self.decoder_block[2])(fmap1, fmap2du_)

        coodslar, coords0, coords1 = self.initialize_flow(image)
        coords1 = coords1.detach()
        mask, coords1 = self.update_block(fmap_out, coords1)
        flow_up = self.upsample_flow(coords1 - coords0, mask)
        bm_up = coodslar + flow_up

        return bm_up