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

from loguru import logger


def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution without padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, padding=0, bias=False)


def conv3x3(in_planes, out_planes, stride=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)


class FinePreprocess(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.config = config
        block_dims = config["backbone"]["block_dims"]
        self.W = self.config["fine_window_size"]
        self.fine_d_model = block_dims[0]

        self.layer3_outconv = conv1x1(block_dims[2], block_dims[2])
        self.layer2_outconv = conv1x1(block_dims[1], block_dims[2])
        self.layer2_outconv2 = nn.Sequential(
            conv3x3(block_dims[2], block_dims[2]),
            nn.BatchNorm2d(block_dims[2]),
            nn.LeakyReLU(),
            conv3x3(block_dims[2], block_dims[1]),
        )
        self.layer1_outconv = conv1x1(block_dims[0], block_dims[1])
        self.layer1_outconv2 = nn.Sequential(
            conv3x3(block_dims[1], block_dims[1]),
            nn.BatchNorm2d(block_dims[1]),
            nn.LeakyReLU(),
            conv3x3(block_dims[1], block_dims[0]),
        )

        self._reset_parameters()

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.kaiming_normal_(p, mode="fan_out", nonlinearity="relu")

    def inter_fpn(self, feat_c, x2, x1, stride):
        feat_c = self.layer3_outconv(feat_c)
        feat_c = F.interpolate(feat_c, scale_factor=2.0, mode="bilinear", align_corners=False)

        x2 = self.layer2_outconv(x2)
        x2 = self.layer2_outconv2(x2 + feat_c)
        x2 = F.interpolate(x2, scale_factor=2.0, mode="bilinear", align_corners=False)

        x1 = self.layer1_outconv(x1)
        x1 = self.layer1_outconv2(x1 + x2)
        x1 = F.interpolate(x1, scale_factor=2.0, mode="bilinear", align_corners=False)
        return x1

    def forward(self, feat_c0, feat_c1, data):
        W = self.W
        stride = data["hw0_f"][0] // data["hw0_c"][0]

        data.update({"W": W})
        if data["b_ids"].shape[0] == 0:
            feat0 = torch.empty(0, self.W**2, self.fine_d_model, device=feat_c0.device)
            feat1 = torch.empty(0, self.W**2, self.fine_d_model, device=feat_c0.device)
            return feat0, feat1

        if data["hw0_i"] == data["hw1_i"]:
            feat_c = rearrange(torch.cat([feat_c0, feat_c1], 0), "b (h w) c -> b c h w", h=data["hw0_c"][0])  # 1/8 feat
            x2 = data["feats_x2"]  # 1/4 feat
            x1 = data["feats_x1"]  # 1/2 feat
            del data["feats_x2"], data["feats_x1"]

            # 1. fine feature extraction
            x1 = self.inter_fpn(feat_c, x2, x1, stride)
            feat_f0, feat_f1 = torch.chunk(x1, 2, dim=0)

            # 2. unfold(crop) all local windows
            feat_f0 = F.unfold(feat_f0, kernel_size=(W, W), stride=stride, padding=0)
            feat_f0 = rearrange(feat_f0, "n (c ww) l -> n l ww c", ww=W**2)
            feat_f1 = F.unfold(feat_f1, kernel_size=(W + 2, W + 2), stride=stride, padding=1)
            feat_f1 = rearrange(feat_f1, "n (c ww) l -> n l ww c", ww=(W + 2) ** 2)

            # 3. select only the predicted matches
            feat_f0 = feat_f0[data["b_ids"], data["i_ids"]]  # [n, ww, cf]
            feat_f1 = feat_f1[data["b_ids"], data["j_ids"]]

            return feat_f0, feat_f1
        else:  # handle different input shapes
            feat_c0, feat_c1 = rearrange(feat_c0, "b (h w) c -> b c h w", h=data["hw0_c"][0]), rearrange(
                feat_c1, "b (h w) c -> b c h w", h=data["hw1_c"][0]
            )  # 1/8 feat
            x2_0, x2_1 = data["feats_x2_0"], data["feats_x2_1"]  # 1/4 feat
            x1_0, x1_1 = data["feats_x1_0"], data["feats_x1_1"]  # 1/2 feat
            del data["feats_x2_0"], data["feats_x1_0"], data["feats_x2_1"], data["feats_x1_1"]

            # 1. fine feature extraction
            feat_f0, feat_f1 = self.inter_fpn(feat_c0, x2_0, x1_0, stride), self.inter_fpn(feat_c1, x2_1, x1_1, stride)

            # 2. unfold(crop) all local windows
            feat_f0 = F.unfold(feat_f0, kernel_size=(W, W), stride=stride, padding=0)
            feat_f0 = rearrange(feat_f0, "n (c ww) l -> n l ww c", ww=W**2)
            feat_f1 = F.unfold(feat_f1, kernel_size=(W + 2, W + 2), stride=stride, padding=1)
            feat_f1 = rearrange(feat_f1, "n (c ww) l -> n l ww c", ww=(W + 2) ** 2)

            # 3. select only the predicted matches
            feat_f0 = feat_f0[data["b_ids"], data["i_ids"]]  # [n, ww, cf]
            feat_f1 = feat_f1[data["b_ids"], data["j_ids"]]

            return feat_f0, feat_f1
