import math

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

from .util import (
    TPS,
    AntiAliasInterpolation2d,
    Hourglass,
    UpBlock2d,
    from_homogeneous,
    kp2gaussian,
    make_coordinate_grid,
    to_homogeneous,
)


class DenseMotionNetwork(nn.Module):
    """
    Module that predicting a dense motion from sparse motion representation given by kp_source and kp_driving
    """

    def __init__(
        self,
        block_expansion,
        num_blocks,
        max_features,
        num_kp,
        num_channels,
        estimate_occlusion_map=True,
        scale_factor=1,
        kp_variance=0.01,
    ):
        super(DenseMotionNetwork, self).__init__()

        infeatures = num_kp + 1
        self.infeatures = infeatures

        self.hourglass = Hourglass(
            block_expansion=block_expansion,
            in_features=infeatures * (num_channels + 1),
            max_features=max_features,
            num_blocks=num_blocks,
        )
        self.mask = nn.Conv2d(
            self.hourglass.out_filters, infeatures, kernel_size=(7, 7), padding=(3, 3)
        )

        if estimate_occlusion_map:
            self.occlusion = nn.Conv2d(
                self.hourglass.out_filters, 1, kernel_size=(7, 7), padding=(3, 3)
            )
        else:
            self.occlusion = None

        self.num_kp = num_kp
        self.scale_factor = scale_factor
        self.kp_variance = kp_variance

        if self.scale_factor != 1:
            self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor)

    def create_heatmap_representations(self, source_image, kp_driving, kp_source):
        spatial_size = source_image.shape[2:]
        gaussian_driving = kp2gaussian(
            kp_driving["kp"], spatial_size=spatial_size, kp_variance=self.kp_variance
        )
        gaussian_source = kp2gaussian(
            kp_source["kp"], spatial_size=spatial_size, kp_variance=self.kp_variance
        )
        heatmap = gaussian_driving - gaussian_source

        # adding background feature
        zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1]).type(
            heatmap.type()
        )
        heatmap = torch.cat([zeros, heatmap], dim=1)
        heatmap = heatmap.unsqueeze(2)
        return heatmap

    def create_sparse_motions(self, source_image, kp_driving, kp_source, bg_param=None):
        bs, _, h, w = source_image.shape
        identity_grid = make_coordinate_grid((h, w), type=kp_source["kp"].type())
        identity_grid = identity_grid.view(1, 1, h, w, 2)
        coordinate_grid = identity_grid - kp_driving["kp"].view(
            bs, self.num_kp, 1, 1, 2
        )
        if "jacobian" in kp_driving:
            jacobian = torch.matmul(
                kp_source["jacobian"], torch.inverse(kp_driving["jacobian"])
            )
            jacobian = jacobian.unsqueeze(-3).unsqueeze(-3)
            jacobian = jacobian.repeat(1, 1, h, w, 1, 1)
            coordinate_grid = torch.matmul(jacobian, coordinate_grid.unsqueeze(-1))
            coordinate_grid = coordinate_grid.squeeze(-1)
            # n * 10 * h * w * 2

        driving_to_source = coordinate_grid + kp_source["kp"].view(
            bs, self.num_kp, 1, 1, 2
        )

        # #adding background feature
        # identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1)

        # adding background feature
        if bg_param is None:
            bg_grid = identity_grid.repeat(bs, 1, 1, 1, 1)
        else:
            bg_grid = identity_grid.repeat(bs, 1, 1, 1, 1)
            bg_grid = to_homogeneous(bg_grid)
            bg_grid = torch.matmul(
                bg_param.view(bs, 1, 1, 1, 3, 3), bg_grid.unsqueeze(-1)
            ).squeeze(-1)
            bg_grid = from_homogeneous(bg_grid)

        sparse_motions = torch.cat([bg_grid, driving_to_source], dim=1)
        return sparse_motions

    def create_deformed_source_image(self, source_image, sparse_motions):
        bs, _, h, w = source_image.shape
        source_repeat = (
            source_image.unsqueeze(1)
            .unsqueeze(1)
            .repeat(1, self.num_kp + 1, 1, 1, 1, 1)
        )
        source_repeat = source_repeat.view(bs * (self.num_kp + 1), -1, h, w)
        sparse_motions = sparse_motions.view((bs * (self.num_kp + 1), h, w, -1))
        sparse_deformed = F.grid_sample(source_repeat, sparse_motions)
        sparse_deformed = sparse_deformed.view((bs, self.num_kp + 1, -1, h, w))
        return sparse_deformed

    def dropout_softmax(self, X, P):
        """
        Dropout for TPS transformations. Eq(7) and Eq(8) in the paper.
        """
        drop = (
            (torch.rand(X.shape[0], X.shape[1]) < (1 - P)).type(X.type()).to(X.device)
        )
        drop[..., 0] = 1
        drop = drop.repeat(X.shape[2], X.shape[3], 1, 1).permute(2, 3, 0, 1)

        maxx = X.max(1).values.unsqueeze_(1)
        X = X - maxx
        X_exp = X.exp()
        X[:, 1:, ...] /= 1 - P
        mask_bool = drop == 0
        X_exp = X_exp.masked_fill(mask_bool, 0)
        partition = X_exp.sum(dim=1, keepdim=True) + 1e-6
        return X_exp / partition

    def forward(
        self,
        source_image,
        kp_driving,
        kp_source,
        bg_param=None,
        dropout_flag=False,
        dropout_p=0,
    ):
        if self.scale_factor != 1:
            source_image = self.down(source_image)

        bs, _, h, w = source_image.shape

        out_dict = dict()
        heatmap_representation = self.create_heatmap_representations(
            source_image, kp_driving, kp_source
        )
        sparse_motion = self.create_sparse_motions(
            source_image, kp_driving, kp_source, bg_param=bg_param
        )
        # n * 11 * h * w * 2
        deformed_source = self.create_deformed_source_image(source_image, sparse_motion)
        # n * 11 * 3 * h * w
        out_dict["sparse_deformed"] = deformed_source

        input = torch.cat([heatmap_representation, deformed_source], dim=2)
        # n * 11(101,22) * 4 * h * w
        input = input.view(bs, -1, h, w)

        prediction = self.hourglass(input)

        mask = self.mask(prediction)
        out_dict["logit_mask"] = mask
        if dropout_flag:
            mask = self.dropout_softmax(mask, dropout_p)
        else:
            mask = F.softmax(mask, dim=1)
        # mask = F.softmax(mask, dim=1)
        out_dict["mask"] = mask
        mask = mask.unsqueeze(2)
        sparse_motion = sparse_motion.permute(0, 1, 4, 2, 3)
        # n * 11 * 2 * h * w
        deformation = (sparse_motion * mask).sum(dim=1)
        deformation = deformation.permute(0, 2, 3, 1)
        # n * h * w * 2

        out_dict["deformation"] = deformation

        if self.occlusion is not None:
            # occlusion_map = torch.sigmoid(self.occlusion(prediction))
            occlusion_map = self.occlusion(prediction)
            out_dict["occlusion"] = occlusion_map

        return out_dict


