/* ******************************************************************************
 *
 *
 * This program and the accompanying materials are made available under the
 * terms of the Apache License, Version 2.0 which is available at
 * https://www.apache.org/licenses/LICENSE-2.0.
 *
 *  See the NOTICE file distributed with this work for additional
 *  information regarding copyright ownership.
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations
 * under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 ******************************************************************************/

//
// @author raver119@gmail.com
//

#ifndef LIBND4J_TADTESTS_H
#define LIBND4J_TADTESTS_H

#include "testlayers.h"
#include <array/NDArray.h>
#include <helpers/TAD.h>
#include <array>
#include <helpers/ConstantTadHelper.h>

using namespace sd;

class TadTests : public testing::Test {
public:
    int numLoops = 100000000;

    int extLoops = 1000;
    int intLoops = 1000;
};

TEST_F(TadTests, Test4DTad1) {

    NDArray*  arraySource = sd::NDArrayFactory::linspace(1.0f, 10000.0f, 10000);

    Nd4jLong badShape[]  = {4, 2, 1, 4, 4, 80, 16, 4, 1, 8192, -1, 99};
    Nd4jLong goodShape[] = {4, 2, 1, 4, 4, 16, 16, 4, 1, 8192,  1, 99};

    std::vector<float> buff = arraySource->getBufferAsVector<float>();

    NDArray* arrayExp = new NDArray(buff.data(), goodShape);
    NDArray* arrayBad = new NDArray(buff.data(), badShape);

    int dim = 1;
    shape::TAD tad;
    tad.init(arrayBad->shapeInfo(), &dim, 1);
    tad.createTadOnlyShapeInfo();
    tad.createOffsets();

    int exp[] = { 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95 };
    for (int e = 0; e < 32; e++)
        ASSERT_EQ((int) tad.tadOffsets[e],  exp[e]);

    delete arrayExp;
    delete arrayBad;
    delete arraySource;
}

TEST_F(TadTests, TestNumTads1) {
    auto x = NDArrayFactory::create<float>('c', {2, 3});
    auto y = NDArrayFactory::create<float>('c', {2, 2});

    std::vector<int> dim({0});

    Nd4jLong tadLengthX = shape::tadLength(x.shapeInfo(), dim.data(), dim.size());
    Nd4jLong numTadsX = x.lengthOf() / tadLengthX;

    Nd4jLong tadLengthY = shape::tadLength(y.shapeInfo(), dim.data(), dim.size());
    Nd4jLong numTadsY = y.lengthOf() / tadLengthY;

    ASSERT_EQ(2, tadLengthX);
    ASSERT_EQ(3, numTadsX);

    ASSERT_EQ(2, tadLengthY);
    ASSERT_EQ(2, numTadsY);
}

TEST_F(TadTests, TestShapeTad_1) {

    float buff[]  = {1,2,3,4,5,6,7,8,9,10,11,12,13,14,16,16,17,18,19,20,21,22,23,24};
    Nd4jLong shapeInfo[] = {3, 2, 3, 4, 12, 4, 1, 8192, 1, 99};

    NDArray input(buff, shapeInfo);

    std::vector<int> dimensions = {0,1,2};
    Nd4jLong tadLength = shape::tadLength(input.shapeInfo(), dimensions.data(), dimensions.size());
    Nd4jLong numTads = input.lengthOf() / tadLength;

    shape::TAD tad;
    tad.init(input.shapeInfo(), dimensions.data(), dimensions.size());
    tad.createTadOnlyShapeInfo();
    tad.createOffsets();

    auto tadShapeInfo = new Nd4jLong[shape::shapeInfoLength(tad.tadOnlyShapeInfo[0])];
    std::memcpy(tadShapeInfo, tad.tadOnlyShapeInfo, shape::shapeInfoByteLength(tad.tadOnlyShapeInfo));

    float* tadBuff = reinterpret_cast<float*>(input.buffer()) + tad.tadOffsets[0];
    NDArray tadArr(tadBuff, tadShapeInfo);

    ASSERT_TRUE(numTads==1);
    ASSERT_TRUE(input.isSameShapeStrict(tadArr));
    ASSERT_TRUE(input.equalsTo(&tadArr));

	delete[] tadShapeInfo;
}

