# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# References:
# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
# DeiT: https://github.com/facebookresearch/deit
# --------------------------------------------------------

from functools import partial

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

from timm.models.vision_transformer import PatchEmbed, Block

from util.pos_embed import get_2d_sincos_pos_embed
from MLP import MLP
from einops import repeat


def generate_mask(x_all, s_id=0):
    mask_s = torch.zeros_like(x_all)
    mask_t = torch.ones_like(x_all)
    mask_s[:, s_id] = 1.
    mask_t[:, s_id] = 0.
    return mask_s, mask_t


def generate_maskmat(N, d, method='UniformNid'):
    maskmat_s = torch.zeros(N, d)
    maskmat_t = torch.zeros(N, d)

    if method == 'UniformNid':
        for i in range(N):
            maskmat_s[i], maskmat_t[i] = generate_1mask_UniformNid(d)
    elif method == 'UniformNset':
        pass

    return maskmat_s, maskmat_t


def generate_1mask_UniformNid(d):
    mask_s = torch.zeros(d)
    mask_t = torch.zeros(d)

    permIdSet_S = torch.randperm(d)
    Nid_S = torch.randint(d, [1, ])
    S = permIdSet_S[:Nid_S]
    mask_s[S] = 1

    permIdSet_T = permIdSet_S[Nid_S:]
    Nid_T = torch.randint(permIdSet_T.numel(), [1, ]) + 1
    T = permIdSet_T[:Nid_T]
    mask_t[T] = 1

    return mask_s, mask_t


