"""
对Kaiming He新提出的MAE的简要实现
"""
# ----------------------------------------------------------------
# borrowed from https://github.com/IcarusWizard/MAE.git
# ----------------------------------------------------------------

import numpy as np, torch
import torch.nn as nn

from .transformer import Encoder

def random_index(size: int):
    """
    产生随机的patch序列 用于产生mask机制
    """
    assert size > 0
    # 用于mask机制的随机索引
    forward_indexes = np.arange(size)
    np.random.shuffle(forward_indexes)
    # 用于重建图像的正确的排序索引
    backward_indexes = np.argsort(forward_indexes)
    return forward_indexes, backward_indexes

class Projection(nn.Module):
    """
    图像和嵌入之间的投射
    """
    def __init__(
        self,
        image_size,
        patch_num,
        embedding_dim,
    ):
        """
        Args:
            image_size: 图像大小
            patch_num: 每条边对应的Patch数目
            embedding_dim: 嵌入的维度
        """
        super(Projection, self).__init__()

        assert image_size % patch_num == 0
        self.image_size = image_size
        self.patch_num = patch_num
        self.embedding_dim = embedding_dim

        self.patch_size = self.image_size // self.patch_num

        self.in_projection = nn.Linear(self.patch_size ** 2 * 3, self.embedding_dim)
        self.out_projection = nn.Linear(self.embedding_dim, self.patch_size ** 2 * 3)

    def img2embeddings(
        self,
        imgs: torch.Tensor
    ):
        """
        将图像投影为嵌入向量: batch * 3 * h * w --> batch * r^2 * 3 * h / r * w / 2 --> batch * r^2 * (3 * h / r * w / r)
        """
        b,c,h,w = imgs.size()
        # ========================================
        # TODO find a faster way for patching
        # ========================================
        patches = torch.zeros(size = (b, self.patch_num ** 2, c, h // self.patch_num, w // self.patch_num)).to(imgs.device)
        for i in range(self.patch_num):
            for j in range(self.patch_num):
                patches[:,i * self.patch_num + j] = imgs[:, :, i * self.patch_size:(i + 1) * self.patch_size, j * self.patch_size:(j + 1) * self.patch_size]
        patches = patches.contiguous()

        b, total_num, c, h, w = patches.size()
        origin_embeddings = patches.view(b, total_num, -1)

        projected_embeddings = self.in_projection(origin_embeddings)

        return origin_embeddings, projected_embeddings
    
    def embeddings2img(
        self,
        embeddings: torch.Tensor,
    ):
        """
        将嵌入向量重新变换为图像: batch * 3 * h * w --> batch * r^2 * 3 * h / r * w / 2 --> batch * r^2 * (3 * h / r * w / r)
        """
        out_embeddings = self.out_projection(embeddings)
        b, total_num, _ = out_embeddings.size()
        embeddings = out_embeddings.view(b, total_num, 3, self.patch_size, self.patch_size)

        # ========================================
        # TODO find a faster way for patching
        # ========================================
        out_img = torch.zeros(size = (b, 3, self.image_size, self.image_size)).to(embeddings.device)
        for i in range(self.patch_num):
            for j in range(self.patch_num):
                index = i * self.patch_num + j
                out_img[:, :, i * self.patch_size:(i + 1) * self.patch_size, j * self.patch_size:(j + 1) * self.patch_size] = embeddings[:, index]
        return out_embeddings, out_img

class PatchShuffle(nn.Module):
    """
    对各个Patch进行挑选并且返回索引
    """
    def __init__(
        self,
        patch_num,
    ):
        """
        Args:
            patch_num: 每条边的Patch数目
        """
        super(PatchShuffle, self).__init__()
        self.patch_num = patch_num

    def forward(self, b):
        indexes = [random_index(self.patch_num ** 2) for i in range(b)]
        forward_indexes = [indexes[i][0] for i in range(b)]  # b * total_num
        backward_indexes = [indexes[i][1] for i in range(b)] # b * total_num
        forward_indexes = torch.tensor(forward_indexes, dtype=torch.int64, requires_grad=False)
        backward_indexes = torch.tensor(backward_indexes, dtype=torch.int64, requires_grad=False)

        return forward_indexes, backward_indexes

class MAE_Encoder(nn.Module):
    """
    MAE中非对称的编码器
    """
    def __init__(
        self,
        layers,
        embedding_dim,
        patch_num,
    ):
        """
        Args:
            layers: Encoder的编码器层数
            embedding_dim: 输入的嵌入维度
            patch_num: 每条边对应的Patch
        """
        super(MAE_Encoder, self).__init__()
        self.layers = layers
        self.patch_num = patch_num
        self.embedding_dim = embedding_dim
        # ========================================================
        # TODO 不同之处在于没有使用cls_token机制 仅仅当作一个特征抽取
        # ========================================================
        self.transformer = Encoder(layers, embedding_dim, embedding_dim, 8, embedding_dim * 4, 0.0)
        self.pos_embeddings = nn.parameter.Parameter(torch.zeros(size = (1, patch_num ** 2, embedding_dim), dtype=torch.float32, requires_grad=True))

    def forward(self, inputs, forward_indexes, mask_ratio):
        """
        Args:
            inputs: 输入嵌入 batch * total_num * embedding_dim
            forward_indexes: shuffle之后的indexes
            mask_ratio: 遮挡的比例
        """
        b = inputs.size()[0]
        inputs = inputs + self.pos_embeddings
        indexes_in_kept = forward_indexes[:, :int((self.patch_num ** 2) * (1 - mask_ratio))] if forward_indexes is not None else None
        if indexes_in_kept is not None:
            unmasked_embeddings = torch.stack(
                [inputs[i][indexes_in_kept[i]] for i in range(b)], dim=0
            )
            inputs = unmasked_embeddings
        return self.transformer(inputs)

class MAE_Decoder(nn.Module):
    """
    MAE中非对称的解码器
    """
    def __init__(
        self,
        layers,
        embedding_dim,
        patch_num,
    ):
        """
        Args:
            layers: 层数
            embedding_dim: 中间向量的维度
            patch_num: 每条边的Patch数目
        """
        super(MAE_Decoder, self).__init__()

        self.layers = layers
        self.embedding_dim = embedding_dim
        self.patch_num = patch_num

        self.pos_embeddings = nn.parameter.Parameter(torch.zeros(size = (self.patch_num ** 2, self.embedding_dim), requires_grad=True))
        self.mask_token = nn.parameter.Parameter(torch.zeros(size = (1, 1, self.embedding_dim), requires_grad=True))
        self.transformer = Encoder(self.layers, self.embedding_dim, self.embedding_dim, 8, self.embedding_dim * 4, dropout=0.0)

    def forward(
        self,
        x,
        forward_indexes
    ):
        b, kept_num, dim = x.size()
        # 补全所有的嵌入向量
        full_embeddings = torch.cat([x, self.mask_token.expand(b, self.patch_num ** 2 - kept_num, dim)], dim = 1) if forward_indexes is not None else x
        
        # 调整正确的位置编码
        full_embeddings = torch.stack(
            [full_embeddings[i] + self.pos_embeddings[forward_indexes[i]] for i in range(b)], dim=0
        ) if forward_indexes is not None else (full_embeddings + self.pos_embeddings.unsqueeze(0))

        output_embeddings = self.transformer(full_embeddings)

        # 得到输出
        return output_embeddings

class MAE(nn.Module):
    """
    Mask-Auto-Encoder
    将输入的图像按照尺寸定义划分为多个Patch
    先线性投射为嵌入
    """
    def __init__(
        self,
        image_size,
        patch_num,
        embedding_dim,
        encoder_layers,
        decoder_layers,
        mask_ratio,
    ):
        """
        Args:
            image_size: 输入的图像大小
            patch_num: 每条边可以划分的Patch数目
            embedding_dim: 编码器内部投影的向量维度
            encoder_layers: 编码器结构对应的层数
            decoder_layers: 解码器结果对应的层数
            mask_ratio: 遮挡比例
        """
        super(MAE, self).__init__()

        self.image_size = image_size
        self.patch_num = patch_num
        self.embedding_dim = embedding_dim
        self.encoder_layers = encoder_layers
        self.decoder_layers = decoder_layers
        self.mask_ratio = mask_ratio

        self.projection = Projection(image_size, patch_num, embedding_dim)
        self.pshuffle = PatchShuffle(self.patch_num)
        self.encoder = MAE_Encoder(self.encoder_layers, self.embedding_dim, self.patch_num)
        self.decoder = MAE_Decoder(self.decoder_layers, self.embedding_dim, self.patch_num)
    
    def forward(
        self,
        imgs,
        encoder_mode = False
    ):
        """
        Args:
            imgs: 多个图像Patch构成的张量
            encoder_mode: 是否处于特征提取模式 即仅仅使用编码器
        """
        b, _, _, _ = imgs.size()

        # b * total_num * (c * h * w), b * total_num * embeddings
        origin_embeddings, input_embeddings = self.projection.img2embeddings(imgs)  
        forward_indexes, backward_indexes = None, None
        if self.training and not encoder_mode:
            # shuffle index, correct index in order for shuffled embeddings
            forward_indexes, backward_indexes = self.pshuffle(b)                       

        # encoding
        middle_feat_embeddings = self.encoder(input_embeddings, forward_indexes, self.mask_ratio)
        if encoder_mode:
            return middle_feat_embeddings

        # decoding
        dec_embeddings = self.decoder(middle_feat_embeddings, forward_indexes)
        dec_embeddings = torch.stack(
            [dec_embeddings[i][backward_indexes[i]] for i in range(b)], dim = 0
        ) if backward_indexes is not None else dec_embeddings
        
        output_embeddings, output_img =  self.projection.embeddings2img(dec_embeddings)

        return origin_embeddings, output_embeddings, output_img, forward_indexes

