#!/usr/bin/env python
# encoding: utf-8
"""
@file fft.py
@brief Utilities for Fourier transforms

@package libtim.fft
@brief Utilities for Fourier transforms
@author Tim van Werkhoven (werkhoven@strw.leidenuniv.nl)
@copyright Creative Commons Attribution-Share Alike license versions 3.0 or higher, see http://creativecommons.org/licenses/by-sa/3.0/
@date 20120403

Package for some utilities for Fourier transforms

Modified by Frank Horowitz from the package libtim found at https://github.com/tvwerkhoven/libtim-py.git
"""
"""from __future__ import division
from future.builtins import zip"""

# ============================================================================
# Import libraries here
# ============================================================================

import numpy as np
from collections.abc import Iterable
from math import ceil, floor

# ============================================================================
# Defines
# ============================================================================

# ============================================================================
# Routines
# ============================================================================


def mk_apod_mask(
    masksz,
    apodpos=None,
    apodsz=None,
    shape="rect",
    wsize=-0.3,
    apod_f=lambda x: 0.5 * (1.0 - np.cos(np.pi * x)),
):
    """
    Generate apodisation mask with custom size, shape, edge.

    The output array mask will be **masksz**, while the actual apodisation masked will **apodsz** big. The position of the mask is given with **apodpos**.

    **apodpos** defaults to the center, **apodsz** defaults to **masksz**

    **apodpos**, **apodsz** and **wsize** can either be given as fraction (if < 0) or as absolute number of pixels (if > 0). If these are given in int or float, the result will be square, if these are tuples, the size can be different in both dimensions.

    If **apodpos** or **apodsz** are fractional, they are relative to **masksz**. Fractional **wsize** is relative to **apodsz**.

    **apod_f** is the windowing function used. It can be a string (see list below), or a lambda function. In the latter case it should take one float coordinate between 1 and 0 as input and return the value of the window at that position.

    Some apodisation functions (for **apod_f**):
    - 'Hann': lambda x: 0.5 * (1.0 - np.cos(np.pi*x))
    - 'Hamming': lambda x: 0.54 - 0.46 *np.cos(np.pi*x)
    - '(Co)sine' window: lambda x: np.sin(np.pi*x*0.5)
    - 'Lanczos': lambda x: np.sinc(x-1.0)

    @param [in] masksz Size of the output array containing the apodisation mask
    @param [in] apodpos Position of the apodisation mask
    @param [in] apodsz Size of the apodisation mask
    @param [in] shape Apodisation mask shape, 'rect' or 'circular'
    @param [in] wsize Size of the apodisation window, i.e. the transition region to go from 0 to 1.
    @param [in] apod_f Apodisation function to use. Can be lambda function
    """

    # Check apodpos and apodsz, if not set, use defaults
    if apodpos == None:
        apodpos = tuple((np.r_[masksz] - 1.0) / 2.0)
    if apodsz == None:
        apodsz = masksz

    apod_func = lambda x: x
    if hasattr(apod_f, "__call__"):
        apod_func = apod_f
    else:
        try:
            apod_f = apod_f.lower()
            if apod_f[:4] == "hann":
                apod_func = lambda x: 0.5 * (1.0 - np.cos(np.pi * x))
            elif apod_f[:4] == "hamm":
                apod_func = lambda x: 0.54 - 0.46 * np.cos(np.pi * x)
            elif apod_f[:3] == "cos" or apod_f[:3] == "sin":
                apod_func = lambda x: np.sin(np.pi * x * 0.5)
            elif apod_f[:4] == "lanc":
                apod_func = lambda x: np.sinc(x - 1.0)
            else:
                raise ValueError("<apod_f> not supported!")
        except ValueError:
            raise
        except:
            raise ValueError("<apod_f> should be a string or callable!")

    # Mask size <masksz> should be iterable (like a list or tuple)
    if not isinstance(masksz, Iterable):
        raise TypeError("<masksz> should be iterable")
    if min(masksz) < 1:
        raise ValueError("All mask size <masksz> dimensions should be >= 1")

    # Only the first 4 letters are significant.
    try:
        shape = shape[:4]
    except:
        raise ValueError("<shape> should be a string!")

    # Check if shape is legal
    if shape not in ("rect", "circ"):
        raise ValueError("<shape> should be 'rectangle' or 'circle'")

    # Check if apodpos, apodsz and wsize are legal. They should either be a
    # scalar (i.e. non-iterable) or the same length as <masksz> (which is iterable). Also, if apodpos, apodsz or wsize are just one int or float, repeat them for each dimension.
    if isinstance(apodpos, Iterable) and len(apodpos) != len(masksz):
        raise TypeError(
            "<apodpos> should be either 1 element per dimension or 1 in total."
        )
    elif not isinstance(apodpos, Iterable):
        apodpos = (apodpos,) * len(masksz)

    if isinstance(apodsz, Iterable) and len(apodsz) != len(masksz):
        raise TypeError(
            "<apodsz> should be either 1 element per dimension or 1 in total."
        )
    elif not isinstance(apodsz, Iterable):
        apodsz = (apodsz,) * len(masksz)

    if isinstance(wsize, Iterable) and len(wsize) != len(masksz):
        raise TypeError(
            "<wsize> should be either 1 element per dimension or 1 in total."
        )
    elif not isinstance(wsize, Iterable):
        wsize = (wsize,) * len(masksz)

    # If apodsz or wsize are fractional, calculate the absolute size.
    if min(apodpos) < 0:
        apodpos *= -np.r_[masksz]
    if min(apodsz) < 0:
        apodsz *= -np.r_[masksz]
    if min(wsize) < 0:
        wsize *= -np.r_[apodsz]

    # Generate base mask, which are (x,y) coordinates around the center
    mask = np.indices(masksz, dtype=float)

    # Center the mask around <apodpos> for any number of dimensions
    for maski, posi in zip(mask, apodpos):
        maski -= posi

    # If the mask shape is circular, calculate the radial distance from
    # <apodpos>
    if shape == "circ":
        mask = np.array([np.sum(mask**2.0, 0) ** 0.5])

    # Scale the pixels such that there is only a band going from 1 to 0 between <masksz>-<wsize> and <masksz>
    for maski, szi, wszi in zip(mask, apodsz, wsize):
        # First take the negative absolute value of the mask, such that 0 is at the origin and the value goes down outward from where the mask should be.
        maski[:] = -np.abs(maski)
        # Next, add the radius of the apodisation mask size to the values, such that the outside edge of the requested mask is exactly zero.
        # TODO Should this be (szi-1)/2 or (szi)/2?
        maski += (szi) / 2.0
        # Now divide the mask by the windowing area inside the apod. mask, such that the inner edge of the mask is 1.
        if wszi != 0:
            maski /= wszi / 2.0
        else:
            maski /= 0.001
        # Store masks for inside and outside the mask area
        inmask = maski > 1
        outmask = maski <= 0
        # Apply function to all data
        maski[:] = apod_func(maski[:])
        # Now everything higher than 1 is inside the mask, and smaller than 0 is outside the mask. Clip these values to (0,1)
        maski[inmask] = 1
        maski[outmask] = 0

    # Apply apodisation function to all elements, and multiply
    if shape == "rect":
        return np.prod(mask, 0)
    elif shape == "circ":
        return mask[0]