class TPSDenseMotionNetwork(nn.Module):
    """
    Module that estimating an optical flow and multi-resolution occlusion masks
                        from K TPS transformations and an affine transformation.
    """

    def __init__(
        self,
        block_expansion,
        num_blocks,
        max_features,
        num_tps,
        num_channels,
        scale_factor=0.25,
        bg=False,
        multi_mask=False,
        kp_variance=0.01,
    ):
        super(TPSDenseMotionNetwork, self).__init__()

        if scale_factor != 1:
            self.down = AntiAliasInterpolation2d(num_channels, scale_factor)
        self.scale_factor = scale_factor
        self.multi_mask = multi_mask

        self.hourglass = Hourglass(
            block_expansion=block_expansion,
            in_features=(num_channels * (num_tps + 1) + num_tps * 5 + 1),
            max_features=max_features,
            num_blocks=num_blocks,
        )

        hourglass_output_size = self.hourglass.out_filters
        # self.maps = nn.Conv2d(hourglass_output_size[-1], num_tps + 1, kernel_size=(7, 7), padding=(3, 3))
        self.maps = nn.Conv2d(
            hourglass_output_size, num_tps + 1, kernel_size=(7, 7), padding=(3, 3)
        )

        if multi_mask:
            up = []
            self.up_nums = int(math.log(1 / scale_factor, 2))
            self.occlusion_num = 4

            channel = [hourglass_output_size[-1] // (2**i) for i in range(self.up_nums)]
            for i in range(self.up_nums):
                up.append(
                    UpBlock2d(channel[i], channel[i] // 2, kernel_size=3, padding=1)
                )
            self.up = nn.ModuleList(up)

            channel = [
                hourglass_output_size[-i - 1]
                for i in range(self.occlusion_num - self.up_nums)[::-1]
            ]
            for i in range(self.up_nums):
                channel.append(hourglass_output_size[-1] // (2 ** (i + 1)))
            occlusion = []

            for i in range(self.occlusion_num):
                occlusion.append(
                    nn.Conv2d(channel[i], 1, kernel_size=(7, 7), padding=(3, 3))
                )
            self.occlusion = nn.ModuleList(occlusion)
        else:
            # occlusion = [nn.Conv2d(hourglass_output_size[-1], 1, kernel_size=(7, 7), padding=(3, 3))]
            occlusion = [
                nn.Conv2d(hourglass_output_size, 1, kernel_size=(7, 7), padding=(3, 3))
            ]
            self.occlusion = nn.ModuleList(occlusion)

        self.num_tps = num_tps
        self.bg = bg
        self.kp_variance = kp_variance

    def create_heatmap_representations(self, source_image, kp_driving, kp_source):
        spatial_size = source_image.shape[2:]
        gaussian_driving = kp2gaussian(
            kp_driving["kp"], spatial_size=spatial_size, kp_variance=self.kp_variance
        )
        gaussian_source = kp2gaussian(
            kp_source["kp"], spatial_size=spatial_size, kp_variance=self.kp_variance
        )
        heatmap = gaussian_driving - gaussian_source

        zeros = (
            torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1])
            .type(heatmap.type())
            .to(heatmap.device)
        )
        heatmap = torch.cat([zeros, heatmap], dim=1)

        return heatmap

    def create_transformations(self, source_image, kp_driving, kp_source, bg_param):
        # K TPS transformaions
        bs, _, h, w = source_image.shape
        kp_1 = kp_driving["kp"]
        kp_2 = kp_source["kp"]
        kp_1 = kp_1.view(bs, -1, 5, 2)
        kp_2 = kp_2.view(bs, -1, 5, 2)
        trans = TPS(mode="kp", bs=bs, kp_1=kp_1, kp_2=kp_2)
        driving_to_source = trans.transform_frame(source_image)

        identity_grid = make_coordinate_grid((h, w), type=kp_1.type()).to(kp_1.device)
        identity_grid = identity_grid.view(1, 1, h, w, 2)
        identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1)

        # affine background transformation
        if bg_param is not None:
            identity_grid = to_homogeneous(identity_grid)
            identity_grid = torch.matmul(
                bg_param.view(bs, 1, 1, 1, 3, 3), identity_grid.unsqueeze(-1)
            ).squeeze(-1)
            identity_grid = from_homogeneous(identity_grid)

        transformations = torch.cat([identity_grid, driving_to_source], dim=1)
        return transformations

    def create_deformed_source_image(self, source_image, transformations):
        bs, _, h, w = source_image.shape
        source_repeat = (
            source_image.unsqueeze(1)
            .unsqueeze(1)
            .repeat(1, self.num_tps + 1, 1, 1, 1, 1)
        )
        source_repeat = source_repeat.view(bs * (self.num_tps + 1), -1, h, w)
        transformations = transformations.view((bs * (self.num_tps + 1), h, w, -1))
        deformed = F.grid_sample(source_repeat, transformations, align_corners=True)
        deformed = deformed.view((bs, self.num_tps + 1, -1, h, w))
        return deformed

    def dropout_softmax(self, X, P):
        """
        Dropout for TPS transformations. Eq(7) and Eq(8) in the paper.
        """
        drop = (
            (torch.rand(X.shape[0], X.shape[1]) < (1 - P)).type(X.type()).to(X.device)
        )
        drop[..., 0] = 1
        drop = drop.repeat(X.shape[2], X.shape[3], 1, 1).permute(2, 3, 0, 1)

        maxx = X.max(1).values.unsqueeze_(1)
        X = X - maxx
        X_exp = X.exp()
        X[:, 1:, ...] /= 1 - P
        mask_bool = drop == 0
        X_exp = X_exp.masked_fill(mask_bool, 0)
        partition = X_exp.sum(dim=1, keepdim=True) + 1e-6
        return X_exp / partition

    def forward(
        self,
        source_image,
        kp_driving,
        kp_source,
        bg_param=None,
        dropout_flag=False,
        dropout_p=0,
    ):
        if self.scale_factor != 1:
            source_image = self.down(source_image)

        bs, _, h, w = source_image.shape

        out_dict = dict()
        heatmap_representation = self.create_heatmap_representations(
            source_image, kp_driving, kp_source
        )
        transformations = self.create_transformations(
            source_image, kp_driving, kp_source, bg_param
        )
        deformed_source = self.create_deformed_source_image(
            source_image, transformations
        )
        out_dict["deformed_source"] = deformed_source
        # out_dict['transformations'] = transformations
        deformed_source = deformed_source.view(bs, -1, h, w)
        input = torch.cat([heatmap_representation, deformed_source], dim=1)
        input = input.view(bs, -1, h, w)

        # prediction = self.hourglass(input, mode = 1)
        prediction = self.hourglass(input)

        # contribution_maps = self.maps(prediction[-1])
        contribution_maps = self.maps(prediction)
        if dropout_flag:
            contribution_maps = self.dropout_softmax(contribution_maps, dropout_p)
        else:
            contribution_maps = F.softmax(contribution_maps, dim=1)
        out_dict["contribution_maps"] = contribution_maps
        out_dict["mask"] = contribution_maps

        # Combine the K+1 transformations
        # Eq(6) in the paper
        contribution_maps = contribution_maps.unsqueeze(2)
        transformations = transformations.permute(0, 1, 4, 2, 3)
        deformation = (transformations * contribution_maps).sum(dim=1)
        deformation = deformation.permute(0, 2, 3, 1)

        out_dict["deformation"] = deformation  # Optical Flow

        occlusion_map = []
        if self.multi_mask:
            for i in range(self.occlusion_num - self.up_nums):
                occlusion_map.append(
                    torch.sigmoid(
                        self.occlusion[i](
                            prediction[self.up_nums - self.occlusion_num + i]
                        )
                    )
                )
            prediction = prediction[-1]
            for i in range(self.up_nums):
                prediction = self.up[i](prediction)
                occlusion_map.append(
                    torch.sigmoid(
                        self.occlusion[i + self.occlusion_num - self.up_nums](
                            prediction
                        )
                    )
                )
        else:
            # occlusion_map.append(torch.sigmoid(self.occlusion[0](prediction[-1])))
            occlusion_map = self.occlusion[0](prediction)

        out_dict["occlusion"] = occlusion_map  # Multi-resolution Occlusion Masks
        return out_dict