"""Padded transformation module.

This module provides two functions, warpPerspectivePadded() and
warpAffinePadded(), which compliment the built-in OpenCV functions
warpPerspective() and warpAffine(). These functions calculate the
extent of the warped image and pads both the destination and the
warped image so both images can be fully displayed together.
"""


import cv2
import numpy as np


def warpPerspectivePadded(
    src, dst, M, flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=0
):
    """Performs a perspective warp with padding.

    Parameters
    ----------
    src : array_like
        source image, to be warped.
    dst : array_like
        destination image, to be padded.
    M : array_like
        `3x3` perspective transformation matrix.

    Returns
    -------
    src_warped : ndarray
        padded and warped source image
    dst_padded : ndarray
        padded destination image, same size as src_warped

    Optional Parameters
    -------------------
    flags : int, optional
        combination of interpolation methods (`cv2.INTER_LINEAR` or
        `cv2.INTER_NEAREST`) and the optional flag `cv2.WARP_INVERSE_MAP`,
        that sets `M` as the inverse transformation (`dst` --> `src`).
    borderMode : int, optional
        pixel extrapolation method (`cv2.BORDER_CONSTANT` or
        `cv2.BORDER_REPLICATE`).
    borderValue : numeric, optional
        value used in case of a constant border; by default, it equals 0.

    See Also
    --------
    warpAffinePadded : for `2x3` affine transformations
    cv2.warpPerspective : original OpenCV function
    """

    assert M.shape == (3, 3), (
        "Perspective transformation shape should be (3, 3).\n"
        + "Use warpAffinePadded() for (2, 3) affine transformations."
    )

    M = M / M[2, 2]  # ensure a legal homography
    if flags in (
        cv2.WARP_INVERSE_MAP,
        cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP,
        cv2.INTER_NEAREST + cv2.WARP_INVERSE_MAP,
    ):
        M = cv2.invert(M)[1]
        flags -= cv2.WARP_INVERSE_MAP

    # it is enough to find where the corners of the image go to find
    # the padding bounds; points in clockwise order from origin
    src_h, src_w = src.shape[:2]
    lin_homg_pts = np.array([[0, src_w, src_w, 0], [0, 0, src_h, src_h], [1, 1, 1, 1]])

    # transform points
    transf_lin_homg_pts = M.dot(lin_homg_pts)
    transf_lin_homg_pts /= transf_lin_homg_pts[2, :]

    # find min and max points
    min_x = np.floor(np.min(transf_lin_homg_pts[0])).astype(int)
    min_y = np.floor(np.min(transf_lin_homg_pts[1])).astype(int)
    max_x = np.ceil(np.max(transf_lin_homg_pts[0])).astype(int)
    max_y = np.ceil(np.max(transf_lin_homg_pts[1])).astype(int)

    # add translation to the transformation matrix to shift to positive values
    anchor_x, anchor_y = 0, 0
    transl_transf = np.eye(3, 3)
    if min_x < 0:
        anchor_x = -min_x
        transl_transf[0, 2] += anchor_x
    if min_y < 0:
        anchor_y = -min_y
        transl_transf[1, 2] += anchor_y
    shifted_transf = transl_transf.dot(M)
    shifted_transf /= shifted_transf[2, 2]

    # create padded destination image
    dst_h, dst_w = dst.shape[:2]

    pad_widths = [
        anchor_y,
        max(max_y, dst_h) - dst_h,
        anchor_x,
        max(max_x, dst_w) - dst_w,
    ]

    dst_padded = cv2.copyMakeBorder(
        dst, *pad_widths, borderType=borderMode, value=borderValue
    )

    dst_pad_h, dst_pad_w = dst_padded.shape[:2]
    src_warped = cv2.warpPerspective(
        src,
        shifted_transf,
        (dst_pad_w, dst_pad_h),
        flags=flags,
        borderMode=borderMode,
        borderValue=borderValue,
    )

    return dst_padded, src_warped


def warpAffinePadded(
    src, dst, M, flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=0
):
    """Performs an affine or Euclidean/rigid warp with padding.

    Parameters
    ----------
    src : array_like
        source image, to be warped.
    dst : array_like
        destination image, to be padded.
    M : array_like
        `2x3` affine transformation matrix.

    Returns
    -------
    src_warped : ndarray
        padded and warped source image
    dst_padded : ndarray
        padded destination image, same size as src_warped

    Optional Parameters
    -------------------
    flags : int, optional
        combination of interpolation methods (`cv2.INTER_LINEAR` or
        `cv2.INTER_NEAREST`) and the optional flag `cv2.WARP_INVERSE_MAP`,
        that sets `M` as the inverse transformation (`dst` --> `src`).
    borderMode : int, optional
        pixel extrapolation method (`cv2.BORDER_CONSTANT` or
        `cv2.BORDER_REPLICATE`).
    borderValue : numeric, optional
        value used in case of a constant border; by default, it equals 0.

    See Also
    --------
    warpPerspectivePadded : for `3x3` perspective transformations
    cv2.warpAffine : original OpenCV function
    """
    assert M.shape == (2, 3), (
        "Affine transformation shape should be (2, 3).\n"
        + "Use warpPerspectivePadded() for (3, 3) homography transformations."
    )

    if flags in (
        cv2.WARP_INVERSE_MAP,
        cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP,
        cv2.INTER_NEAREST + cv2.WARP_INVERSE_MAP,
    ):
        M = cv2.invertAffineTransform(M)
        flags -= cv2.WARP_INVERSE_MAP

    # it is enough to find where the corners of the image go to find
    # the padding bounds; points in clockwise order from origin
    src_h, src_w = src.shape[:2]
    lin_pts = np.array([[0, src_w, src_w, 0], [0, 0, src_h, src_h]])

    # transform points
    transf_lin_pts = M[:, :2].dot(lin_pts) + M[:, 2].reshape(2, 1)

    # find min and max points
    min_x = np.floor(np.min(transf_lin_pts[0])).astype(int)
    min_y = np.floor(np.min(transf_lin_pts[1])).astype(int)
    max_x = np.ceil(np.max(transf_lin_pts[0])).astype(int)
    max_y = np.ceil(np.max(transf_lin_pts[1])).astype(int)

    # add translation to the transformation matrix to shift to positive values
    anchor_x, anchor_y = 0, 0
    if min_x < 0:
        anchor_x = -min_x
    if min_y < 0:
        anchor_y = -min_y
    shifted_transf = M + [[0, 0, anchor_x], [0, 0, anchor_y]]

    # create padded destination image
    dst_h, dst_w = dst.shape[:2]

    pad_widths = [
        anchor_y,
        max(max_y, dst_h) - dst_h,
        anchor_x,
        max(max_x, dst_w) - dst_w,
    ]

    dst_padded = cv2.copyMakeBorder(
        dst, *pad_widths, borderType=borderMode, value=borderValue
    )

    dst_pad_h, dst_pad_w = dst_padded.shape[:2]
    src_warped = cv2.warpAffine(
        src,
        shifted_transf,
        (dst_pad_w, dst_pad_h),
        flags=flags,
        borderMode=borderMode,
        borderValue=borderValue,
    )

    return dst_padded, src_warped


"""
References
----------
See the following question and my answer on Stack Overflow for an
idea of how this was conceptualized and to read the mathematics
behind the functions: https://stackoverflow.com/a/44459869/5087436

"""
