#! /usr/bin/env python
# -*- coding: utf-8 -*-

"""
Image deformation using moving least squares.

    * Affine deformation
    * Similarity deformation
    * Rigid deformation

For more details please refer to the Chinese documentation: 
    
    ./doc/Image Deformation.pdf

or the original paper: 
    
    Image deformation using moving least squares
    Schaefer, Mcphail, Warren. 

Note:
    In the original paper, the author missed the weight w_j in formular (5).
    In addition, all the formulars in section 2.1 miss the w_j. 
    And I have corrected this point in my documentation.

@author: Jian-Wei ZHANG
@email: zjw.cs@zju.edu.cn
@date: 2022/01/12: PyTorch implementation
"""

import torch

from interp_torch import interp


def mls_affine_deformation(vy, vx, p, q, alpha=1.0, eps=1e-8):
    """
    Affine deformation

    Parameters
    ----------
    vy, vx: torch.Tensor
        coordinate grid, generated by torch.meshgrid(gridX, gridY)
    p: torch.Tensor
        an array with size [n, 2], original control points, in (y, x) formats
    q: torch.Tensor
        an array with size [n, 2], final control points, in (y, x) formats
    alpha: float
        parameter used by weights
    eps: float
        epsilon
    
    Return
    ------
        A deformed image.
    """

    device = q.device
    # Change (x, y) to (row, col)
    q = q.short()
    p = p.short()

    # Exchange p and q and hence we transform destination pixels to the corresponding source pixels.
    p, q = q, p

    grow = vx.shape[0]  # grid rows
    gcol = vx.shape[1]  # grid cols
    ctrls = p.shape[0]  # control points

    # Precompute
    reshaped_p = p.reshape(ctrls, 2, 1, 1)                                              # [ctrls, 2, 1, 1]
    reshaped_v = torch.cat((vx.reshape(1, grow, gcol), vy.reshape(1, grow, gcol)), dim=0)      # [2, grow, gcol]

    w = 1.0 / (torch.sum((reshaped_p - reshaped_v).float() ** 2, dim=1) + eps) ** alpha    # [ctrls, grow, gcol]
    w /= torch.sum(w, dim=0, keepdim=True)                                               # [ctrls, grow, gcol]

    pstar = torch.zeros((2, grow, gcol), dtype=torch.float32).to(device)
    for i in range(ctrls):
        pstar += w[i] * reshaped_p[i]                                                   # [2, grow, gcol]

    phat = reshaped_p - pstar                                                           # [ctrls, 2, grow, gcol]
    phat = phat.reshape(ctrls, 2, 1, grow, gcol)                                        # [ctrls, 2, 1, grow, gcol]
    phat1 = phat.reshape(ctrls, 1, 2, grow, gcol)                                       # [ctrls, 1, 2, grow, gcol]
    reshaped_w = w.reshape(ctrls, 1, 1, grow, gcol)                                     # [ctrls, 1, 1, grow, gcol]
    pTwp = torch.zeros((2, 2, grow, gcol), dtype=torch.float32).to(device)
    for i in range(ctrls):
        pTwp += phat[i] * reshaped_w[i] * phat1[i]
    del phat1

    try:
        inv_pTwp = torch.inverse(pTwp.permute(2, 3, 0, 1))                            # [grow, gcol, 2, 2]
        flag = False
    except RuntimeError as error:
        # print(error)
        flag = True             
        det = torch.det(pTwp.permute(2, 3, 0, 1))                                 # [grow, gcol]
        det[det < 1e-8] = float("Inf")
        reshaped_det = det.reshape(1, 1, grow, gcol)                                    # [1, 1, grow, gcol]
        adjoint = pTwp[[[1, 0], [1, 0]], [[1, 1], [0, 0]], :, :]                        # [2, 2, grow, gcol]
        adjoint[[0, 1], [1, 0], :, :] = -adjoint[[0, 1], [1, 0], :, :]                  # [2, 2, grow, gcol]
        inv_pTwp = (adjoint / reshaped_det).permute(2, 3, 0, 1)                       # [grow, gcol, 2, 2]
    
    mul_left = reshaped_v - pstar                                                       # [2, grow, gcol]
    reshaped_mul_left = mul_left.reshape(1, 2, grow, gcol).permute(2, 3, 0, 1)        # [grow, gcol, 1, 2]
    mul_right = torch.mul(reshaped_w, phat, out=phat)                                 # [ctrls, 2, 1, grow, gcol]
    reshaped_mul_right = mul_right.permute(0, 3, 4, 1, 2)                             # [ctrls, grow, gcol, 2, 1]
    out_A = mul_right.reshape(2, ctrls, grow, gcol, 1, 1)[0]                            # [ctrls, grow, gcol, 1, 1]
    A = torch.matmul(torch.matmul(reshaped_mul_left, inv_pTwp), reshaped_mul_right, out=out_A)    # [ctrls, grow, gcol, 1, 1]
    A = A.reshape(ctrls, 1, grow, gcol)                                                 # [ctrls, 1, grow, gcol]
    del mul_right, reshaped_mul_right, phat

    # Calculate q
    reshaped_q = q.reshape((ctrls, 2, 1, 1))                                            # [ctrls, 2, 1, 1]
    qstar = torch.zeros((2, grow, gcol), dtype=torch.float32).to(device)
    for i in range(ctrls):
        qstar += w[i] * reshaped_q[i]                                                   # [2, grow, gcol]
    del w, reshaped_w

    # Get final image transfomer -- 3-D array
    transformers = torch.zeros((2, grow, gcol), dtype=torch.float32).to(device)
    for i in range(ctrls):
        transformers += A[i] * (reshaped_q[i] - qstar)
    transformers += qstar
    del A

    # Correct the points where pTwp is singular
    if flag:
        blidx = det == float('Inf')    # bool index
        transformers[0][blidx] = vx[blidx] + qstar[0][blidx] - pstar[0][blidx]
        transformers[1][blidx] = vy[blidx] + qstar[1][blidx] - pstar[1][blidx]

    # Removed the points outside the border
    transformers[transformers < 0] = 0
    transformers[0][transformers[0] > grow - 1] = 0
    transformers[1][transformers[1] > gcol - 1] = 0

    return transformers.long()