class MaskedAutoencoderViT(nn.Module):
    """ Masked Autoencoder with VisionTransformer backbone
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3,
                 embed_dim=1024, depth=24, num_heads=16,
                 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
                 mlp_ratio=4., norm_layer=nn.LayerNorm, norm_pix_loss=False, num_class=10,
                 GLastNorm='LN', toImg='vit', num_channel_bfconv=32):
        super().__init__()

        self.num_class = num_class
        self.in_chans = in_chans
        self.L = (img_size // patch_size) ** 2 + 1
        self.toImg = toImg
        self.num_channel_bfconv = num_channel_bfconv
        self.patch_size = patch_size
        self.img_size = img_size

        # --------------------------------------------------------------------------
        # MAE encoder specifics
        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
        num_patches = self.patch_embed.num_patches

        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim),
                                      requires_grad=False)  # fixed sin-cos embedding

        self.blocks = nn.ModuleList([
            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)
            for i in range(depth)])
        self.norm = norm_layer(embed_dim)
        # --------------------------------------------------------------------------

        # --------------------------------------------------------------------------
        # MAE decoder specifics
        self.decoder_embed_dim = decoder_embed_dim
        self.decoder_embed = nn.Linear(embed_dim, embed_dim, bias=True)

        self.cls_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
        self.nos_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.mask_token_s = nn.Parameter(torch.zeros(1, 1, embed_dim))

        self.label_embed = nn.Embedding(self.num_class, decoder_embed_dim)
        self.decoder_pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim),
                                              requires_grad=False)  # fixed sin-cos embedding

        # self.nos_pos_embed = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim), requires_grad=False)
        # self.decoder_nos_embed = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim), requires_grad=False)
        self.decoder_depth = decoder_depth
        self.decoder_blocks = nn.ModuleList([
            Block(decoder_embed_dim, decoder_num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)
            for i in range(decoder_depth)])

        if GLastNorm == 'LN':
            self.decoder_norm = norm_layer(decoder_embed_dim)
        elif GLastNorm == 'NO':
            self.decoder_norm = nn.Identity(decoder_embed_dim)
        elif GLastNorm == 'IN':
            self.decoder_norm = nn.InstanceNorm1d(self.L, affine=True)
        elif GLastNorm == 'LR':
            self.decoder_norm = nn.LeakyReLU(0.2)

        # for convs
        if self.toImg == '2conv':
            self.decoder_pred_proj = nn.Linear(decoder_embed_dim, patch_size ** 2 * num_channel_bfconv, bias=True)
            self.decoder_pred_head = nn.Sequential(
                nn.BatchNorm2d(num_channel_bfconv),
                nn.LeakyReLU(0.2),
                nn.Conv2d(num_channel_bfconv, num_channel_bfconv, 3, stride=1, padding=1),
                nn.BatchNorm2d(num_channel_bfconv),
                nn.LeakyReLU(0.2),
                nn.Conv2d(num_channel_bfconv, 3, 5, stride=1, padding=2),
                nn.Tanh()
            )
        elif self.toImg == '4conv':
            self.decoder_pred_proj = nn.Linear(decoder_embed_dim, patch_size ** 2 * num_channel_bfconv, bias=True)
            self.decoder_pred_head = nn.Sequential(
                nn.BatchNorm2d(num_channel_bfconv),
                nn.LeakyReLU(0.2),
                nn.Conv2d(num_channel_bfconv, num_channel_bfconv, 3, stride=1, padding=1),
                nn.BatchNorm2d(num_channel_bfconv),
                nn.LeakyReLU(0.2),
                nn.Conv2d(num_channel_bfconv, num_channel_bfconv, 3, stride=1, padding=1),
                nn.BatchNorm2d(num_channel_bfconv),
                nn.LeakyReLU(0.2),
                nn.Conv2d(num_channel_bfconv, num_channel_bfconv, 3, stride=1, padding=1),
                nn.BatchNorm2d(num_channel_bfconv),
                nn.LeakyReLU(0.2),
                nn.Conv2d(num_channel_bfconv, 3, 5, stride=1, padding=2),
                nn.Tanh()
            )
        elif self.toImg == 'vit':
            self.decoder_pred = nn.Sequential(
                nn.Linear(decoder_embed_dim, patch_size ** 2 * in_chans, bias=True),  # decoder to patch
                nn.Tanh()
            )
        elif self.toImg == '3conv':
            self.decoder_pred_proj = nn.Linear(decoder_embed_dim, patch_size ** 2 * num_channel_bfconv,
                                               bias=True)  # 512*8192
            self.decoder_pred_head = nn.Sequential(
                nn.BatchNorm2d(num_channel_bfconv),
                nn.LeakyReLU(0.2),
                nn.Conv2d(num_channel_bfconv, num_channel_bfconv, 3, stride=1, padding=1),
                nn.BatchNorm2d(num_channel_bfconv),
                nn.LeakyReLU(0.2),
                nn.Conv2d(num_channel_bfconv, num_channel_bfconv, 3, stride=1, padding=1),
                nn.BatchNorm2d(num_channel_bfconv),
                nn.LeakyReLU(0.2),
                nn.Conv2d(num_channel_bfconv, 3, 5, stride=1, padding=2),
                nn.Tanh()
            )

        else:
            pass

        self.len_patch = self.decoder_pos_embed.shape[1] - 1

        self.mapping_params = {'layers': [512, 512, 512, 512, 512, 512], 'activation': 'gelu',
                               'dropout_rate': 0.0,
                               'in_features': embed_dim, 'out_features': 512}
        self.mapping_mlp = MLP(**self.mapping_params)
        self.to_noise = nn.Sequential(
            torch.nn.LeakyReLU(0.2),
            nn.Linear(512, embed_dim * self.len_patch)
        )
        # --------------------------------------------------------------------------
        # --------------------------------------------------------------------------
        # MAE classification specifics
        self.head = nn.Linear(decoder_embed_dim, self.num_class)
        # --------------------------------------------------------------------------
        self.norm_pix_loss = norm_pix_loss

        self.initialize_weights()

    def initialize_weights(self):
        # initialization
        # initialize (and freeze) pos_embed by sin-cos embedding
        # # type#1, [class_token, postion_embed, noise_token]
        # pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), noise_token=True, cls_token=True)
        # self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
        # # type#2, [class_token, postion_embed], noise_token
        pos_embed_all = get_2d_sincos_pos_embed(
            self.pos_embed.shape[-1], int(self.patch_embed.num_patches ** .5),
            nos_token=True, cls_token=True)  # [class_token, postion_embed, noise_token]
        pos_embed = pos_embed_all[:-1, ...]  # [class_token, postion_embed]
        # nos_pos_embed = pos_embed_all[-1:, ...]
        # self.nos_pos_embed.data.copy_(torch.from_numpy(nos_pos_embed).float().unsqueeze(0))
        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))

        # # type#1, [class_token, postion_embed, noise_token]
        # decoder_pos_embed = get_2d_sincos_pos_embed(self.decoder_pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), noise_token=True, cls_token=True)
        # self.decoder_pos_embed.data.copy_(torch.from_numpy(decoder_pos_embed).float().unsqueeze(0))
        # # type#2, [class_token, postion_embed], noise_token
        decoder_pos_embed_all = get_2d_sincos_pos_embed(
            self.decoder_pos_embed.shape[-1], int(self.patch_embed.num_patches ** .5),
            nos_token=True, cls_token=True)  # [class_token, postion_embed, noise_token]
        decoder_pos_embed = decoder_pos_embed_all[:-1, ...]  # [class_token, postion_embed]
        nos_pos_embed = decoder_pos_embed_all[-1:, ...]
        # self.nos_pos_embed.data.copy_(torch.from_numpy(nos_pos_embed).float().unsqueeze(0))
        self.decoder_pos_embed.data.copy_(torch.from_numpy(decoder_pos_embed).float().unsqueeze(0))

        # initialize patch_embed like nn.Linear (instead of nn.Conv2d)
        w = self.patch_embed.proj.weight.data
        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))

        # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)
        torch.nn.init.normal_(self.cls_token, std=.02)
        torch.nn.init.normal_(self.nos_token, std=.02)
        torch.nn.init.normal_(self.mask_token, std=.02)
        torch.nn.init.normal_(self.mask_token_s, std=.02)
        # torch.nn.init.normal_(self.label_embed.weight, std=.02) # ???

        # initialize nn.Linear and nn.LayerNorm
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            # we use xavier_uniform following official JAX ViT:
            # torch.nn.init.xavier_uniform_(m.weight, gain=0.6)
            torch.nn.init.xavier_uniform_(m.weight)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
            if m.weight is not None:
                nn.init.constant_(m.weight, 1.0)

    def patchify(self, imgs):
        """
        imgs: (N, 3, H, W)
        x: (N, L, patch_size**2 *3)
        """
        p = self.patch_embed.patch_size[0]
        assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0

        h = w = imgs.shape[2] // p
        x = imgs.reshape(shape=(imgs.shape[0], self.in_chans, h, p, w, p))
        x = torch.einsum('nchpwq->nhwpqc', x)
        x = x.reshape(shape=(imgs.shape[0], h * w, p ** 2 * self.in_chans))
        return x

    def unpatchify(self, x):
        """
        x: (N, L, patch_size**2 *3)
        imgs: (N, 3, H, W)
        """
        p = self.patch_embed.patch_size[0]
        h = w = int(x.shape[1] ** .5)
        assert h * w == x.shape[1]

        x = x.reshape(shape=(x.shape[0], h, w, p, p, self.in_chans))
        x = torch.einsum('nhwpqc->nchpwq', x)
        imgs = x.reshape(shape=(x.shape[0], self.in_chans, h * p, h * p))
        return imgs

    def unpatchify_feamap(self, x):
        """
        x: (N, L, (patch_size)**2 *C)
        imgs: (N, C, H, W)
        """
        p = self.patch_embed.patch_size[0]
        c = x.shape[-1] // (p ** 2)
        h = w = int(x.shape[1] ** .5)
        assert h * w == x.shape[1]

        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))  # self.in_chans
        x = torch.einsum('nhwpqc->nchpwq', x)
        imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p))
        return imgs

    def fourier_pos_embedding(self, hidden, fea_dim=128):
        device = hidden.device
        # Create input pixel coordinates in the unit square
        coords = np.linspace(-1, 1, self.img_size, endpoint=True)
        pos = np.stack(np.meshgrid(coords, coords), -1)
        pos = torch.tensor(pos, dtype=torch.float, device=device)
        pos = self.lff(pos).reshape([self.img_size ** 2, fea_dim])
        result = repeat(pos, 'p h -> n p h', n=hidden.shape[0])
        return result.to(device)

    def random_masking(self, x, r_ind=None, Sratio=0., Tratio=1., Lst=None):
        # mask_ratio=0.75, maskType=4, len_s=None):
        """
        Perform per-sample random masking by per-sample shuffling.
        Per-sample shuffling is done by argsort random noise.
        x: [N, L, D], sequence
        """
        N, L, D = x.shape  # batch, length, dim

        ''' ----------- change ----------- '''
        if r_ind is None:
            r_ind = [torch.randperm(L).to(x.device) for i in range(N)]
            r_ind = torch.cat(r_ind, dim=0).view(N, L)
        ids_restore = torch.argsort(r_ind, dim=1)

        mask_s = torch.zeros([N, L], device=x.device)
        mask_t = torch.zeros([N, L], device=x.device)

        if (Tratio == -10.) and (Lst is not None):
            len_s = int(Lst * Sratio)
            len_t = Lst - len_s
        else:
            len_s = int(L * Sratio)
            len_t = int(np.ceil((L - len_s) * Tratio))

        mask_s[:, :len_s] = 1
        mask_t[:, len_s:len_s + len_t] = 1

        x_masked = torch.gather(x, dim=1, index=r_ind[:, :len_s].unsqueeze(-1).repeat(1, 1, D))

        mask_s = torch.gather(mask_s, dim=1, index=ids_restore)
        mask_t = torch.gather(mask_t, dim=1, index=ids_restore)

        return x_masked, mask_s, mask_t, ids_restore, r_ind

    def add_mask_token(self, x, ids_restore):
        x = self.decoder_embed(x)  # B, L, E
        # append mask tokens to sequence
        mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] - x.shape[1], 1)
        x_ = torch.cat([x, mask_tokens], dim=1)  # no cls and noise token
        x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]))  # unshuffle
        return x_

    def f_encode(self, x):
        for blk in self.blocks:
            x = blk(x)
        x = self.norm(x)
        return x

    def forward_encoder(self, x, r_ind=None, Sratio=0., Tratio=1., Lst=None):

        x = self.patch_embed(x)  # embed patches  B, L, E
        x = x + self.pos_embed[:, 1:, :]  # add pos embed w/o cls and noise token
        x, mask_s, mask_t, ids_restore, r_ind = self.random_masking(x, r_ind=r_ind, Sratio=Sratio, Tratio=Tratio,
                                                                    Lst=Lst)

        if Sratio > 0.:
            x = self.f_encode(x)

        return x, mask_s, mask_t, ids_restore, r_ind

    def forward_loss(self, imgs, pred, mask_t):
        """
        imgs: [N, 3, H, W]
        pred: [N, L, p*p*3]
        mask: [N, L], 0 is keep, 1 is remove,
        """
        target = self.patchify(imgs)
        # pred_patch = (1 - mask.unsqueeze(-1)) * target + mask.unsqueeze(-1) * pred
        if self.norm_pix_loss:
            mean = target.mean(dim=-1, keepdim=True)
            var = target.var(dim=-1, keepdim=True)
            target = (target - mean) / (var + 1.e-6) ** .5

        loss = (pred - target) ** 2
        loss = loss.mean(dim=-1)  # [N, L], mean loss per patch

        loss = (loss * mask_t).sum() / mask_t.sum()  # mean loss on removed patches
        # pred_img = self.unpatchify(pred_patch)
        return loss

    def forward(self, imgs, z=None, label=None, r_ind=None, Sratio=0., Tratio=1., method='generate', Lst=None):
        # MUST specify "Sratio [0,1), Tratio (0,1], method" outside the function
        # NOTE: Tratio = -10. NOW indicates communication's forward
        assert Sratio >= 0.  # and Tratio >= 0.
        assert (Tratio > 0) or ((Tratio == -10.) and (Lst is not None) and (r_ind is not None))
        N, L = imgs.shape[0], self.L - 1

        if method in ['generate', 'conditional_generate']:
            assert Sratio < 1.
            w = self.mapping_mlp(z)
            cc_noise = self.to_noise(w)
            cc_noise = cc_noise.reshape(N, L, -1)

        if method in ['classify']:
            assert Sratio > 0.
            cc_noise = self.nos_token.expand(N, L, -1)

        #  Encoding
        latent, mask_s, mask_t, ids_restore, r_ind = self.forward_encoder(imgs, r_ind=r_ind, Sratio=Sratio,
                                                                          Tratio=Tratio, Lst=Lst)

        #  Decoding
        x_ = self.add_mask_token(latent, ids_restore)

        if method in ['classify']:
            cls_token = self.cls_token[:, :, :x_.shape[-1]].expand(N, -1, -1)
        else:
            cls_token = self.label_embed(label)[:, :x_.shape[-1]].unsqueeze(1)  # B,1,768

        x = torch.cat([cls_token, x_], dim=1)  # append cls token and noise token

        # add pos embed
        x = x + self.decoder_pos_embed

        # cat noise + mask_noise
        if Sratio == 0.:
            x_down = cc_noise
        else:
            cc_masks = self.mask_token_s.expand(N, L, -1)
            # x_down = cc_masks * mask_s.unsqueeze(-1) + cc_noise * mask_t.unsqueeze(-1)
            x_down = cc_masks * mask_s.unsqueeze(-1) + cc_noise * (1. - mask_s).unsqueeze(-1)  # generate S U T = L
        x_down = torch.cat((self.cls_token[:, :, x.shape[-1]:].expand(N, -1, -1), x_down), dim=1)
        x = torch.cat([x, x_down], dim=-1)

        flag = 0
        for blk in self.decoder_blocks:
            x = blk(x)  # apply noise to every layer
            if (method in ['generate', 'conditional_generate']) and (flag < self.decoder_depth - 1):
                x[:, :, x.shape[-1] // 2:] = x_down
                flag += 1

        x = self.decoder_norm(x)

        cls_token_out = x[:, :1, :]
        if 'conv' in self.toImg:
            tmp = self.unpatchify_feamap(self.decoder_pred_proj(x[:, 1:, :]))  # NLD -> NCHW
            pred_patch = self.patchify(self.decoder_pred_head(tmp))  # N3HW -> NLD
        elif self.toImg == 'vit':
            pred_patch = self.decoder_pred(x[:, 1:, :])
        else:
            pass

        # forming the predicted image
        if Sratio == 0.:
            pred_patch_st = pred_patch
        else:
            target = self.patchify(imgs)
            # pred_patch_st = mask_s.unsqueeze(-1) * target + mask_t.unsqueeze(-1) * pred_patch
            pred_patch_st = mask_s.unsqueeze(-1) * target + (1. - mask_s).unsqueeze(
                -1) * pred_patch  # generate S U T = L
        pred_img_st = self.unpatchify(pred_patch_st)

        return pred_img_st, cls_token_out, mask_s, mask_t, r_ind


def mae_vit_base_patch16_dec512d8b(**kwargs):
    model = MaskedAutoencoderViT(
        patch_size=16, embed_dim=768, depth=12, num_heads=12,
        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model


def mae_vit_base_patch16_dec512d12h16(**kwargs):
    model = MaskedAutoencoderViT(
        patch_size=16, embed_dim=768, depth=12, num_heads=12,
        decoder_embed_dim=512, decoder_depth=12, decoder_num_heads=16,
        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model


def mae_vit_large_patch16_dec512d8b(**kwargs):
    model = MaskedAutoencoderViT(
        patch_size=16, embed_dim=1024, depth=24, num_heads=16,
        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model


def mae_vit_huge_patch14_dec512d8b(**kwargs):
    model = MaskedAutoencoderViT(
        patch_size=14, embed_dim=1280, depth=32, num_heads=16,
        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model


def mae_vit_small_patch16_dec512d8b(**kwargs):
    model = MaskedAutoencoderViT(
        patch_size=16, embed_dim=768, depth=6, num_heads=6,
        decoder_embed_dim=512, decoder_depth=4, decoder_num_heads=8,
        mlp_ratio=2, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model


# set recommended archs
mae_vit_base_patch16 = mae_vit_base_patch16_dec512d8b  # decoder: 512 dim, 8 blocks, 86M
mae_vit_large_patch16 = mae_vit_large_patch16_dec512d8b  # decoder: 512 dim, 8 blocks, 300M
mae_vit_huge_patch14 = mae_vit_huge_patch14_dec512d8b  # decoder: 512 dim, 8 blocks

mae_vit_small = mae_vit_small_patch16_dec512d8b  # decoder: 512 dim, 8 blocks
mae_vit_largeDec = mae_vit_base_patch16_dec512d12h16