TEST_F(TadTests, TadNoAxis_1) {
    auto array = NDArrayFactory::create<float>('c', {2, 3});

    shape::TAD tad;
    tad.init(array.shapeInfo(), nullptr, 0);
    tad.createTadOnlyShapeInfo();
    tad.createOffsets();

    ASSERT_TRUE(tad.wholeThing);

    ASSERT_TRUE(shape::equalsStrict(tad.tadOnlyShapeInfo, array.shapeInfo()));
}

TEST_F(TadTests, TadEdgeCase_1) {
    auto array = NDArrayFactory::create<float>('c', {5, 4, 1});
    auto exp = NDArrayFactory::create<float>('c', {5, 4});
    array.linspace(1);

    auto tad = array(0, {2});

    ASSERT_TRUE(exp.isSameShape(tad));
}

TEST_F(TadTests, TestEdgeCase_2) {

    auto array = NDArrayFactory::create<float>('f', {2, 3, 1}, {1, 4, 2, 5, 3, 6});

    for (int e = 0 ; e < array.lengthOf(); e++) {
        auto tad = array(e, {0,1});
        ASSERT_NEAR(tad.e<float>(0), array.e<float>(e), 1e-5);
    }
}

TEST_F(TadTests, TadEdgeCase_2) {
    auto array = NDArrayFactory::create<float>('c', {2, 3, 4});

    auto tad = array(0, {0,2});

    ASSERT_EQ(3, tad.lengthOf());
}


TEST_F(TadTests, test_Tad_Ews_optimization_1) {
    shape::TAD xTad;

    std::array<int,2> array = {1,2};
    ASSERT_TRUE(xTad.dimensionsDescending(3, array.data(), array.size()));
}

TEST_F(TadTests, test_Tad_Ews_optimization_2) {
    shape::TAD xTad;

    std::array<int,2> array = {0,2};
    ASSERT_FALSE(xTad.dimensionsDescending(3, array.data(), array.size()));
}

TEST_F(TadTests, test_Tad_Ews_optimization_3) {
    shape::TAD xTad;

    std::array<int,1> array = {1};
    ASSERT_TRUE(xTad.dimensionsDescending(2, array.data(), array.size()));
}

TEST_F(TadTests, test_Tad_Ews_optimization_4) {
    shape::TAD xTad;

    std::array<int,1> array = {0};
    ASSERT_TRUE(xTad.dimensionsDescending(1, array.data(), array.size()));
}

TEST_F(TadTests, test_Tad_Ews_optimization_5) {
    shape::TAD xTad;

    std::array<int,2> array = {2,3};
    ASSERT_TRUE(xTad.dimensionsDescending(4, array.data(), array.size()));
}

TEST_F(TadTests, test_TAD_empty_dims_1) {
    Nd4jLong xShape[8] = {2, 150, 1, 3, 1, 16384, 3, 99};
    shape::TAD xTad;
    xTad.init(xShape, reinterpret_cast<int*>(112L), 0);
    xTad.createTadOnlyShapeInfo();
    xTad.createOffsets();
}

TEST_F(TadTests, test_tad_order_1) {
    Nd4jLong xShape[8] = {2, 150, 10, 10, 1, 8192, 1, 99};
    Nd4jLong tShape[8] = {2, 1, 10, 1, 1, 8192, 1, 99};
    shape::TAD xTad;
    int dim = 1;
    xTad.init(xShape, &dim, 1);
    xTad.createTadOnlyShapeInfo();

    ASSERT_TRUE(shape::equalsStrict(tShape, xTad.tadOnlyShapeInfo));
}

TEST_F(TadTests, test_tad_order_2) {
    Nd4jLong xShape[8] = {2, 150, 10, 10, 1, 8192, 1, 99};
    Nd4jLong tShape[8] = {2, 1, 150, 1, 10, 8192, 10, 99};
    shape::TAD xTad;
    int dim = 0;
    xTad.init(xShape, &dim, 1);
    xTad.createTadOnlyShapeInfo();

    ASSERT_TRUE(shape::equalsStrict(tShape, xTad.tadOnlyShapeInfo));
}