def mls_similarity_deformation(vy, vx, p, q, alpha=1.0, eps=1e-8):
    """ Similarity deformation
    
    Parameters
    ----------
    vx, vy: torch.Tensor
        coordinate grid, generated by torch.meshgrid(gridX, gridY)
    p: torch.Tensor
        an array with size [n, 2], original control points, in (y, x) formats
    q: torch.Tensor
        an array with size [n, 2], final control points, in (y, x) formats
    alpha: float
        parameter used by weights
    eps: float
        epsilon
    
    Return
    ------
        A deformed image.
    """
    device = q.device
    q = q.short()
    p = p.short()

    # Exchange p and q and hence we transform destination pixels to the corresponding source pixels.
    p, q = q, p

    grow = vx.shape[0]  # grid rows
    gcol = vx.shape[1]  # grid cols
    ctrls = p.shape[0]  # control points

    # Compute
    reshaped_p = p.reshape(ctrls, 2, 1, 1)                                              # [ctrls, 2, 1, 1]
    reshaped_v = torch.cat((vx.reshape(1, grow, gcol), vy.reshape(1, grow, gcol)), dim=0)      # [2, grow, gcol]
    
    w = 1.0 / (torch.sum((reshaped_p - reshaped_v).float() ** 2, dim=1) + eps) ** alpha    # [ctrls, grow, gcol]
    w /= torch.sum(w, dim=0, keepdim=True)                                               # [ctrls, grow, gcol]

    pstar = torch.zeros((2, grow, gcol), dtype=torch.float32).to(device)
    for i in range(ctrls):
        pstar += w[i] * reshaped_p[i]                                                   # [2, grow, gcol]

    phat = reshaped_p - pstar                                                           # [ctrls, 2, grow, gcol]
    reshaped_phat = phat.reshape(ctrls, 1, 2, grow, gcol)                               # [ctrls, 1, 2, grow, gcol]
    reshaped_w = w.reshape(ctrls, 1, 1, grow, gcol)                                     # [ctrls, 1, 1, grow, gcol]
    
    mu = torch.zeros((grow, gcol), dtype=torch.float32).to(device)
    for i in range(ctrls):
        mu += w[i] * (phat[i] ** 2).sum(0)
    reshaped_mu = mu.reshape(1, grow, gcol)                                             # [1, grow, gcol]

    vpstar = reshaped_v - pstar                                                         # [2, grow, gcol]
    reshaped_vpstar = vpstar.reshape(2, 1, grow, gcol)                                  # [2, 1, grow, gcol]
    neg_vpstar_verti = vpstar[[1, 0],...]                                               # [2, grow, gcol]
    neg_vpstar_verti[1,...] = -neg_vpstar_verti[1,...]                                  
    reshaped_neg_vpstar_verti = neg_vpstar_verti.reshape(2, 1, grow, gcol)              # [2, 1, grow, gcol]
    mul_right = torch.cat((reshaped_vpstar, reshaped_neg_vpstar_verti), dim=1)    # [2, 2, grow, gcol]

    # Calculate q
    reshaped_q = q.reshape((ctrls, 2, 1, 1))                                            # [ctrls, 2, 1, 1]
    qstar = torch.zeros((2, grow, gcol), dtype=torch.float32).to(device)
    for i in range(ctrls):
        qstar += w[i] * reshaped_q[i]                                                   # [2, grow, gcol]
    
    # Get final image transfomer -- 3-D array
    temp = torch.zeros((grow, gcol, 2), dtype=torch.float32).to(device)
    for i in range(ctrls):
        neg_phat_verti = phat[i, [1, 0]]                                                # [2, grow, gcol]
        neg_phat_verti[1] = -neg_phat_verti[1]
        reshaped_neg_phat_verti = neg_phat_verti.reshape(1, 2, grow, gcol)              # [1, 2, grow, gcol]
        mul_left = torch.cat((reshaped_phat[i], reshaped_neg_phat_verti), dim=0)  # [2, 2, grow, gcol]
        
        A = torch.matmul((reshaped_w[i] * mul_left).permute(2, 3, 0, 1), 
                      mul_right.permute(2, 3, 0, 1))                                  # [grow, gcol, 2, 2]
    
        qhat = reshaped_q[i] - qstar                                                    # [2, grow, gcol]
        reshaped_qhat = qhat.reshape(1, 2, grow, gcol).permute(2, 3, 0, 1)            # [grow, gcol, 1, 2]

        # Get final image transfomer -- 3-D array
        temp += torch.matmul(reshaped_qhat, A).reshape(grow, gcol, 2)                      # [grow, gcol, 2]

    transformers = temp.permute(2, 0, 1) / reshaped_mu  + qstar                       # [2, grow, gcol]

    # Removed the points outside the border
    transformers[transformers < 0] = 0
    transformers[0][transformers[0] > grow - 1] = 0
    transformers[1][transformers[1] > gcol - 1] = 0

    return transformers.long()


