/*******************************************************
 * 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 <gtest/gtest.h>
#include <arrayfire.h>
#include <af/dim4.hpp>
#include <af/defines.h>
#include <af/traits.hpp>
#include <vector>
#include <iostream>
#include <complex>
#include <string>
#include <testHelpers.hpp>

using std::vector;
using std::string;
using std::cout;
using std::endl;
using af::cfloat;
using af::cdouble;

template<typename T>
class Iota : public ::testing::Test
{
    public:
        virtual void SetUp() {
            subMat0.push_back(af_make_seq(0, 4, 1));
            subMat0.push_back(af_make_seq(2, 6, 1));
            subMat0.push_back(af_make_seq(0, 2, 1));
        }
        vector<af_seq> subMat0;
};

// create a list of types to be tested
typedef ::testing::Types<float, double, int, unsigned int, intl, uintl, unsigned char, short, ushort> TestTypes;

// register the type list
TYPED_TEST_CASE(Iota, TestTypes);

template<typename T>
void iotaTest(const af::dim4 idims, const af::dim4 tdims)
{
    if (noDoubleTests<T>()) return;

    af_array outArray = 0;

    ASSERT_EQ(AF_SUCCESS, af_iota(&outArray, idims.ndims(), idims.get(),
               tdims.ndims(), tdims.get(), (af_dtype) af::dtype_traits<T>::af_type));

    af_array temp0 = 0, temp1 = 0, temp2 = 0;
    af::dim4 tempdims(idims.elements());
    af::dim4 fulldims;
    for(unsigned i = 0; i < 4; i++) {
        fulldims[i] = idims[i] * tdims[i];
    }
    ASSERT_EQ(AF_SUCCESS, af_range(&temp2, tempdims.ndims(), tempdims.get(), 0, (af_dtype) af::dtype_traits<T>::af_type));
    ASSERT_EQ(AF_SUCCESS, af_moddims(&temp1, temp2, idims.ndims(), idims.get()));
    ASSERT_EQ(AF_SUCCESS, af_tile(&temp0, temp1, tdims[0], tdims[1], tdims[2], tdims[3]));

    // Get result
    T* outData = new T[fulldims.elements()];
    ASSERT_EQ(AF_SUCCESS, af_get_data_ptr((void*)outData, outArray));

    T* tileData = new T[fulldims.elements()];
    ASSERT_EQ(AF_SUCCESS, af_get_data_ptr((void*)tileData, temp0));

    // Compare result
    for(int i = 0; i < (int) fulldims.elements(); i++)
        ASSERT_EQ(tileData[i], outData[i]) << "at: " << i << std::endl;

    // Delete
    delete[] outData;
    delete[] tileData;

    if(outArray  != 0) af_release_array(outArray);
    if(temp0     != 0) af_release_array(temp0);
    if(temp1     != 0) af_release_array(temp1);
    if(temp2     != 0) af_release_array(temp2);
}

#define IOTA_INIT(desc, x, y, z, w, a, b, c, d)                                             \
    TYPED_TEST(Iota, desc)                                                                  \
    {                                                                                       \
        iotaTest<TypeParam>(af::dim4(x, y, z, w), af::dim4(a, b, c, d));                    \
    }

    IOTA_INIT(Iota1D0, 100,  1, 1, 1, 2, 3, 1, 1);

    IOTA_INIT(Iota2D0,  10, 20, 1, 1, 3, 1, 2, 1);
    IOTA_INIT(Iota2D1, 100,  5, 1, 1, 1, 2, 4, 2);

    IOTA_INIT(Iota3D0,  20,  6, 3, 1, 1, 1, 1, 1);
    IOTA_INIT(Iota3D1,  10, 12, 5, 1, 2, 3, 4, 5);
    IOTA_INIT(Iota3D2,  25, 30, 2, 1, 1, 2, 2, 1);

    IOTA_INIT(Iota4D0,  20,  6, 3, 2, 2, 3, 1, 2);
    IOTA_INIT(Iota4D1,  10, 12, 5, 2, 1, 2, 2, 2);
    IOTA_INIT(Iota4D2,  25, 30, 2, 2, 3, 2, 1, 1);
    IOTA_INIT(Iota4D3,  25, 30, 2, 2, 4, 2, 4, 2);

///////////////////////////////// CPP ////////////////////////////////////
//
TEST(Iota, CPP)
{
    if (noDoubleTests<float>()) return;

    af::dim4 idims(23, 15, 1, 1);
    af::dim4 tdims(2, 2, 1, 1);
    af::dim4 fulldims;
    for(unsigned i = 0; i < 4; i++) {
        fulldims[i] = idims[i] * tdims[i];
    }

    af::array output = af::iota(idims, tdims);
    af::array tileArray = af::tile(af::moddims(af::range(af::dim4(idims.elements()), 0), idims), tdims);

    // Get result
    float* outData = new float[fulldims.elements()];
    output.host((void*)outData);

    float* tileData = new float[fulldims.elements()];
    tileArray.host((void*)tileData);

    // Compare result

    // Compare result
    for(int i = 0; i < (int)fulldims.elements(); i++)
        ASSERT_EQ(tileData[i], outData[i]) << "at: " << i << std::endl;

    // Delete
    delete[] outData;
    delete[] tileData;
}