def descramble(data, direction=1):
    """
    (de)scramble **data**, usually used for Fourier transform.

    'Scrambling' data means to swap around quadrant 1 with 3 and 2 with 4 in
    a data matrix. The effect is that the zero frequency is no longer at
    **data[0,0]** but in the middle of the matrix

    @param [in] data Data to (de)scramble
    @param [in] direction 1: scramble, -1: descramble
    @return (de)scrambled data
    """

    if direction == 1:
        return np.fft.fftshift(data)
    else:
        return np.fft.ifftshift(data)


def embed_data(indata, big_shape=(None, None), pad_size=None):
    """
    Embed **indata** in a zero-filled rectangular array of shape big_shape.

    For apodization purposes, a padding of size pad_size will surround the embedded image.

    To prevent wrapping artifacts in Fourier analysis, this function can  embed data in a zero-filled rectangular array of a larger size.

    @param [in] indata Data to embed
    @param [in] big_shape Size of embedding matrix
    @param [in] pad_size width of apodization pad, which will be filled with indata edge values
    @return embedded data, and x/y-direction slicing objects
    """
    if big_shape == (None, None):
        raise ValueError("Need to specify size of returned array!")
    if big_shape[0] < indata.shape[0]:
        raise ValueError("Returned Y direction must be at least as big as indata's.")
    if big_shape[1] < indata.shape[1]:
        raise ValueError("Returned X direction must be at least as big as indata's.")

    s = big_shape

    # Generate empty array
    retdat = np.zeros(s, dtype=indata.dtype)

    inmean = np.mean(indata)

    # These slices denote the central region where <indata> will go
    x_pad_2 = s[1] - indata.shape[1]
    x_pad = floor(x_pad_2 / 2.0)
    slice_x = slice(x_pad, x_pad + indata.shape[1], 1)
    if pad_size != None:
        pad_slice_x = slice(slice_x.start - pad_size, slice_x.stop + pad_size)

    y_pad_2 = s[0] - indata.shape[0]
    y_pad = floor(y_pad_2 / 2.0)
    slice_y = slice(y_pad, y_pad + indata.shape[0], 1)
    if pad_size != None:
        pad_slice_y = slice(slice_y.start - pad_size, slice_y.stop + pad_size)

    #
    # Insert the data and return it
    if pad_size == None:
        retdat[slice_y, slice_x] = indata - inmean
    else:
        retdat[pad_slice_y, pad_slice_x] = np.pad(
            indata - inmean, pad_width=pad_size, mode="edge"
        )
    return retdat, slice_y, slice_x


