import torch


def close_form_affine(moving_kp, target_kp):
    """
    Obtain affine matrix to align moving keypoints to target keypoints.
    Affine matrix computed in a close form solution. 
    
    Arguments
    ---------
    moving_kp : keypoints from the moving image [n_batch, 3, n_keypoints]
    target_kp : keypoints from the fixed/target image [n_batch, 3, n_keypoints]

    Return
    ------
        out : affine matrix [n_batch, 3, 4]
    """
    Y_cm = moving_kp
    Y_tg = target_kp
    
    # Initialize 
    one = torch.ones(Y_cm.shape[0], 1, Y_cm.shape[2]).float() #Add a row of ones
    one = one.cuda() if Y_cm.is_cuda else one 
    _Y_cm = torch.cat([Y_cm, one],1)    
    
    out = torch.bmm(_Y_cm, torch.transpose(_Y_cm,-2,-1))
    out = torch.inverse(out)
    out = torch.bmm(torch.transpose(_Y_cm,-2,-1), out)
    out = torch.bmm(Y_tg, out)
    return out