# utils.py

# Copyright (c) 2014-?, Matěj Týč
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in the
#   documentation and/or other materials provided with the distribution.
# * Neither the name of the copyright holders nor the names of any
#   contributors may be used to endorse or promote products derived
#   from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

"""
FFT based image registration. --- utility functions
"""

import numpy as np
import numpy.fft as fft
import scipy.ndimage as ndi


def wrap_angle(angles, ceil=2 * np.pi):
    """
    Parameters
    ----------
    angles (float or ndarray, unit depends on kwarg ``ceil``)
    ceil (float)
        Turnaround value
    """
    angles += ceil / 2.0
    angles %= ceil
    angles -= ceil / 2.0
    return angles


def rot180(arr):
    """
    Rotate the input array over 180°
    """
    ret = np.rot90(arr, 2)
    return ret


def _get_angles(shape):
    """
    In the log-polar spectrum, the (first) coord corresponds to an angle.
    This function returns a mapping of (the two) coordinates
    to the respective angle.
    """
    ret = np.zeros(shape, dtype=np.float64)
    ret -= np.linspace(0, np.pi, shape[0], endpoint=False)[:, np.newaxis]
    return ret


def _get_lograd(shape, log_base):
    """
    In the log-polar spectrum, the (second) coord corresponds to an angle.
    This function returns a mapping of (the two) coordinates
    to the respective scale.

    Returns
    -------
    2D np.ndarray of shape ``shape``, -1 coord contains scales
    from 0 to log_base ** (shape[1] - 1)
    """
    ret = np.zeros(shape, dtype=np.float64)
    ret += np.power(log_base, np.arange(shape[1], dtype=float))[np.newaxis, :]
    return ret


def _get_constraint_mask(shape, log_base, constraints=None):
    """
    Prepare mask to apply to constraints to a cross-power spectrum.
    """
    if constraints is None:
        constraints = {}

    mask = np.ones(shape, float)

    # Here, we create masks that modulate picking the best correspondence.
    # Generally, we look at the log-polar array and identify mapping of
    # coordinates to values of quantities.
    if "scale" in constraints:
        scale, sigma = constraints["scale"]
        scales = fft.ifftshift(_get_lograd(shape, log_base))
        # vvv This issome kind of transformation of result of _get_lograd
        # vvv (log radius in pixels) to the linear scale.
        scales *= log_base ** (-shape[1] / 2.0)
        # This makes the scales array low near where scales is near 'scale'
        scales -= 1.0 / scale
        if sigma == 0:
            # there isn't: ascales = np.abs(scales - scale)
            # because scales are already low for values near 'scale'
            ascales = np.abs(scales)
            scale_min = ascales.min()
            mask[ascales > scale_min] = 0
        elif sigma is None:
            pass
        else:
            mask *= np.exp(-(scales**2) / sigma**2)

    if "angle" in constraints:
        angle, sigma = constraints["angle"]
        angles = _get_angles(shape)
        # We flip the sign on purpose
        # TODO: ^^^ Why???
        angles += np.deg2rad(angle)
        # TODO: Check out the wrapping. It may be tricky since pi+1 != 1
        wrap_angle(angles, np.pi)
        angles = np.rad2deg(angles)
        if sigma == 0:
            aangles = np.abs(angles)
            angle_min = aangles.min()
            mask[aangles > angle_min] = 0
        elif sigma is None:
            pass
        else:
            mask *= np.exp(-(angles**2) / sigma**2)

    mask = fft.fftshift(mask)
    return mask


def argmax_angscale(array, log_base, exponent, constraints=None):
    """
    Given a power spectrum, we choose the best fit.

    The power spectrum is treated with constraint masks and then
    passed to `_argmax_ext`.
    """
    mask = _get_constraint_mask(array.shape, log_base, constraints)
    array_orig = array.copy()

    array *= mask
    ret = _argmax_ext(array, exponent)
    ret_final = _interpolate(array, ret)

    success = _get_success(array_orig, tuple(ret_final), 0)

    return ret_final, success


