import torch
import torch.nn as nn
from einops.einops import rearrange
from .backbone import ResNet_8_2
from .utils.position_encoding import PositionEncodingSine
from .xoftr_module import LocalFeatureTransformer, FineProcess


class XoFTR_Pretrain(nn.Module):
    def __init__(self, config):
        super().__init__()
        # Misc
        self.config = config
        self.patch_size = config["pretrain_patch_size"]

        # Modules
        self.backbone = ResNet_8_2(config["resnet"])
        self.pos_encoding = PositionEncodingSine(config["coarse"]["d_model"])
        self.loftr_coarse = LocalFeatureTransformer(config["coarse"])
        self.fine_process = FineProcess(config)
        self.mask_token_f = nn.Parameter(torch.zeros(1, config["resnet"]["block_dims"][0], 1, 1))
        self.mask_token_m = nn.Parameter(torch.zeros(1, config["resnet"]["block_dims"][1], 1, 1))
        self.mask_token_c = nn.Parameter(torch.zeros(1, config["resnet"]["block_dims"][2], 1, 1))
        self.out_proj = nn.Linear(config["resnet"]["block_dims"][0], 4)

        torch.nn.init.normal_(self.mask_token_f, std=0.02)
        torch.nn.init.normal_(self.mask_token_m, std=0.02)
        torch.nn.init.normal_(self.mask_token_c, std=0.02)

    def upsample_mae_mask(self, mae_mask, scale):
        assert len(mae_mask.shape) == 2
        p = int(mae_mask.shape[1] ** 0.5)
        return mae_mask.reshape(-1, p, p).repeat_interleave(scale, axis=1).repeat_interleave(scale, axis=2)

    def upsample_mask(self, mask, scale):
        return mask.repeat_interleave(scale, axis=1).repeat_interleave(scale, axis=2)

    def mask_layer(self, feat, mae_mask, mae_mask_scale, mask=None, mask_scale=None, mask_token=None):
        """Mask the feature map and replace with trainable inpu tokens if available
        Args:
            feat (torch.Tensor): [N, C, H, W]
            mae_mask (torch.Tensor): (N, L) mask for masked image modeling
            mae_mask_scale (int): the scale of layer to mae mask
            mask (torch.Generator): mask for padded input image
            mask_scale (int): the scale of layer to mask (mask is created on course scale)
            mask_token (torch.Tensor): [1, C, 1, 1] learnable mae mask token
        Returns:
            feat (torch.Tensor): [N, C, H, W]
        """
        mae_mask = self.upsample_mae_mask(mae_mask, mae_mask_scale)
        mae_mask = mae_mask.unsqueeze(1).type_as(feat)
        if mask is not None:
            mask = self.upsample_mask(mask, mask_scale)
            mask = mask.unsqueeze(1).type_as(feat)
            mae_mask = mask * mae_mask
        feat = feat * (1.0 - mae_mask)
        if mask_token is not None:
            mask_token = mask_token.repeat(feat.shape[0], 1, feat.shape[2], feat.shape[3])
            feat += mask_token * mae_mask
        return feat

    def forward(self, data):
        """
        Update:
            data (dict): {
                'image0': (torch.Tensor): (N, 1, H, W)
                'image1': (torch.Tensor): (N, 1, H, W)
                'mask0'(optional) : (torch.Tensor): (N, H, W) '0' indicates a padded position
                'mask1'(optional) : (torch.Tensor): (N, H, W)
            }
        """
        # 1. Local Feature CNN
        data.update(
            {"bs": data["image0"].size(0), "hw0_i": data["image0"].shape[2:], "hw1_i": data["image1"].shape[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"]

        mask0 = mask1 = None  # mask fro madded images
        if "mask0" in data:
            mask0, mask1 = data["mask0"], data["mask1"]

        # mask input images
        image0 = self.mask_layer(image0, data["mae_mask0"], mae_mask_scale=self.patch_size, mask=mask0, mask_scale=8)
        image1 = self.mask_layer(image1, data["mae_mask1"], mae_mask_scale=self.patch_size, mask=mask1, mask_scale=8)
        data.update({"masked_image0": image0.clone().detach().cpu(), "masked_image1": image1.clone().detach().cpu()})

        if data["hw0_i"] == data["hw1_i"]:  # faster & better BN convergence
            feats_c, feats_m, feats_f = self.backbone(torch.cat([image0, image1], dim=0))
            (feat_c0, feat_c1) = feats_c.split(data["bs"])
            (feat_m0, feat_m1) = feats_m.split(data["bs"])
            (feat_f0, feat_f1) = feats_f.split(data["bs"])
        else:  # handle different input shapes
            feat_c0, feat_m0, feat_f0 = self.backbone(image0)
            feat_c1, feat_m1, feat_f1 = self.backbone(image1)

        # mask output layers of backbone and replace with trainable token
        feat_c0 = self.mask_layer(
            feat_c0,
            data["mae_mask0"],
            mae_mask_scale=self.patch_size // 8,
            mask=mask0,
            mask_scale=1,
            mask_token=self.mask_token_c,
        )
        feat_c1 = self.mask_layer(
            feat_c1,
            data["mae_mask1"],
            mae_mask_scale=self.patch_size // 8,
            mask=mask1,
            mask_scale=1,
            mask_token=self.mask_token_c,
        )
        feat_m0 = self.mask_layer(
            feat_m0,
            data["mae_mask0"],
            mae_mask_scale=self.patch_size // 4,
            mask=mask0,
            mask_scale=2,
            mask_token=self.mask_token_m,
        )
        feat_m1 = self.mask_layer(
            feat_m1,
            data["mae_mask1"],
            mae_mask_scale=self.patch_size // 4,
            mask=mask1,
            mask_scale=2,
            mask_token=self.mask_token_m,
        )
        feat_f0 = self.mask_layer(
            feat_f0,
            data["mae_mask0"],
            mae_mask_scale=self.patch_size // 2,
            mask=mask0,
            mask_scale=4,
            mask_token=self.mask_token_f,
        )
        feat_f1 = self.mask_layer(
            feat_f1,
            data["mae_mask1"],
            mae_mask_scale=self.patch_size // 2,
            mask=mask1,
            mask_scale=4,
            mask_token=self.mask_token_f,
        )
        data.update(
            {
                "hw0_c": feat_c0.shape[2:],
                "hw1_c": feat_c1.shape[2:],
                "hw0_m": feat_m0.shape[2:],
                "hw1_m": feat_m1.shape[2:],
                "hw0_f": feat_f0.shape[2:],
                "hw1_f": feat_f1.shape[2:],
            }
        )

        # save coarse features for fine matching module
        feat_c0_pre, feat_c1_pre = feat_c0.clone(), feat_c1.clone()

        # 2. Coarse-level loftr module
        # add featmap with positional encoding, then flatten it to sequence [N, HW, C]
        feat_c0 = rearrange(self.pos_encoding(feat_c0), "n c h w -> n (h w) c")
        feat_c1 = rearrange(self.pos_encoding(feat_c1), "n c h w -> n (h w) c")

        mask_c0 = mask_c1 = None  # mask is useful in training
        if "mask0" in data:
            mask_c0, mask_c1 = data["mask0"].flatten(-2), data["mask1"].flatten(-2)
        feat_c0, feat_c1 = self.loftr_coarse(feat_c0, feat_c1, mask_c0, mask_c1)

        # 3. Fine-level maching module as decoder
        # generate window locations from mae mask to reconstruct
        mae_mask_c0 = self.upsample_mae_mask(data["mae_mask0"], self.patch_size // 8)
        if mask0 is not None:
            mae_mask_c0 = mae_mask_c0 * mask0.type_as(mae_mask_c0)

        mae_mask_c1 = self.upsample_mae_mask(data["mae_mask1"], self.patch_size // 8)
        if mask1 is not None:
            mae_mask_c1 = mae_mask_c1 * mask1.type_as(mae_mask_c1)

        mae_mask_c = torch.logical_or(mae_mask_c0, mae_mask_c1)

        b_ids, i_ids = mae_mask_c.flatten(1).nonzero(as_tuple=True)
        j_ids = i_ids

        # b_ids, i_ids and j_ids are masked location for both images
        # ids_image0 and ids_image1 determines which indeces belogs to which image
        ids_image0 = mae_mask_c0.flatten(1)[b_ids, i_ids]
        ids_image1 = mae_mask_c1.flatten(1)[b_ids, j_ids]

        data.update(
            {
                "b_ids": b_ids,
                "i_ids": i_ids,
                "j_ids": j_ids,
                "ids_image0": ids_image0 == 1,
                "ids_image1": ids_image1 == 1,
            }
        )

        # fine level matching module
        feat_f0_unfold, feat_f1_unfold = self.fine_process(
            feat_f0, feat_f1, feat_m0, feat_m1, feat_c0, feat_c1, feat_c0_pre, feat_c1_pre, data
        )

        # output projection 5x5 window to 10x10 window
        pred0 = self.out_proj(feat_f0_unfold)
        pred1 = self.out_proj(feat_f1_unfold)

        data.update({"pred0": pred0, "pred1": pred1})

    def load_state_dict(self, state_dict, *args, **kwargs):
        for k in list(state_dict.keys()):
            if k.startswith("matcher."):
                state_dict[k.replace("matcher.", "", 1)] = state_dict.pop(k)
        return super().load_state_dict(state_dict, *args, **kwargs)
