import torch
import torch.nn as nn
from einops.einops import rearrange
import torch.nn.functional as F


def generate_random_masks(batch, patch_size, mask_ratio, generator=None, margins=[0, 0, 0, 0]):
    mae_mask0 = _gen_random_mask(batch["image0"], patch_size, mask_ratio, generator, margins=margins)
    mae_mask1 = _gen_random_mask(batch["image1"], patch_size, mask_ratio, generator, margins=margins)
    batch.update({"mae_mask0": mae_mask0, "mae_mask1": mae_mask1})


def _gen_random_mask(image, patch_size, mask_ratio, generator=None, margins=[0, 0, 0, 0]):
    """Random mask generator
    Args:
        image (torch.Tensor): [N, C, H, W]
        patch_size (int)
        mask_ratio (float)
        generator (torch.Generator): RNG to create the same random masks for validation
        margins [float, float, float, float]: unused part for masking (up bottom left right)
    Returns:
        mask (torch.Tensor): (N, L)
    """
    N = image.shape[0]
    l = image.shape[2] // patch_size
    L = l**2
    len_keep = int(L * (1 - mask_ratio * (1 - sum(margins))))

    margins = [int(margin * l) for margin in margins]

    noise = torch.rand(N, l, l, device=image.device, generator=generator)
    if margins[0] > 0:
        noise[:, : margins[0], :] = 0
    if margins[1] > 0:
        noise[:, -margins[1] :, :] = 0
    if margins[2] > 0:
        noise[:, :, : margins[2]] = 0
    if margins[3] > 0:
        noise[:, :, -margins[3] :] = 0
    noise = noise.flatten(1)

    # sort noise for each sample
    ids_shuffle = torch.argsort(noise, dim=1)
    ids_restore = torch.argsort(ids_shuffle, dim=1)

    # generate the binary mask: 0 is keep 1 is remove
    mask = torch.ones([N, L], device=image.device)
    mask[:, :len_keep] = 0
    # unshuffle to get the binary mask
    mask = torch.gather(mask, dim=1, index=ids_restore)
    return mask


def patchify(data):
    """Split images into small overlapped patches
    Args:
        data (dict):{
            'image0_norm' (torch.Tensor): [N, C, H, W] normalized image,
            'image1_norm' (torch.Tensor): [N, C, H, W] normalized image,
    Returns:
        image0 (torch.Tensor): [N, K, W_f**2, -1] (K: num of windows)
        image1 (torch.Tensor): [N, K, W_f**2, -1] (K: num of windows)
    """
    stride = data["hw0_i"][0] // data["hw0_c"][0]
    scale = data["hw0_i"][0] // data["hw0_f"][0]
    W_f = data["W_f"]
    kernel_size = [int(W_f * scale), int(W_f * scale)]
    padding = kernel_size[0] // 2 - 1 if kernel_size[0] % 2 == 0 else kernel_size[0] // 2

    image0 = data["image0_norm"] if "image0_norm" in data else data["image0"]
    image1 = data["image1_norm"] if "image1_norm" in data else data["image1"]

    image0 = F.unfold(image0, kernel_size=kernel_size, stride=stride, padding=padding)
    image0 = rearrange(image0, "n (c h p w q) l -> n l h w p q c", h=W_f, w=W_f, p=scale, q=scale)
    image0 = image0.flatten(4)
    image0 = image0.reshape(*image0.shape[:2], W_f**2, -1)

    image1 = F.unfold(image1, kernel_size=kernel_size, stride=stride, padding=padding)
    image1 = rearrange(image1, "n (c h p w q) l -> n l h w p q c", h=W_f, w=W_f, p=scale, q=scale)
    image1 = image1.flatten(4)
    image1 = image1.reshape(*image1.shape[:2], W_f**2, -1)

    return image0, image1


def get_target(data):
    """Create target patches for mae"""
    target0, target1 = patchify(data)
    data.update({"target0": target0, "target1": target1})