def argmax_translation(array, filter_pcorr, constraints=None):
    if constraints is None:
        constraints = dict(tx=(0, None), ty=(0, None))

    # We want to keep the original and here is obvious that
    # it won't get changed inadvertently
    array_orig = array.copy()
    if filter_pcorr > 0:
        array = ndi.minimum_filter(array, filter_pcorr)

    ashape = np.array(array.shape, int)
    mask = np.ones(ashape, float)
    # first goes Y, then X
    for dim, key in enumerate(("ty", "tx")):
        if constraints.get(key, (0, None))[1] is None:
            continue
        pos, sigma = constraints[key]
        alen = ashape[dim]
        dom = np.linspace(-alen // 2, -alen // 2 + alen, alen, False)
        if sigma == 0:
            # generate a binary array closest to the position
            idx = np.argmin(np.abs(dom - pos))
            vals = np.zeros(dom.size)
            vals[idx] = 1.0
        else:
            vals = np.exp(-((dom - pos) ** 2) / sigma**2)
        if dim == 0:
            mask *= vals[:, np.newaxis]
        else:
            mask *= vals[np.newaxis, :]

    array *= mask

    # WE ARE FFTSHIFTED already.
    # ban translations that are too big
    aporad = (ashape // 6).min()
    mask2 = get_apofield(ashape, aporad)
    array *= mask2
    # Find what we look for
    tvec = _argmax_ext(array, "inf")
    tvec = _interpolate(array_orig, tvec)

    # If we use constraints or min filter,
    # array_orig[tvec] may not be the maximum
    success = _get_success(array_orig, tuple(tvec), 2)

    return tvec, success


def _get_success(array, coord, radius=2):
    """
    Given a coord, examine the array around it and return a number signifying
    how good is the "match".

    Args:
        radius: Get the success as a sum of neighbor of coord of this radius
        coord: Coordinates of the maximum. Float numbers are allowed
            (and converted to int inside)

    Returns:
        Success as float between 0 and 1 (can get slightly higher than 1).
        The meaning of the number is loose, but the higher the better.
    """
    coord = np.round(coord).astype(int)
    coord = tuple(coord)

    subarr = _get_subarr(array, coord, radius)

    theval = subarr.sum()
    theval2 = array[coord]
    # bigval = np.percentile(array, 97)
    # success = theval / bigval
    # TODO: Think this out
    success = np.sqrt(theval * theval2)
    return success


def _argmax2D(array):
    """
    Simple 2D argmax function with simple sharpness indication
    """
    amax = np.argmax(array)
    ret = list(np.unravel_index(amax, array.shape))

    return np.array(ret)


def _get_subarr(array, center, rad):
    """
    Parameters
    ----------
    array (ndarray)
        The array to search
    center (2-tuple)
        The point in the array to search around
    rad (int)
        Search radius, no radius (i.e. get the single point) implies rad == 0
    """
    dim = 1 + 2 * rad
    subarr = np.zeros((dim,) * 2)
    corner = np.array(center) - rad
    for ii in range(dim):
        yidx = corner[0] + ii
        yidx %= array.shape[0]
        for jj in range(dim):
            xidx = corner[1] + jj
            xidx %= array.shape[1]
            subarr[ii, jj] = array[yidx, xidx]
    return subarr


def _interpolate(array, rough, rad=2):
    """
    Returns index that is in the array after being rounded.

    The result index tuple is in each of its components between zero and the
    array's shape.
    """
    rough = np.round(rough).astype(int)
    surroundings = _get_subarr(array, rough, rad)
    com = _argmax_ext(surroundings, 1)
    offset = com - rad
    ret = rough + offset
    # similar to win.wrap, so
    # -0.2 becomes 0.3 and then again -0.2, which is rounded to 0
    # -0.8 becomes - 0.3 -> len() - 0.3 and then len() - 0.8,
    # which is rounded to len() - 1. Yeah!
    ret += 0.5
    ret %= np.array(array.shape).astype(int)
    ret -= 0.5
    return ret


def _argmax_ext(array, exponent):
    """
    Calculate coordinates of the COM (center of mass) of the provided array.

    Args:
        array (ndarray): The array to be examined.
        exponent (float or 'inf'): The exponent we power the array with. If the
            value 'inf' is given, the coordinage of the array maximum is taken.

    Returns:
        np.ndarray: The COM coordinate tuple, float values are allowed!
    """

    # When using an integer exponent for _argmax_ext, it is good to have the
    # neutral rotation/scale in the center rather near the edges

    ret = None
    if exponent == "inf":
        ret = _argmax2D(array)
    else:
        col = np.arange(array.shape[0])[:, np.newaxis]
        row = np.arange(array.shape[1])[np.newaxis, :]

        arr2 = array**exponent
        arrsum = arr2.sum()
        if arrsum == 0:
            # We have to return SOMETHING, so let's go for (0, 0)
            return np.zeros(2)
        arrprody = np.sum(arr2 * col) / arrsum
        arrprodx = np.sum(arr2 * row) / arrsum
        ret = [arrprody, arrprodx]
        # We don't use it, but it still tells us about value distribution

    return np.array(ret)


def imfilter(img, low=None, high=None, cap=None):
    """
    Given an image, it a high-pass and/or low-pass filters on its
    Fourier spectrum.

    Parameters
    ----------
    img (ndarray)
        The image to be filtered
    low (tuple)
        The low-pass filter parameters, 0..1
    high (tuple)
        The high-pass filter parameters, 0..1
    cap (tuple)
        The quantile cap parameters, 0..1. A filtered image will have extremes below
        the lower quantile and above the upper one cut.

    Returns
    -------
    np.ndarray
    The real component of the image after filtering
    """
    dft = fft.fft2(img)

    if low is not None:
        _lowpass(dft, low[0], low[1])
    if high is not None:
        _highpass(dft, high[0], high[1])

    ret = fft.ifft2(dft)
    # if the input was a real number array, return real numbers,
    # otherwise let it be complex.
    if not np.iscomplexobj(img):
        ret = np.real(ret)

    if cap is None:
        cap = (0, 1)

    low, high = cap
    if low > 0.0:
        low_val = np.percentile(ret, low * 100.0)
        ret[ret < low_val] = low_val
    if high < 1.0:
        high_val = np.percentile(ret, high * 100.0)
        ret[ret > high_val] = high_val

    return ret


def _highpass(dft, lo, hi):
    mask = _xpass((dft.shape), lo, hi)
    dft *= 1 - mask


def _lowpass(dft, lo, hi):
    mask = _xpass((dft.shape), lo, hi)
    dft *= mask


def _xpass(shape, lo, hi):
    """
    Compute a pass-filter mask with values ranging from 0 to 1.0
    The mask is low-pass, application has to be handled by a calling funcion.
    """
    assert lo <= hi, "Filter order wrong, low '%g', high '%g'" % (lo, hi)
    assert lo >= 0, "Low filter lower than zero (%g)" % lo
    # High can be as high as possible

    dom_x = np.fft.fftfreq(shape[0])[:, np.newaxis]
    dom_y = np.fft.fftfreq(shape[1])[np.newaxis, :]

    # freq goes 0..0.5, we want from 0..1, so we multiply it by 2.
    dom = np.sqrt(dom_x**2 + dom_y**2) * 2

    res = np.ones(dom.shape)
    res[dom >= hi] = 0.0
    mask = (dom > lo) * (dom < hi)
    res[mask] = 1 - (dom[mask] - lo) / (hi - lo)

    return res


def _apodize(what, aporad=None, ratio=None):
    """
    Given an image, it apodizes it (so it becomes quasi-seamless).
    When ``ratio`` is None, color near the edges will converge
    to the same colour, whereas when ratio is a float number, a blurred
    original image will serve as background.

    Args:
        what: The original image
        aporad (int): Radius [px], width of the band near the edges
            that will get modified
        ratio (float or None): When None, the apodization background will
            be a flat color.
            When a float number, the background will be the image itself
            convolved with Gaussian kernel of sigma (aporad / ratio).

    Returns:
        The apodized image
    """
    if aporad is None:
        mindim = min(what.shape)
        aporad = int(mindim * 0.12)
    apofield = get_apofield(what.shape, aporad)
    res = what * apofield
    if ratio is not None:
        ratio = float(ratio)
        bg = ndi.gaussian_filter(what, aporad / ratio, mode="wrap")
    else:
        bg = get_borderval(what, aporad // 2)
    res += bg * (1 - apofield)
    return res


def get_apofield(shape, aporad):
    """
    Returns an array between 0 and 1 that goes to zero close to the edges.
    """
    if aporad == 0:
        return np.ones(shape, dtype=float)
    apos = np.hanning(aporad * 2)
    vecs = []
    for dim in shape:
        assert dim > aporad * 2, "Apodization radius %d too big for shape dim. %d" % (
            aporad,
            dim,
        )
        toapp = np.ones(dim)
        toapp[:aporad] = apos[:aporad]
        toapp[-aporad:] = apos[-aporad:]
        vecs.append(toapp)
    apofield = np.outer(vecs[0], vecs[1])
    return apofield


def frame_img(img, mask, dst, apofield=None):
    """
    Given an array, a mask (floats between 0 and 1), and a distance,
    alter the area where the mask is low (and roughly within dst from the edge)
    so it blends well with the area where the mask is high.
    The purpose of this is removal of spurious frequencies in the image's
    Fourier spectrum.

    Args:
        img (np.array): What we want to alter
        maski (np.array): The indicator what can be altered (0)
            and what can not (1)
        dst (int): Parameter controlling behavior near edges, value could be
            probably deduced from the mask.
    """
    import scipy.ndimage as ndimg

    radius = dst / 1.8

    convmask0 = mask + 1e-10

    krad_max = radius * 6
    convimg = img
    convmask = convmask0
    convimg0 = img
    krad0 = 0.8
    krad = krad0

    while krad < krad_max:
        convimg = ndimg.gaussian_filter(convimg0 * convmask0, krad, mode="wrap")
        convmask = ndimg.gaussian_filter(convmask0, krad, mode="wrap")
        convimg /= convmask

        convimg = convimg * (convmask - convmask0) + convimg0 * (
            1 - convmask + convmask0
        )
        krad *= 1.8

        convimg0 = convimg
        convmask0 = convmask

    if apofield is not None:
        ret = convimg * (1 - apofield) + img * apofield
    else:
        ret = convimg
        ret[mask >= 1] = img[mask >= 1]

    return ret


def get_borderval(img, radius=None):
    """
    Given an image and a radius, examine the average value of the image
    at most radius pixels from the edge
    """
    if radius is None:
        mindim = min(img.shape)
        radius = max(1, mindim // 20)
    mask = np.zeros_like(img, dtype=bool)
    mask[:, :radius] = True
    mask[:, -radius:] = True
    mask[:radius, :] = True
    mask[-radius:, :] = True

    mean = np.median(img[mask])
    return mean


def transform_2d_coord_arrays(homography, in_coords):
    aug_coords = np.dstack([in_coords, np.ones(in_coords.shape[:2])]).squeeze()
    out_coordinates = np.einsum("ik,abk", homography, aug_coords)
    out_coordinates = (out_coordinates / out_coordinates[:, :, 2][:, :, None])[:, :, :2]

    return out_coordinates
