/*******************************************************
 * 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/data.h>

#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wparentheses"
#include <half.hpp>
#pragma GCC diagnostic pop

#include <af/arith.h>
#include <af/array.h>
#include <af/complex.h>
#include <af/defines.h>
#include <af/gfor.h>
#include <af/half.h>
#include <af/traits.hpp>
#include "error.hpp"

#include <type_traits>

using af::array;
using af::dim4;
using af::dtype;
using std::enable_if;

namespace {
// NOTE: we are repeating this here so that we don't need to access the
// is_complex types in backend/common. This is done to isolate the C++ API from
// the internal API
template<typename T>
struct is_complex {
    static const bool value = false;
};
template<>
struct is_complex<af::cfloat> {
    static const bool value = true;
};
template<>
struct is_complex<af::cdouble> {
    static const bool value = true;
};

array constant(af_half val, const dim4 &dims, const dtype type) {
    af_array res;
    UNUSED(val);
    AF_THROW(af_constant(&res, 0,  //(double)val,
                         dims.ndims(), dims.get(), type));
    return array(res);
}

template<typename T, typename = typename enable_if<
                         !static_cast<bool>(is_complex<T>::value), T>::type>
array constant(T val, const dim4 &dims, dtype type) {
    af_array res;
    if (type != s64 && type != u64) {
        AF_THROW(
            af_constant(&res, (double)val, dims.ndims(), dims.get(), type));
    } else if (type == s64) {
        AF_THROW(
            af_constant_long(&res, (long long)val, dims.ndims(), dims.get()));
    } else {
        AF_THROW(af_constant_ulong(&res, (unsigned long long)val, dims.ndims(),
                                   dims.get()));
    }
    return array(res);
}

template<typename T>
typename enable_if<static_cast<bool>(is_complex<T>::value), array>::type
constant(T val, const dim4 &dims, const dtype type) {
    if (type != c32 && type != c64) {
        return ::constant(real(val), dims, type);
    }
    af_array res;
    AF_THROW(af_constant_complex(&res, real(val), imag(val), dims.ndims(),
                                 dims.get(), type));
    return array(res);
}
}  // namespace

namespace af {
template<typename T>
array constant(T val, const dim4 &dims, const af::dtype type) {
    return ::constant(val, dims, type);
}

template<typename T>
array constant(T val, const dim_t d0, const af::dtype ty) {
    return ::constant(val, dim4(d0), ty);
}

template<typename T>
array constant(T val, const dim_t d0, const dim_t d1, const af::dtype ty) {
    return ::constant(val, dim4(d0, d1), ty);
}

template<typename T>
array constant(T val, const dim_t d0, const dim_t d1, const dim_t d2,
               const af::dtype ty) {
    return ::constant(val, dim4(d0, d1, d2), ty);
}

template<typename T>
array constant(T val, const dim_t d0, const dim_t d1, const dim_t d2,
               const dim_t d3, const af::dtype ty) {
    return ::constant(val, dim4(d0, d1, d2, d3), ty);
}

#define CONSTANT(TYPE)                                                       \
    template AFAPI array constant<TYPE>(TYPE val, const dim4 &dims,          \
                                        const af::dtype ty);                 \
    template AFAPI array constant<TYPE>(TYPE val, const dim_t d0,            \
                                        const af::dtype ty);                 \
    template AFAPI array constant<TYPE>(TYPE val, const dim_t d0,            \
                                        const dim_t d1, const af::dtype ty); \
    template AFAPI array constant<TYPE>(TYPE val, const dim_t d0,            \
                                        const dim_t d1, const dim_t d2,      \
                                        const af::dtype ty);                 \
    template AFAPI array constant<TYPE>(TYPE val, const dim_t d0,            \
                                        const dim_t d1, const dim_t d2,      \
                                        const dim_t d3, const af::dtype ty);
CONSTANT(double);
CONSTANT(float);
CONSTANT(int);
CONSTANT(unsigned);
CONSTANT(char);
CONSTANT(signed char);
CONSTANT(unsigned char);
CONSTANT(cfloat);
CONSTANT(cdouble);
CONSTANT(long);
CONSTANT(unsigned long);
CONSTANT(long long);
CONSTANT(unsigned long long);
CONSTANT(bool);
CONSTANT(short);
CONSTANT(unsigned short);
CONSTANT(half);
CONSTANT(half_float::half);

#undef CONSTANT

array range(const dim4 &dims, const int seq_dim, const af::dtype ty) {
    af_array out;
    AF_THROW(af_range(&out, dims.ndims(), dims.get(), seq_dim, ty));
    return array(out);
}

array range(const dim_t d0, const dim_t d1, const dim_t d2, const dim_t d3,
            const int seq_dim, const af::dtype ty) {
    return range(dim4(d0, d1, d2, d3), seq_dim, ty);
}

array iota(const dim4 &dims, const dim4 &tile_dims, const af::dtype ty) {
    af_array out;
    AF_THROW(af_iota(&out, dims.ndims(), dims.get(), tile_dims.ndims(),
                     tile_dims.get(), ty));
    return array(out);
}

array identity(const dim4 &dims, const af::dtype type) {
    af_array res;
    AF_THROW(af_identity(&res, dims.ndims(), dims.get(), type));
    return array(res);
}

array identity(const dim_t d0, const af::dtype ty) {
    return identity(dim4(d0), ty);
}

array identity(const dim_t d0, const dim_t d1, const af::dtype ty) {
    return identity(dim4(d0, d1), ty);
}

array identity(const dim_t d0, const dim_t d1, const dim_t d2,
               const af::dtype ty) {
    return identity(dim4(d0, d1, d2), ty);
}

array identity(const dim_t d0, const dim_t d1, const dim_t d2, const dim_t d3,
               const af::dtype ty) {
    return identity(dim4(d0, d1, d2, d3), ty);
}

array diag(const array &in, const int num, const bool extract) {
    af_array res;
    if (extract) {
        AF_THROW(af_diag_extract(&res, in.get(), num));
    } else {
        AF_THROW(af_diag_create(&res, in.get(), num));
    }

    return array(res);
}

array moddims(const array &in, const unsigned ndims, const dim_t *const dims) {
    af_array out = 0;
    AF_THROW(af_moddims(&out, in.get(), ndims, dims));
    return array(out);
}

array moddims(const array &in, const dim4 &dims) {
    return af::moddims(in, dims.ndims(), dims.get());
}

array moddims(const array &in, const dim_t d0, const dim_t d1, const dim_t d2,
              const dim_t d3) {
    dim_t dims[4] = {d0, d1, d2, d3};
    return af::moddims(in, 4, dims);
}

array flat(const array &in) {
    af_array out = 0;
    AF_THROW(af_flat(&out, in.get()));
    return array(out);
}

array join(const int dim, const array &first, const array &second) {
    af_array out = 0;
    AF_THROW(af_join(&out, dim, first.get(), second.get()));
    return array(out);
}

array join(const int dim, const array &first, const array &second,
           const array &third) {
    af_array out       = 0;
    af_array inputs[3] = {first.get(), second.get(), third.get()};
    AF_THROW(af_join_many(&out, dim, 3, inputs));
    return array(out);
}

array join(const int dim, const array &first, const array &second,
           const array &third, const array &fourth) {
    af_array out       = 0;
    af_array inputs[4] = {first.get(), second.get(), third.get(), fourth.get()};
    AF_THROW(af_join_many(&out, dim, 4, inputs));
    return array(out);
}

array tile(const array &in, const unsigned x, const unsigned y,
           const unsigned z, const unsigned w) {
    af_array out = 0;
    AF_THROW(af_tile(&out, in.get(), x, y, z, w));
    return array(out);
}

array tile(const array &in, const af::dim4 &dims) {
    af_array out = 0;
    AF_THROW(af_tile(&out, in.get(), dims[0], dims[1], dims[2], dims[3]));
    return array(out);
}

array reorder(const array &in, const unsigned x, const unsigned y,
              const unsigned z, const unsigned w) {
    af_array out = 0;
    AF_THROW(af_reorder(&out, in.get(), x, y, z, w));
    return array(out);
}

array shift(const array &in, const int x, const int y, const int z,
            const int w) {
    af_array out = 0;
    AF_THROW(af_shift(&out, in.get(), x, y, z, w));
    return array(out);
}

array flip(const array &in, const unsigned dim) {
    af_array out = 0;
    AF_THROW(af_flip(&out, in.get(), dim));
    return array(out);
}

array lower(const array &in, bool is_unit_diag) {
    af_array res;
    AF_THROW(af_lower(&res, in.get(), is_unit_diag));
    return array(res);
}

array upper(const array &in, bool is_unit_diag) {
    af_array res;
    AF_THROW(af_upper(&res, in.get(), is_unit_diag));
    return array(res);
}

array select(const array &cond, const array &a, const array &b) {
    af_array res;
    AF_THROW(af_select(&res, cond.get(), a.get(), b.get()));
    return array(res);
}

array select(const array &cond, const array &a, const double &b) {
    af_array res;
    AF_THROW(af_select_scalar_r(&res, cond.get(), a.get(), b));
    return array(res);
}

array select(const array &cond, const double &a, const array &b) {
    af_array res;
    AF_THROW(af_select_scalar_l(&res, cond.get(), a, b.get()));
    return array(res);
}

void replace(array &a, const array &cond, const array &b) {
    AF_THROW(af_replace(a.get(), cond.get(), b.get()));
}

void replace(array &a, const array &cond, const double &b) {
    AF_THROW(af_replace_scalar(a.get(), cond.get(), b));
}

void replace(array &a, const array &cond, const long long b) {
    AF_THROW(af_replace_scalar_long(a.get(), cond.get(), b));
}

void replace(array &a, const array &cond, const unsigned long long b) {
    AF_THROW(af_replace_scalar_ulong(a.get(), cond.get(), b));
}

array select(const array &cond, const array &a, const long long b) {
    af_array res;
    AF_THROW(af_select_scalar_r_long(&res, cond.get(), a.get(), b));
    return array(res);
}

array select(const array &cond, const array &a, const unsigned long long b) {
    af_array res;
    AF_THROW(af_select_scalar_r_ulong(&res, cond.get(), a.get(), b));
    return array(res);
}

array select(const array &cond, const long long a, const array &b) {
    af_array res;
    AF_THROW(af_select_scalar_l_long(&res, cond.get(), a, b.get()));
    return array(res);
}

array select(const array &cond, const unsigned long long a, const array &b) {
    af_array res;
    AF_THROW(af_select_scalar_l_ulong(&res, cond.get(), a, b.get()));
    return array(res);
}

array pad(const array &in, const dim4 &beginPadding, const dim4 &endPadding,
          const borderType padFillType) {
    af_array out = 0;
    // FIXME(pradeep) Cannot use dim4::ndims() since that will
    //               always return 0 if any one of dimensions
    //               has no padding completely
    AF_THROW(af_pad(&out, in.get(), 4, beginPadding.get(), 4, endPadding.get(),
                    padFillType));
    return array(out);
}

}  // namespace af