def mls_rigid_deformation(vy, vx, p, q, alpha=1.0, eps=1e-8):
    """ Rigid deformation
    
    Parameters
    ----------
    vx, vy: torch.Tensor
        coordinate grid, generated by torch.meshgrid(gridX, gridY)
    p: torch.Tensor
        an array with size [n, 2], original control points, in (y, x) formats
    q: torch.Tensor
        an array with size [n, 2], final control points, in (y, x) formats
    alpha: float
        parameter used by weights
    eps: float
        epsilon
    
    Return
    ------
        A deformed image.
    """
    device = q.device
    q = q.short()
    p = p.short()

    # Exchange p and q and hence we transform destination pixels to the corresponding source pixels.
    p, q = q, p

    grow = vx.shape[0]  # grid rows
    gcol = vx.shape[1]  # grid cols
    ctrls = p.shape[0]  # control points

    # Compute
    reshaped_p = p.reshape(ctrls, 2, 1, 1)                                              # [ctrls, 2, 1, 1]
    reshaped_v = torch.cat((vx.reshape(1, grow, gcol), vy.reshape(1, grow, gcol)), dim=0)      # [2, grow, gcol]
    
    w = 1.0 / (torch.sum((reshaped_p - reshaped_v).float() ** 2, dim=1) + eps) ** alpha    # [ctrls, grow, gcol]
    w /= torch.sum(w, dim=0, keepdim=True)                                               # [ctrls, grow, gcol]
    
    pstar = torch.zeros((2, grow, gcol), dtype=torch.float32).to(device)
    for i in range(ctrls):
        pstar += w[i] * reshaped_p[i]                                                   # [2, grow, gcol]

    vpstar = reshaped_v - pstar                                                         # [2, grow, gcol]
    reshaped_vpstar = vpstar.reshape(2, 1, grow, gcol)                                  # [2, 1, grow, gcol]
    neg_vpstar_verti = vpstar[[1, 0],...]                                               # [2, grow, gcol]
    neg_vpstar_verti[1,...] = -neg_vpstar_verti[1,...]                                  
    reshaped_neg_vpstar_verti = neg_vpstar_verti.reshape(2, 1, grow, gcol)              # [2, 1, grow, gcol]
    mul_right = torch.cat((reshaped_vpstar, reshaped_neg_vpstar_verti), dim=1)    # [2, 2, grow, gcol]
    reshaped_mul_right = mul_right.reshape(2, 2, grow, gcol)                            # [2, 2, grow, gcol]

    # Calculate q
    reshaped_q = q.reshape((ctrls, 2, 1, 1))                                            # [ctrls, 2, 1, 1]
    qstar = torch.zeros((2, grow, gcol), dtype=torch.float32).to(device)
    for i in range(ctrls):
        qstar += w[i] * reshaped_q[i]                                                   # [2, grow, gcol]
    
    temp = torch.zeros((grow, gcol, 2), dtype=torch.float32).to(device)
    for i in range(ctrls):
        phat = reshaped_p[i] - pstar                                                    # [2, grow, gcol]
        reshaped_phat = phat.reshape(1, 2, grow, gcol)                                  # [1, 2, grow, gcol]
        reshaped_w = w[i].reshape(1, 1, grow, gcol)                                     # [1, 1, grow, gcol]
        neg_phat_verti = phat[[1, 0]]                                                   # [2, grow, gcol]
        neg_phat_verti[1] = -neg_phat_verti[1]
        reshaped_neg_phat_verti = neg_phat_verti.reshape(1, 2, grow, gcol)              # [1, 2, grow, gcol]
        mul_left = torch.cat((reshaped_phat, reshaped_neg_phat_verti), dim=0)     # [2, 2, grow, gcol]
        
        A = torch.matmul((reshaped_w * mul_left).permute(2, 3, 0, 1), 
                        reshaped_mul_right.permute(2, 3, 0, 1))                       # [grow, gcol, 2, 2]

        qhat = reshaped_q[i] - qstar                                                    # [2, grow, gcol]
        reshaped_qhat = qhat.reshape(1, 2, grow, gcol).permute(2, 3, 0, 1)            # [grow, gcol, 1, 2]

        # Get final image transfomer -- 3-D array
        temp += torch.matmul(reshaped_qhat, A).reshape(grow, gcol, 2)                      # [grow, gcol, 2]

    temp = temp.permute(2, 0, 1)                                                      # [2, grow, gcol]
    normed_temp = torch.norm(temp, dim=0, keepdim=True)                           # [1, grow, gcol]
    normed_vpstar = torch.norm(vpstar, dim=0, keepdim=True)                       # [1, grow, gcol]
    transformers = temp / normed_temp * normed_vpstar  + qstar                          # [2, grow, gcol]
    nan_mask = normed_temp[0] == 0

    # Replace nan values by interpolated values
    nan_mask_flat = torch.nonzero(nan_mask.view(-1), as_tuple=True)[0]
    nan_mask_anti_flat = torch.nonzero(~nan_mask.view(-1), as_tuple=True)[0]
    transformers[0][nan_mask] = interp(nan_mask_flat, nan_mask_anti_flat, transformers[0][~nan_mask])
    transformers[1][nan_mask] = interp(nan_mask_flat, nan_mask_anti_flat, transformers[1][~nan_mask])

    # Remove the points outside the border
    transformers[transformers < 0] = 0
    transformers[0][transformers[0] > grow - 1] = 0
    transformers[1][transformers[1] > gcol - 1] = 0
    
    return transformers.long()
