/*******************************************************
 * Copyright (c) 2014, ArrayFire
 * All rights reserved.
 *
 * This file is distributed under 3-clause BSD license.
 * The complete license agreement can be obtained at:
 * http://arrayfire.com/licenses/BSD-3-Clause
 ********************************************************/

#include <af/array.h>
#include <af/dim4.hpp>
#include <af/signal.h>
#include "error.hpp"

using af::array;
using af::dim4;

namespace af {
array fftNorm(const array& in, const double norm_factor, const dim_t odim0) {
    af_array out = 0;
    AF_THROW(af_fft(&out, in.get(), norm_factor, odim0));
    return array(out);
}

array fft2Norm(const array& in, const double norm_factor, const dim_t odim0,
               const dim_t odim1) {
    af_array out = 0;
    AF_THROW(af_fft2(&out, in.get(), norm_factor, odim0, odim1));
    return array(out);
}

array fft3Norm(const array& in, const double norm_factor, const dim_t odim0,
               const dim_t odim1, const dim_t odim2) {
    af_array out = 0;
    AF_THROW(af_fft3(&out, in.get(), norm_factor, odim0, odim1, odim2));
    return array(out);
}

array fft(const array& in, const dim_t odim0) {
    return fftNorm(in, 1.0, odim0);
}

array fft2(const array& in, const dim_t odim0, const dim_t odim1) {
    return fft2Norm(in, 1.0, odim0, odim1);
}

array fft3(const array& in, const dim_t odim0, const dim_t odim1,
           const dim_t odim2) {
    return fft3Norm(in, 1.0, odim0, odim1, odim2);
}

// NOLINTNEXTLINE(performance-unnecessary-value-param)
array dft(const array& in, const double norm_factor, const dim4 outDims) {
    array temp;
    switch (in.dims().ndims()) {
        case 1: temp = fftNorm(in, norm_factor, outDims[0]); break;
        case 2: temp = fft2Norm(in, norm_factor, outDims[0], outDims[1]); break;
        case 3:
            temp =
                fft3Norm(in, norm_factor, outDims[0], outDims[1], outDims[2]);
            break;
        default: AF_THROW(AF_ERR_NOT_SUPPORTED);
    }
    return temp;
}

// NOLINTNEXTLINE(performance-unnecessary-value-param)
array dft(const array& in, const dim4 outDims) { return dft(in, 1.0, outDims); }

array dft(const array& in) { return dft(in, 1.0, dim4(0, 0, 0, 0)); }

array ifftNorm(const array& in, const double norm_factor, const dim_t odim0) {
    af_array out = 0;
    AF_THROW(af_ifft(&out, in.get(), norm_factor, odim0));
    return array(out);
}

array ifft2Norm(const array& in, const double norm_factor, const dim_t odim0,
                const dim_t odim1) {
    af_array out = 0;
    AF_THROW(af_ifft2(&out, in.get(), norm_factor, odim0, odim1));
    return array(out);
}

array ifft3Norm(const array& in, const double norm_factor, const dim_t odim0,
                const dim_t odim1, const dim_t odim2) {
    af_array out = 0;
    AF_THROW(af_ifft3(&out, in.get(), norm_factor, odim0, odim1, odim2));
    return array(out);
}

array ifft(const array& in, const dim_t odim0) {
    const dim4 dims    = in.dims();
    dim_t dim0         = odim0 == 0 ? dims[0] : odim0;
    double norm_factor = 1.0 / static_cast<double>(dim0);
    return ifftNorm(in, norm_factor, odim0);
}

array ifft2(const array& in, const dim_t odim0, const dim_t odim1) {
    const dim4 dims    = in.dims();
    dim_t dim0         = odim0 == 0 ? dims[0] : odim0;
    dim_t dim1         = odim1 == 0 ? dims[1] : odim1;
    double norm_factor = 1.0 / static_cast<double>(dim0 * dim1);
    return ifft2Norm(in, norm_factor, odim0, odim1);
}

array ifft3(const array& in, const dim_t odim0, const dim_t odim1,
            const dim_t odim2) {
    const dim4 dims    = in.dims();
    dim_t dim0         = odim0 == 0 ? dims[0] : odim0;
    dim_t dim1         = odim1 == 0 ? dims[1] : odim1;
    dim_t dim2         = odim2 == 0 ? dims[2] : odim2;
    double norm_factor = 1.0 / static_cast<double>(dim0 * dim1 * dim2);
    return ifft3Norm(in, norm_factor, odim0, odim1, odim2);
}

// NOLINTNEXTLINE(performance-unnecessary-value-param)
array idft(const array& in, const double norm_factor, const dim4 outDims) {
    array temp;
    switch (in.dims().ndims()) {
        case 1: temp = ifftNorm(in, norm_factor, outDims[0]); break;
        case 2:
            temp = ifft2Norm(in, norm_factor, outDims[0], outDims[1]);
            break;
        case 3:
            temp =
                ifft3Norm(in, norm_factor, outDims[0], outDims[1], outDims[2]);
            break;
        default: AF_THROW(AF_ERR_NOT_SUPPORTED);
    }
    return temp;
}

// NOLINTNEXTLINE(performance-unnecessary-value-param)
array idft(const array& in, const dim4 outDims) {
    return idft(in, 1.0, outDims);
}

array idft(const array& in) { return idft(in, 1.0, dim4(0, 0, 0, 0)); }

void fftInPlace(array& in, const double norm_factor) {
    AF_THROW(af_fft_inplace(in.get(), norm_factor));
}

void fft2InPlace(array& in, const double norm_factor) {
    AF_THROW(af_fft2_inplace(in.get(), norm_factor));
}

void fft3InPlace(array& in, const double norm_factor) {
    AF_THROW(af_fft3_inplace(in.get(), norm_factor));
}

void ifftInPlace(array& in, const double norm_factor) {
    const dim4 dims = in.dims();
    double norm     = norm_factor * (1.0 / static_cast<double>(dims[0]));
    AF_THROW(af_ifft_inplace(in.get(), norm));
}

void ifft2InPlace(array& in, const double norm_factor) {
    const dim4 dims = in.dims();
    double norm = norm_factor * (1.0 / static_cast<double>(dims[0] * dims[1]));
    AF_THROW(af_ifft2_inplace(in.get(), norm));
}

void ifft3InPlace(array& in, const double norm_factor) {
    const dim4 dims = in.dims();
    double norm =
        norm_factor * (1.0 / static_cast<double>(dims[0] * dims[1] * dims[2]));
    AF_THROW(af_ifft3_inplace(in.get(), norm));
}

template<>
AFAPI array fftR2C<1>(const array& in, const dim4& dims,
                      const double norm_factor) {
    af_array res;
    AF_THROW(af_fft_r2c(&res, in.get(), norm_factor == 0 ? 1.0 : norm_factor,
                        dims[0]));
    return array(res);
}

template<>
AFAPI array fftR2C<2>(const array& in, const dim4& dims,
                      const double norm_factor) {
    af_array res;
    AF_THROW(af_fft2_r2c(&res, in.get(), norm_factor == 0 ? 1.0 : norm_factor,
                         dims[0], dims[1]));
    return array(res);
}

template<>
AFAPI array fftR2C<3>(const array& in, const dim4& dims,
                      const double norm_factor) {
    af_array res;
    AF_THROW(af_fft3_r2c(&res, in.get(), norm_factor == 0 ? 1.0 : norm_factor,
                         dims[0], dims[1], dims[2]));
    return array(res);
}

inline dim_t getOrigDim(dim_t d, bool is_odd) {
    return 2 * (d - 1) + (is_odd ? 1 : 0);
}

template<>
AFAPI array fftC2R<1>(const array& in, const bool is_odd,
                      const double norm_factor) {
    double norm = norm_factor;

    if (norm == 0) {
        dim4 idims = in.dims();
        dim_t dim0 = getOrigDim(idims[0], is_odd);
        norm       = 1.0 / static_cast<double>(dim0);
    }

    af_array res;
    AF_THROW(af_fft_c2r(&res, in.get(), norm, is_odd));
    return array(res);
}

template<>
AFAPI array fftC2R<2>(const array& in, const bool is_odd,
                      const double norm_factor) {
    double norm = norm_factor;

    if (norm == 0) {
        dim4 idims = in.dims();
        dim_t dim0 = getOrigDim(idims[0], is_odd);
        dim_t dim1 = idims[1];
        norm       = 1.0 / static_cast<double>(dim0 * dim1);
    }

    af_array res;
    AF_THROW(af_fft2_c2r(&res, in.get(), norm, is_odd));
    return array(res);
}

template<>
AFAPI array fftC2R<3>(const array& in, const bool is_odd,
                      const double norm_factor) {
    double norm = norm_factor;

    if (norm == 0) {
        dim4 idims = in.dims();
        dim_t dim0 = getOrigDim(idims[0], is_odd);
        dim_t dim1 = idims[1];
        dim_t dim2 = idims[2];
        norm       = 1.0 / static_cast<double>(dim0 * dim1 * dim2);
    }

    af_array res;
    AF_THROW(af_fft3_c2r(&res, in.get(), norm, is_odd));
    return array(res);
}

#define FFT_REAL(rank)                                                    \
    template<>                                                            \
    AFAPI array fftR2C<rank>(const array& in, const double norm_factor) { \
        return fftR2C<rank>(in, in.dims(), norm_factor);                  \
    }

FFT_REAL(1)
FFT_REAL(2)
FFT_REAL(3)

void setFFTPlanCacheSize(size_t cacheSize) {
    AF_THROW(af_set_fft_plan_cache_size(cacheSize));
}
}  // namespace af