def embed_data_old(indata, direction=1, scale=2):
    """
    Embed **indata** in a zero-filled rectangular array **scale** times as big as **indata**.

    To prevent wrapping artifacts in Fourier analysis, this function can  embed data in a zero-filled rectangular array of twice the size.

    If **direction** = 1, **indata** will be embedded, if **direction** = -1, it will be dis-embedded.

    @param [in] indata Data to embed
    @param [in] direction 1: embed, -1: dis-embed
    @param [in] scale Size of embedding matrix wrt **indata**
    @return (dis)-embedded data, either scale*indata.shape or 1/scale*indata.shape
    """

    s = np.r_[indata.shape]

    if direction == 1:
        # Generate empty array
        retdat = np.zeros(np.r_[s] * scale, dtype=indata.dtype)
        # These slices denote the central region where <indata> will go
        slice0 = slice(
            retdat.shape[0] / 2 - floor(s[0] / 2.0),
            retdat.shape[0] / 2 + ceil(s[0] / 2.0),
        )
        slice1 = slice(
            retdat.shape[1] / 2 - floor(s[1] / 2.0),
            retdat.shape[1] / 2 + ceil(s[1] / 2.0),
        )

        #
        # Insert the data and return it
        retdat[slice0, slice1] = indata
        return retdat, slice0, slice1
    else:
        # These slices give the central area of the data
        slice0 = slice(
            s[0] / 2 - floor(s[0] / (2.0 * scale)),
            s[0] / 2 + ceil(s[0] / (2.0 * scale)),
        )
        slice1 = slice(
            s[1] / 2 - floor(s[1] / (2.0 * scale)),
            s[1] / 2 + ceil(s[1] / (2.0 * scale)),
        )

        # Slice out the center and return it
        return indata[slice0, slice1]
