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

from .util import DownBlock2d, ResBlock2d, SameBlock2d, UpBlock2d


class InpaintingNetwork(nn.Module):
    """
    Inpaint the missing regions and reconstruct the Driving image.
    """

    def __init__(
        self,
        num_channels,
        block_expansion,
        max_features,
        num_down_blocks,
        multi_mask=True,
        **kwargs,
    ):
        super(InpaintingNetwork, self).__init__()

        self.num_down_blocks = num_down_blocks
        self.multi_mask = multi_mask
        self.first = SameBlock2d(
            num_channels, block_expansion, kernel_size=(7, 7), padding=(3, 3)
        )

        down_blocks = []
        up_blocks = []
        resblock = []
        for i in range(num_down_blocks):
            in_features = min(max_features, block_expansion * (2**i))
            out_features = min(max_features, block_expansion * (2 ** (i + 1)))
            down_blocks.append(
                DownBlock2d(
                    in_features, out_features, kernel_size=(3, 3), padding=(1, 1)
                )
            )
            decoder_in_feature = out_features * 2
            if i == num_down_blocks - 1:
                decoder_in_feature = out_features
            up_blocks.append(
                UpBlock2d(
                    decoder_in_feature, in_features, kernel_size=(3, 3), padding=(1, 1)
                )
            )
            resblock.append(
                ResBlock2d(decoder_in_feature, kernel_size=(3, 3), padding=(1, 1))
            )
            resblock.append(
                ResBlock2d(decoder_in_feature, kernel_size=(3, 3), padding=(1, 1))
            )
        self.down_blocks = nn.ModuleList(down_blocks)
        self.up_blocks = nn.ModuleList(up_blocks[::-1])
        self.resblock = nn.ModuleList(resblock[::-1])

        self.final = nn.Conv2d(
            block_expansion, num_channels, kernel_size=(7, 7), padding=(3, 3)
        )
        self.num_channels = num_channels

    def deform_input(self, inp, deformation):
        _, h_old, w_old, _ = deformation.shape
        _, _, h, w = inp.shape
        if h_old != h or w_old != w:
            deformation = deformation.permute(0, 3, 1, 2)
            deformation = F.interpolate(
                deformation, size=(h, w), mode="bilinear", align_corners=True
            )
            deformation = deformation.permute(0, 2, 3, 1)
        return F.grid_sample(inp, deformation, align_corners=True)

    def occlude_input(self, inp, occlusion_map):
        if not self.multi_mask:
            if (
                inp.shape[2] != occlusion_map.shape[2]
                or inp.shape[3] != occlusion_map.shape[3]
            ):
                occlusion_map = F.interpolate(
                    occlusion_map,
                    size=inp.shape[2:],
                    mode="bilinear",
                    align_corners=True,
                )
        out = inp * occlusion_map
        return out

    def forward(self, source_image, dense_motion):
        out = self.first(source_image)
        encoder_map = [out]
        for i in range(len(self.down_blocks)):
            out = self.down_blocks[i](out)
            encoder_map.append(out)

        output_dict = {}
        output_dict["contribution_maps"] = dense_motion["contribution_maps"]
        output_dict["deformed_source"] = dense_motion["deformed_source"]

        occlusion_map = dense_motion["occlusion_map"]
        output_dict["occlusion_map"] = occlusion_map

        deformation = dense_motion["deformation"]
        out_ij = self.deform_input(out.detach(), deformation)
        out = self.deform_input(out, deformation)

        out_ij = self.occlude_input(out_ij, occlusion_map[0].detach())
        out = self.occlude_input(out, occlusion_map[0])

        warped_encoder_maps = []
        warped_encoder_maps.append(out_ij)

        for i in range(self.num_down_blocks):
            out = self.resblock[2 * i](out)
            out = self.resblock[2 * i + 1](out)
            out = self.up_blocks[i](out)

            encode_i = encoder_map[-(i + 2)]
            encode_ij = self.deform_input(encode_i.detach(), deformation)
            encode_i = self.deform_input(encode_i, deformation)

            occlusion_ind = 0
            if self.multi_mask:
                occlusion_ind = i + 1
            encode_ij = self.occlude_input(
                encode_ij, occlusion_map[occlusion_ind].detach()
            )
            encode_i = self.occlude_input(encode_i, occlusion_map[occlusion_ind])
            warped_encoder_maps.append(encode_ij)

            if i == self.num_down_blocks - 1:
                break

            out = torch.cat([out, encode_i], 1)

        deformed_source = self.deform_input(source_image, deformation)
        output_dict["deformed"] = deformed_source
        output_dict["warped_encoder_maps"] = warped_encoder_maps

        occlusion_last = occlusion_map[-1]
        if not self.multi_mask:
            occlusion_last = F.interpolate(
                occlusion_last, size=out.shape[2:], mode="bilinear", align_corners=True
            )

        out = out * (1 - occlusion_last) + encode_i
        out = self.final(out)
        out = torch.sigmoid(out)
        out = out * (1 - occlusion_last) + deformed_source * occlusion_last
        output_dict["prediction"] = out

        return output_dict

    def get_encode(self, driver_image, occlusion_map):
        out = self.first(driver_image)
        encoder_map = []
        encoder_map.append(self.occlude_input(out.detach(), occlusion_map[-1].detach()))
        for i in range(len(self.down_blocks)):
            out = self.down_blocks[i](out.detach())
            out_mask = self.occlude_input(out.detach(), occlusion_map[2 - i].detach())
            encoder_map.append(out_mask.detach())

        return encoder_map