import numpy as np
import torch
import torch.nn as nn

from ..registry import TPSWARP


@TPSWARP.register_module
class TPSWarp(nn.Module):

    def __init__(self,
                 out_h=256,
                 out_w=192,
                 grid_size=3,
                 use_regular_grid=True,
                 reg_factor=0):
        super(TPSWarp, self).__init__()
        self.out_h = out_h
        self.out_w = out_w
        self.reg_factor = reg_factor

        # create grid in numpy
        self.grid = np.zeros((self.out_h, self.out_w, 3), dtype=np.float32)
        # sampling grid with dim-0 (Y)
        self.gridX, self.gridY = np.meshgrid(
            np.linspace(-1, 1, out_w), np.linspace(-1, 1, out_h))
        # gridX, gridY: size [1, H, W, 1, 1]
        self.gridX = torch.FloatTensor(self.gridX).unsqueeze(0).unsqueeze(3)
        self.gridY = torch.FloatTensor(self.gridY).unsqueeze(0).unsqueeze(3)
        # put on gpu
        self.gridX = self.gridX.cuda()
        self.gridY = self.gridY.cuda()

        # init regular grid for control points P_i
        if use_regular_grid:
            axis_coords = np.linspace(-1, 1, grid_size)
            self.N = grid_size * grid_size
            P_Y, P_X = np.meshgrid(axis_coords, axis_coords)
            P_X = np.reshape(P_X, (-1, 1))  # size (N,1)
            P_Y = np.reshape(P_Y, (-1, 1))  # size (N,1)
            P_X = torch.FloatTensor(P_X)
            P_Y = torch.FloatTensor(P_Y)
            self.P_X_base = P_X.clone()
            self.P_Y_base = P_Y.clone()
            self.Li = self.compute_L_inverse(P_X, P_Y).unsqueeze(0)
            self.P_X = P_X.unsqueeze(2).unsqueeze(3).unsqueeze(4).transpose(
                0, 4)
            self.P_Y = P_Y.unsqueeze(2).unsqueeze(3).unsqueeze(4).transpose(
                0, 4)
            # put on gpu
            self.P_X = self.P_X.cuda()
            self.P_X_base = self.P_X_base.cuda()
            self.P_Y = self.P_Y.cuda()
            self.P_Y_base = self.P_Y_base.cuda()

    def compute_L_inverse(self, X, Y):
        N = X.size()[0]  # num of points (along dim 0)
        # construct matrix K
        Xmat = X.expand(N, N)
        Ymat = Y.expand(N, N)
        P_dist_squared = torch.pow(Xmat - Xmat.transpose(0, 1), 2) + torch.pow(
            Ymat - Ymat.transpose(0, 1), 2)
        # make diagonal 1 to avoid NaN in log computation
        P_dist_squared[P_dist_squared == 0] = 1
        K = torch.mul(P_dist_squared, torch.log(P_dist_squared))

        # construct matrix L
        one = torch.FloatTensor(N, 1).fill_(1)
        Z = torch.FloatTensor(3, 3).fill_(0)
        P = torch.cat((one, X, Y), 1)
        L = torch.cat((torch.cat(
            (K, P), 1), torch.cat((P.transpose(0, 1), Z), 1)), 0)
        Li = torch.inverse(L)
        Li = Li.cuda()
        return Li

    def apply_transformation(self, theta, points):
        if theta.dim() == 2:
            theta = theta.unsqueeze(2).unsqueeze(3)

        # points should be in the [B,H,W,2] format,
        # where points[:,:,:,0] are the X coords
        # and points[:,:,:,1] are the Y coords
        # input are the corresponding control points P_i
        batch_size = theta.size(0)
        # split theta into point coords
        Q_X = theta[:, :self.N, :, :].squeeze(3)
        Q_Y = theta[:, self.N:, :, :].squeeze(3)
        Q_X = Q_X + self.P_X_base.expand_as(Q_X)
        Q_Y = Q_Y + self.P_Y_base.expand_as(Q_Y)

        # get spatial dimensions of points
        points_b = points.size(0)
        points_h = points.size(1)
        points_w = points.size(2)

        # repeat pre-defined control points along
        # spatial dimensions of points to be transformed
        P_X = self.P_X.expand((1, points_h, points_w, 1, self.N))
        P_Y = self.P_Y.expand((1, points_h, points_w, 1, self.N))

        # compute weigths for non-linear part
        W_X = torch.bmm(
            self.Li[:, :self.N, :self.N].expand((batch_size, self.N, self.N)),
            Q_X)
        W_Y = torch.bmm(
            self.Li[:, :self.N, :self.N].expand((batch_size, self.N, self.N)),
            Q_Y)
        # reshape
        # W_X,W,Y: size [B,H,W,1,N]
        W_X = W_X.unsqueeze(3).unsqueeze(4).transpose(1, 4).repeat(
            1, points_h, points_w, 1, 1)
        W_Y = W_Y.unsqueeze(3).unsqueeze(4).transpose(1, 4).repeat(
            1, points_h, points_w, 1, 1)
        # compute weights for affine part
        A_X = torch.bmm(
            self.Li[:, self.N:, :self.N].expand((batch_size, 3, self.N)), Q_X)
        A_Y = torch.bmm(
            self.Li[:, self.N:, :self.N].expand((batch_size, 3, self.N)), Q_Y)
        # reshape
        # A_X,A,Y: size [B,H,W,1,3]
        A_X = A_X.unsqueeze(3).unsqueeze(4).transpose(1, 4).repeat(
            1, points_h, points_w, 1, 1)
        A_Y = A_Y.unsqueeze(3).unsqueeze(4).transpose(1, 4).repeat(
            1, points_h, points_w, 1, 1)

        # compute distance P_i - (grid_X,grid_Y)
        # grid is expanded in point dim 4, but not in batch dim 0,
        # as points P_X,P_Y are fixed for all batch
        points_X_for_summation = points[:, :, :, 0].unsqueeze(3).unsqueeze(
            4).expand(points[:, :, :, 0].size() + (1, self.N))
        points_Y_for_summation = points[:, :, :, 1].unsqueeze(3).unsqueeze(
            4).expand(points[:, :, :, 1].size() + (1, self.N))

        if points_b == 1:
            delta_X = points_X_for_summation - P_X
            delta_Y = points_Y_for_summation - P_Y
        else:
            # use expanded P_X,P_Y in batch dimension
            delta_X = points_X_for_summation - P_X.expand_as(
                points_X_for_summation)
            delta_Y = points_Y_for_summation - P_Y.expand_as(
                points_Y_for_summation)

        dist_squared = torch.pow(delta_X, 2) + torch.pow(delta_Y, 2)
        # U: size [1,H,W,1,N]
        dist_squared[dist_squared == 0] = 1  # avoid NaN in log computation
        U = torch.mul(dist_squared, torch.log(dist_squared))

        # expand grid in batch dimension if necessary
        points_X_batch = points[:, :, :, 0].unsqueeze(3)
        points_Y_batch = points[:, :, :, 1].unsqueeze(3)
        if points_b == 1:
            points_X_batch = points_X_batch.expand((batch_size, ) +
                                                   points_X_batch.size()[1:])
            points_Y_batch = points_Y_batch.expand((batch_size, ) +
                                                   points_Y_batch.size()[1:])

        points_X_prime = A_X[:, :, :, :, 0] + torch.mul(
            A_X[:, :, :, :, 1], points_X_batch) + torch.mul(
                A_X[:, :, :, :, 2], points_Y_batch) + torch.sum(
                    torch.mul(W_X, U.expand_as(W_X)), 4)

        points_Y_prime = A_Y[:, :, :, :, 0] + torch.mul(
            A_Y[:, :, :, :, 1], points_X_batch) + torch.mul(
                A_Y[:, :, :, :, 2], points_Y_batch) + torch.sum(
                    torch.mul(W_Y, U.expand_as(W_Y)), 4)

        return torch.cat((points_X_prime, points_Y_prime), 3)

    def forward(self, theta):
        points = torch.cat((self.gridX, self.gridY), 3)
        warped_grid = self.apply_transformation(theta, points)
        return warped_grid
