/*******************************************************

* Copyright (c) 2015, 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 <backend.hpp>

#include <common/ArrayInfo.hpp>

#include <common/err_common.hpp>

#include <common/half.hpp>

#include <handle.hpp>

#include <implicit.hpp>

#include <optypes.hpp>

#include <select.hpp>

#include <af/arith.h>

#include <af/array.h>

#include <af/data.h>

#include <af/defines.h>

using af::dim4;

using common::half;

using detail::Array;

using detail::cdouble;

using detail::cfloat;

using detail::createSelectNode;

using detail::intl;

using detail::uchar;

using detail::uint;

using detail::uintl;

using detail::ushort;

template<typename T>

af_array select(const af_array cond, const af_array a, const af_array b,

const dim4& odims) {

Array<T> out = createSelectNode(getArray<char>(cond), getArray<T>(a),

getArray<T>(b), odims);

return getHandle<T>(out);

}

af_err af_select(af_array* out, const af_array cond, const af_array a,

const af_array b) {

try {

const ArrayInfo& ainfo = getInfo(a);

const ArrayInfo& binfo = getInfo(b);

const ArrayInfo& cond_info = getInfo(cond);

if (cond_info.ndims() == 0) { return af_retain_array(out, cond); }

ARG_ASSERT(2, ainfo.getType() == binfo.getType());

ARG_ASSERT(1, cond_info.getType() == b8);

dim4 adims = ainfo.dims();

dim4 bdims = binfo.dims();

dim4 cond_dims = cond_info.dims();

dim4 odims(1, 1, 1, 1);

for (int i = 0; i < 4; i++) {

DIM_ASSERT(2, (adims[i] == bdims[i] && adims[i] == cond_dims[i]) ||

adims[i] == 1 || bdims[i] == 1 ||

cond_dims[i] == 1);

odims[i] = std::max(std::max(adims[i], bdims[i]), cond_dims[i]);

}

af_array res;

switch (ainfo.getType()) {

case f32: res = select<float>(cond, a, b, odims); break;

case f64: res = select<double>(cond, a, b, odims); break;

case c32: res = select<cfloat>(cond, a, b, odims); break;

case c64: res = select<cdouble>(cond, a, b, odims); break;

case s32: res = select<int>(cond, a, b, odims); break;

case u32: res = select<uint>(cond, a, b, odims); break;

case s64: res = select<intl>(cond, a, b, odims); break;

case u64: res = select<uintl>(cond, a, b, odims); break;

case s16: res = select<short>(cond, a, b, odims); break;

case u16: res = select<ushort>(cond, a, b, odims); break;

case u8: res = select<uchar>(cond, a, b, odims); break;

case b8: res = select<char>(cond, a, b, odims); break;

case f16: res = select<half>(cond, a, b, odims); break;

default: TYPE_ERROR(2, ainfo.getType());

}

std::swap(*out, res);

}

CATCHALL;

return AF_SUCCESS;

}

template<typename T, bool flip>

af_array select_scalar(const af_array cond, const af_array a, const double b,

const dim4& odims) {

Array<T> out = createSelectNode<T, flip>(getArray<char>(cond),

getArray<T>(a), b, odims);

return getHandle<T>(out);

}

af_err af_select_scalar_r(af_array* out, const af_array cond, const af_array a,

const double b) {

try {

const ArrayInfo& ainfo = getInfo(a);

const ArrayInfo& cinfo = getInfo(cond);

ARG_ASSERT(1, cinfo.getType() == b8);

dim4 adims = ainfo.dims();

dim4 cond_dims = cinfo.dims();

dim4 odims(1);

for (int i = 0; i < 4; i++) {

DIM_ASSERT(1, cond_dims[i] == adims[i] || cond_dims[i] == 1 ||

adims[i] == 1);

odims[i] = std::max(cond_dims[i], adims[i]);

}

af_array res;

switch (ainfo.getType()) {

case f16:

res = select_scalar<half, false>(cond, a, b, odims);

break;

case f32:

res = select_scalar<float, false>(cond, a, b, odims);

break;

case f64:

res = select_scalar<double, false>(cond, a, b, odims);

break;

case c32:

res = select_scalar<cfloat, false>(cond, a, b, odims);

break;

case c64:

res = select_scalar<cdouble, false>(cond, a, b, odims);

break;

case s32: res = select_scalar<int, false>(cond, a, b, odims); break;

case u32:

res = select_scalar<uint, false>(cond, a, b, odims);

break;

case s16:

res = select_scalar<short, false>(cond, a, b, odims);

break;

case u16:

res = select_scalar<ushort, false>(cond, a, b, odims);

break;

case s64:

res = select_scalar<intl, false>(cond, a, b, odims);

break;

case u64:

res = select_scalar<uintl, false>(cond, a, b, odims);

break;

case u8:

res = select_scalar<uchar, false>(cond, a, b, odims);

break;

case b8: res = select_scalar<char, false>(cond, a, b, odims); break;

default: TYPE_ERROR(2, ainfo.getType());

}

std::swap(*out, res);

}

CATCHALL;

return AF_SUCCESS;

}

af_err af_select_scalar_l(af_array* out, const af_array cond, const double a,

const af_array b) {

try {

const ArrayInfo& binfo = getInfo(b);

const ArrayInfo& cinfo = getInfo(cond);

ARG_ASSERT(1, cinfo.getType() == b8);

dim4 bdims = binfo.dims();

dim4 cond_dims = cinfo.dims();

dim4 odims(1);

for (int i = 0; i < 4; i++) {

DIM_ASSERT(1, cond_dims[i] == bdims[i] || cond_dims[i] == 1 ||

bdims[i] == 1);

odims[i] = std::max(cond_dims[i], bdims[i]);

}

af_array res;

switch (binfo.getType()) {

case f16: res = select_scalar<half, true>(cond, b, a, odims); break;

case f32:

res = select_scalar<float, true>(cond, b, a, odims);

break;

case f64:

res = select_scalar<double, true>(cond, b, a, odims);

break;

case c32:

res = select_scalar<cfloat, true>(cond, b, a, odims);

break;

case c64:

res = select_scalar<cdouble, true>(cond, b, a, odims);

break;

case s32: res = select_scalar<int, true>(cond, b, a, odims); break;

case u32: res = select_scalar<uint, true>(cond, b, a, odims); break;

case s16:

res = select_scalar<short, true>(cond, b, a, odims);

break;

case u16:

res = select_scalar<ushort, true>(cond, b, a, odims);

break;

case s64: res = select_scalar<intl, true>(cond, b, a, odims); break;

case u64:

res = select_scalar<uintl, true>(cond, b, a, odims);

break;

case u8: res = select_scalar<uchar, true>(cond, b, a, odims); break;

case b8: res = select_scalar<char, true>(cond, b, a, odims); break;

default: TYPE_ERROR(2, binfo.getType());

}

std::swap(*out, res);

}

CATCHALL;

return AF_SUCCESS;

}