TEST_F(TadTests, test_tad_order_3) {
    Nd4jLong xShape[10] = {3, 10, 20, 30, 600 ,30, 1, 8192, 1, 99};
    Nd4jLong tShape[8] = {2, 1, 30, 1, 1, 8192, 1, 99};
    shape::TAD xTad;
    int dim = 2;
    xTad.init(xShape, &dim, 1);
    xTad.createTadOnlyShapeInfo();

    ASSERT_TRUE(shape::equalsStrict(tShape, xTad.tadOnlyShapeInfo));
}


TEST_F(TadTests, test_tad_order_4) {
    Nd4jLong xShape[10] = {3, 10, 20, 30, 600 ,30, 1, 8192, 1, 99};
    Nd4jLong tShape[8] = {2, 20, 30, 30, 1, 8192, 1, 99};
    shape::TAD xTad;
    int dim[2] = {1, 2};
    xTad.init(xShape, dim, 2);
    xTad.createTadOnlyShapeInfo();

    ASSERT_TRUE(shape::equalsStrict(tShape, xTad.tadOnlyShapeInfo));
}

TEST_F(TadTests, test_column_1) {
    auto x = NDArrayFactory::create<float>('c', {5, 2});
    auto tadPack = sd::ConstantTadHelper::getInstance().tadForDimensions(x.shapeInfo(), 0);

    ASSERT_EQ(1, shape::rank(tadPack.primaryShapeInfo()));
    ASSERT_EQ(5, shape::length(tadPack.primaryShapeInfo()));
    ASSERT_TRUE(shape::isVector(tadPack.primaryShapeInfo()));

    auto scalarViewPack = sd::ConstantTadHelper::getInstance().tadForDimensions(tadPack.primaryShapeInfo(), 0);

    ASSERT_TRUE(shape::equalsStrict(tadPack.primaryShapeInfo(), scalarViewPack.primaryShapeInfo()));
}

///////////////////////////////////////////////////////////////////
TEST_F(TadTests, calcOffsets_1) {

    Nd4jLong shapeInfoF[10]  = {3, 2,3,4,  1,2,6,   8192, 1, 102};
    Nd4jLong shapeInfoC[10]  = {3, 2,3,4,  12,4,1,  8192, 1, 99};
    Nd4jLong shapeInfoFC[10] = {3, 2,3,4,  1,2,6,   8192, 1, 99};;

    Nd4jLong expOffsetsF[24] = {0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23};
    Nd4jLong expOffsetsC[24] = {0,12,4,16,8,20,1,13,5,17,9,21,2,14,6,18,10,22,3,15,7,19,11,23};

    Nd4jLong offsets[24];

    shape::calcOffsets(shapeInfoF, offsets, 'f');

    for (int e = 0; e < 24; e++)
        ASSERT_TRUE(offsets[e] == expOffsetsF[e]);

    shape::calcOffsets(shapeInfoC, offsets, 'f');

    for (int e = 0; e < 24; e++)
        ASSERT_TRUE(offsets[e] == expOffsetsC[e]);

    shape::calcOffsets(shapeInfoFC, offsets, 'f');

    for (int e = 0; e < 24; e++)
        ASSERT_TRUE(offsets[e] == expOffsetsF[e]);
}


