#!/usr/bin/env python
#
# Henry Gomersall
# heng@kedevelopments.co.uk
#
# 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 holder nor the names of its 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 HOLDER 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.
#

# This code has been adapted to add support for the out argument in rfftn, irfftn
# to allow for reusing existing array buffers
# Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>

import threading

import pyfftw
import numpy as np
import pyfftw.builders as builders
from pyfftw.interfaces import cache
from pyfftw.builders._utils import _norm_args, _default_effort, _default_threads


def _Xfftn(
    a,
    s,
    axes,
    overwrite_input,
    planner_effort,
    threads,
    auto_align_input,
    auto_contiguous,
    calling_func,
    normalise_idft=True,
    ortho=False,
    real_direction_flag=None,
    output_array=None,
):

    work_with_copy = False

    a = np.asanyarray(a)

    try:
        s = tuple(s)
    except TypeError:
        pass

    try:
        axes = tuple(axes)
    except TypeError:
        pass

    if calling_func in ("dct", "dst"):
        # real-to-real transforms require passing an additional flag argument
        avoid_copy = False
        args = (
            overwrite_input,
            planner_effort,
            threads,
            auto_align_input,
            auto_contiguous,
            avoid_copy,
            real_direction_flag,
        )
    elif calling_func in ("irfft2", "irfftn"):
        # overwrite_input is not an argument to irfft2 or irfftn
        args = (planner_effort, threads, auto_align_input, auto_contiguous)

        if not overwrite_input:
            # Only irfft2 and irfftn have overwriting the input
            # as the default (and so require the input array to
            # be reloaded).
            work_with_copy = True
    else:
        args = (
            overwrite_input,
            planner_effort,
            threads,
            auto_align_input,
            auto_contiguous,
        )

        if not a.flags.writeable:
            # Special case of a locked array - always work with a
            # copy.  See issue #92.
            work_with_copy = True

            if overwrite_input:
                raise ValueError(
                    "overwrite_input cannot be True when the "
                    + "input array flags.writeable is False"
                )

    if work_with_copy:
        # We make the copy before registering the key so that the
        # copy's stride information will be cached since this will be
        # used for planning.  Make sure the copy is byte aligned to
        # prevent further copying
        a_original = a
        a = pyfftw.empty_aligned(shape=a.shape, dtype=a.dtype)
        a[...] = a_original

    if cache.is_enabled():
        alignment = a.ctypes.data % pyfftw.simd_alignment

        key = (
            calling_func,
            a.shape,
            a.strides,
            a.dtype,
            s.__hash__(),
            axes.__hash__(),
            alignment,
            args,
            threading.get_ident(),
        )

        try:
            if key in cache._fftw_cache:
                FFTW_object = cache._fftw_cache.lookup(key)
            else:
                FFTW_object = None

        except KeyError:
            # This occurs if the object has fallen out of the cache between
            # the check and the lookup
            FFTW_object = None

    if not cache.is_enabled() or FFTW_object is None:

        # If we're going to create a new FFTW object and are not
        # working with a copy, then we need to copy the input array to
        # preserve it, otherwise we can't actually  take the transform
        # of the input array! (in general, we have to assume that the
        # input array will be destroyed during planning).
        if not work_with_copy:
            a_copy = a.copy()

        planner_args = (a, s, axes) + args

        FFTW_object = getattr(builders, calling_func)(*planner_args)

        # Only copy if the input array is what was actually used
        # (otherwise it shouldn't be overwritten)
        if not work_with_copy and FFTW_object.input_array is a:
            a[:] = a_copy

        if cache.is_enabled():
            cache._fftw_cache.insert(FFTW_object, key)

        output_array = FFTW_object(normalise_idft=normalise_idft, ortho=ortho)

    else:
        orig_output_array = FFTW_object.output_array
        output_shape = orig_output_array.shape
        output_dtype = orig_output_array.dtype
        output_alignment = FFTW_object.output_alignment

        if output_array is None:
            output_array = pyfftw.empty_aligned(
                output_shape, output_dtype, n=output_alignment
            )

        FFTW_object(
            input_array=a,
            output_array=output_array,
            normalise_idft=normalise_idft,
            ortho=ortho,
        )

    return output_array


def rfftn(
    a,
    s=None,
    axes=None,
    norm=None,
    overwrite_input=False,
    planner_effort=None,
    threads=None,
    auto_align_input=True,
    auto_contiguous=True,
    out=None,
):
    """Perform an n-D real FFT.

    The first four arguments are as per :func:`numpy.fft.rfftn`;
    the rest of the arguments are documented
    in the :ref:`additional arguments docs<interfaces_additional_args>`.
    """
    calling_func = "rfftn"
    planner_effort = _default_effort(planner_effort)
    threads = _default_threads(threads)

    return _Xfftn(
        a,
        s,
        axes,
        overwrite_input,
        planner_effort,
        threads,
        auto_align_input,
        auto_contiguous,
        calling_func,
        **_norm_args(norm),
        output_array=out,
    )


def irfftn(
    a,
    s=None,
    axes=None,
    norm=None,
    overwrite_input=False,
    planner_effort=None,
    threads=None,
    auto_align_input=True,
    auto_contiguous=True,
    out=None,
):
    """Perform an n-D real inverse FFT.

    The first four arguments are as per :func:`numpy.fft.rfftn`;
    the rest of the arguments are documented
    in the :ref:`additional arguments docs<interfaces_additional_args>`.
    """
    calling_func = "irfftn"
    planner_effort = _default_effort(planner_effort)
    threads = _default_threads(threads)

    return _Xfftn(
        a,
        s,
        axes,
        overwrite_input,
        planner_effort,
        threads,
        auto_align_input,
        auto_contiguous,
        calling_func,
        **_norm_args(norm),
        output_array=out,
    )