/////////////////////////////////////////////////////////////////
TEST_F(TadTests, outerArrayIndexes_1) {

    NDArray x('c', {2,3,4,5}, sd::DataType::FLOAT32);
    int maxIdxs[120];

    NDArray y1('c', {3,5}, sd::DataType::FLOAT32);
    const std::vector<int> dimsToExclude1 = {0,2};
    const int n1[] = {20,25,30,35,  80,85,90,95};
    int minIdx = 5;

    int N = shape::outerArrayIndexes(maxIdxs, minIdx, x.shapeInfo(), y1.shapeInfo(), dimsToExclude1.data());
    ASSERT_TRUE(N == x.lengthOf()/y1.lengthOf());
    for(int i = 0; i < N; ++i)
        ASSERT_TRUE(n1[i] == maxIdxs[i]);

    NDArray y2('c', {4,5}, sd::DataType::FLOAT32);
    const std::vector<int> dimsToExclude2 = {0,1};
    const int n2[] = {12,32,52,  72,92,112};
    minIdx = 12;

    N = shape::outerArrayIndexes(maxIdxs, minIdx, x.shapeInfo(), y2.shapeInfo(), dimsToExclude2.data());
    ASSERT_TRUE(N == x.lengthOf()/y2.lengthOf());
    for(int i = 0; i < N; ++i)
        ASSERT_TRUE(n2[i] == maxIdxs[i]);

    NDArray y3('c', {2,5}, sd::DataType::FLOAT32);
    const std::vector<int> dimsToExclude3 = {1,2};
    const int n3[] = {64,69,74,79,84,89,94,99,104,109,114,119};
    minIdx = 9;

    N = shape::outerArrayIndexes(maxIdxs, minIdx, x.shapeInfo(), y3.shapeInfo(), dimsToExclude3.data());
    ASSERT_TRUE(N == x.lengthOf()/y3.lengthOf());
    for(int i = 0; i < N; ++i)
        ASSERT_TRUE(n3[i] == maxIdxs[i]);

    NDArray y4('c', {2,3}, sd::DataType::FLOAT32);
    const std::vector<int> dimsToExclude4 = {2,3};
    const int n4[] = {20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39};
    minIdx = 1;

    N = shape::outerArrayIndexes(maxIdxs, minIdx, x.shapeInfo(), y4.shapeInfo(), dimsToExclude4.data());
    ASSERT_TRUE(N == x.lengthOf()/y4.lengthOf());
    for(int i = 0; i < N; ++i)
        ASSERT_TRUE(n4[i] == maxIdxs[i]);

    NDArray y5('c', {2,4}, sd::DataType::FLOAT32);
    const std::vector<int> dimsToExclude5 = {1,3};
    const int n5[] = {65,66,67,68,69, 85,86,87,88,89, 105,106,107,108,109};
    minIdx = 5;

    N = shape::outerArrayIndexes(maxIdxs, minIdx, x.shapeInfo(), y5.shapeInfo(), dimsToExclude5.data());
    ASSERT_TRUE(N == x.lengthOf()/y5.lengthOf());
    for(int i = 0; i < N; ++i)
        ASSERT_TRUE(n5[i] == maxIdxs[i]);

    NDArray y6('c', {2,3,4}, sd::DataType::FLOAT32);
    const std::vector<int> dimsToExclude6 = {3};
    const int n6[] = {65,66,67,68,69};
    minIdx = 13;

    N = shape::outerArrayIndexes(maxIdxs, minIdx, x.shapeInfo(), y6.shapeInfo(), dimsToExclude6.data());
    ASSERT_TRUE(N == x.lengthOf()/y6.lengthOf());
    for(int i = 0; i < N; ++i)
        ASSERT_TRUE(n6[i] == maxIdxs[i]);

    NDArray y7('c', {4}, sd::DataType::FLOAT32);
    const std::vector<int> dimsToExclude7 = {0,1,3};
    const int n7[] = {15,16,17,18,19, 35,36,37,38,39, 55,56,57,58,59, 75,76,77,78,79, 95,96,97,98,99, 115,116,117,118,119};
    minIdx = 3;

    N = shape::outerArrayIndexes(maxIdxs, minIdx, x.shapeInfo(), y7.shapeInfo(), dimsToExclude7.data());
    ASSERT_TRUE(N == x.lengthOf()/y7.lengthOf());
    for(int i = 0; i < N; ++i)
        ASSERT_TRUE(n7[i] == maxIdxs[i]);

    NDArray y8('c', {5}, sd::DataType::FLOAT32);
    const std::vector<int> dimsToExclude8 = {0,1,2};
    const int n8[] = {0,5,10,15,  20,25,30,35, 40,45,50,55, 60,65,70,75, 80,85,90,95, 100,105,110,115};
    minIdx = 0;

    N = shape::outerArrayIndexes(maxIdxs, minIdx, x.shapeInfo(), y8.shapeInfo(), dimsToExclude8.data());
    ASSERT_TRUE(N == x.lengthOf()/y8.lengthOf());
    for(int i = 0; i < N; ++i)
        ASSERT_TRUE(n8[i] == maxIdxs[i]);

    NDArray y9('c', {2}, sd::DataType::FLOAT32);
    const std::vector<int> dimsToExclude9 = {1,2,3};
    const int n9[] = {60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119};
    minIdx = 1;

    N = shape::outerArrayIndexes(maxIdxs, minIdx, x.shapeInfo(), y9.shapeInfo(), dimsToExclude9.data());
    ASSERT_TRUE(N == x.lengthOf()/y9.lengthOf());
    for(int i = 0; i < N; ++i)
        ASSERT_TRUE(n9[i] == maxIdxs[i]);

    NDArray y10('c', {3,4,5}, sd::DataType::FLOAT32);
    const std::vector<int> dimsToExclude10 = {0};
    const int n10[] = {11, 71};
    minIdx = 11;

    N = shape::outerArrayIndexes(maxIdxs, minIdx, x.shapeInfo(), y10.shapeInfo(), dimsToExclude10.data());
    ASSERT_TRUE(N == x.lengthOf()/y10.lengthOf());
    for(int i = 0; i < N; ++i)
        ASSERT_TRUE(n10[i] == maxIdxs[i]);

    NDArray y11('c', {2,4,5}, sd::DataType::FLOAT32);
    const std::vector<int> dimsToExclude11 = {1};
    const int n11[] = {66, 86, 106};
    minIdx = 26;

    N = shape::outerArrayIndexes(maxIdxs, minIdx, x.shapeInfo(), y11.shapeInfo(), dimsToExclude11.data());
    ASSERT_TRUE(N == x.lengthOf()/y11.lengthOf());
    for(int i = 0; i < N; ++i)
        ASSERT_TRUE(n11[i] == maxIdxs[i]);

    NDArray y12('c', {3,2}, sd::DataType::FLOAT32);
    const std::vector<int> dimsToExclude12 = {0,2};
    const int n12[] = {0,2,4,5,7,9,10,12,14,15,17,19,60,62,64,65,67,69,70,72,74,75,77,79};
    minIdx = 0;

    N = shape::outerArrayIndexes(maxIdxs, minIdx, x.shapeInfo(), y12.shapeInfo(), dimsToExclude12.data());
    for(int i = 0; i < N; ++i)
        ASSERT_TRUE(n12[i] == maxIdxs[i]);

    NDArray y13('c', {3,2}, sd::DataType::FLOAT32);
    const std::vector<int> dimsToExclude13 = {0,2};
    const int n13[] = {1,3,6,8,11,13,16,18,61,63,66,68,71,73,76,78};
    minIdx = 1;

    N = shape::outerArrayIndexes(maxIdxs, minIdx, x.shapeInfo(), y13.shapeInfo(), dimsToExclude13.data());
    for(int i = 0; i < N; ++i)
        ASSERT_TRUE(n13[i] == maxIdxs[i]);

    NDArray y14('c', {4,5}, sd::DataType::FLOAT32);
    const int n14[] = {12,32,52,  72,92,112};
    minIdx = 12;

    N = shape::outerArrayIndexes(maxIdxs, minIdx, x.shapeInfo(), y14.shapeInfo(), nullptr);
    ASSERT_TRUE(N == x.lengthOf()/y14.lengthOf());
    for(int i = 0; i < N; ++i)
        ASSERT_TRUE(n14[i] == maxIdxs[i]);

    NDArray y15('c', {3,4,5}, sd::DataType::FLOAT32);
    const int n15[] = {11, 71};
    minIdx = 11;

    N = shape::outerArrayIndexes(maxIdxs, minIdx, x.shapeInfo(), y15.shapeInfo(), nullptr);
    ASSERT_TRUE(N == x.lengthOf()/y15.lengthOf());
    for(int i = 0; i < N; ++i)
        ASSERT_TRUE(n15[i] == maxIdxs[i]);
}



#endif //LIBND4J_TADTESTS_H