// Copyright (C) 2018-2020 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#include "plg_test.hpp"
#include "common_test_utils/xml_net_builder/xml_net_builder.hpp"

#include <cmath>
#include <vector>
#include <string>

using namespace InferenceEngine;
using std::map;
using std::pair;
using std::vector;
using std::string;

static const size_t DataSize = 10;  // Data size
static const size_t T = 1;   // Sequence length (for single LSTM cell is 1)
static const size_t StateSize = 5;   // State size
static const size_t G = 4;   // Number of gate

class LSTMCellNet : CommonTestUtils::V2NetBuilder {
public:
    LSTMCellNet(size_t N, size_t S, size_t D): CommonTestUtils::V2NetBuilder(buildNetworkWithOneInput("LSTM_Cell_Net"
            , { N,D }, "FP32")) {
        const size_t wSz = S*G*(D+S);
        const size_t bSz = S*G;
        const size_t wSz_b = wSz * sizeof(float);
        const size_t bSz_b = bSz * sizeof(float);

        weights = make_shared_blob<uint8_t>(TensorDesc(Precision::U8, SizeVector{ (wSz_b + bSz_b) }, Layout::C));
        weights->allocate();

        auto ptr = weights->buffer().as<float*>();
        w_blob = make_shared_blob<float>(TensorDesc(Precision::FP32, SizeVector{wSz}, 
            TensorDesc::getLayoutByDims(SizeVector{wSz})), ptr);
        b_blob = make_shared_blob<float>(TensorDesc(Precision::FP32, SizeVector{bSz}, 
            TensorDesc::getLayoutByDims(SizeVector{bSz})), ptr + wSz);

        // input layer #1 and #2
        addInputLayer("FP32", { N,S });
        addInputLayer("FP32", { N,S });

        // layer #3
        map<string, string> lstm_p {{"hidden_size", std::to_string(S)}};
        addLayer("LSTMCell", "FP32", &lstm_p,
                { {{ N,D }, { N,S }, { N,S }},
                  {{ N,S }, { N,S }} },
                wSz_b, bSz_b);

        vector<pair<string, string>> edges = {
                {"0,0", "3,3"},
                {"1,1", "3,4"},
                {"2,2", "3,5"}
        };
        model = finish(&edges);
    }

    using Filler = std::function<void(Blob::Ptr)>;

    CNNNetwork net(Filler w_filler = nullptr, Filler b_filler = nullptr) {
        w_filler(w_blob);
        b_filler(b_blob);

        Core ie;
        return ie.ReadNetwork(model, weights);
    }

private:
    string model;
    TBlob<uint8_t>::Ptr weights;
    Blob::Ptr w_blob;
    Blob::Ptr b_blob;
};

static inline bool cmp_near(float res, float ref, float eq_threshold) {
    constexpr float eps = 1e-5;

    if (eq_threshold != eps) {
        return std::abs(res-ref) < eq_threshold;
    }

    auto ref_abs = std::abs(ref);
    if (ref_abs > eps)
        return std::abs(res-ref)/ref_abs < eps;
    else
        return std::abs(res-ref) < eps;
}

/**********************************************/
/***     Test Body       **********************/
/**********************************************/

struct lstm_cell_param {
    size_t N;    // Batch size
    size_t S;    // State size
    size_t D;    // Data  size
};

class LSTMCellTestBase : public PlgTest<lstm_cell_param> {
 public:
    void runSingleLSTMTest(const std::map<std::string, std::string> & config = {},
                           float eq_threshold = 1e-5) {
        auto p = param();
        const size_t N = p.N;
        const size_t S = p.S;
        const size_t D = p.D;

        /* Broadcast value through tensors */
        const float H0 = 0.3, C0 = 0.77;

        const float Wf = 0.1, Bf = 0.35;
        const float Wi = 0.2, Bi = 0.25;
        const float Wc = 0.5, Bc = 0.15;
        const float Wo = 0.7, Bo = 0.05;

        using Vals = float[T+1];
        Vals f, i, c, o, C, H, X;

        auto _f = [](float x) { return 1/(1 + std::exp(-x)); };  // sigmoid
        auto _h = [](float x) { return std::tanh(x); };          // tanh

        H[0] = H0; C[0] = C0;

        for (int t = 1; t < T+1; t++) {  // t=0 - initial state. So time index starts from 1.
            X[t] = t;
            f[t] = _f(Wf*(H[t-1] + X[t]) + Bf);
            i[t] = _f(Wi*(H[t-1] + X[t]) + Bi);
            c[t] = _h(Wc*(H[t-1] + X[t]) + Bc);
            o[t] = _f(Wo*(H[t-1] + X[t]) + Bo);

            C[t] = f[t] * C[t-1] + i[t] * c[t];
            H[t] = o[t] * _h(C[t]);
        }

        /********  Weight and Input blob filler *****************/

        auto w_filler = [=](Blob::Ptr blob) {
            assert(blob->size() == G*S*(S+D));
            auto ptr = blob->buffer().as<float*>();

            float W[] = {Wf, Wi, Wc, Wo};
            for (int g = 0; g < G; g++)
                for (int s = 0; s < S; s++) {
                    for (int i = 0; i < D; i++) *ptr++ = W[g] / D;
                    for (int i = 0; i < S; i++) *ptr++ = W[g] / S;
                }
        };

        auto b_filler = [=](Blob::Ptr blob) {
            assert(blob->size() == G*S);
            auto ptr = blob->buffer().as<float*>();

            float B[] = {Bf, Bi, Bc, Bo};
            for (int g = 0; g < G; g++)
                for (int s = 0; s < S; s++) *ptr++ = B[g];
        };

        auto stat_filler = [=](Blob::Ptr blob, float val) {
            assert(blob->size() == N*S);
            auto ptr = blob->buffer().as<float*>();

            for (int n = 0; n < N; n++)
                for (int s = 0; s < S; s++) *ptr++ = val;
        };

        auto data_filler = [&](Blob::Ptr blob) {
            assert(blob->size() == N*T*D);
            auto ptr = blob->buffer().as<float*>();

            for (int n = 0; n < N; n++)
                for (int d = 0; d < D; d++) *ptr++ = X[1];
        };

        /*****  Output blob checkers  ************************/
        auto stat_checker = [=](Blob::Ptr blob, float val) {
            assert(blob->size() == N*S);
            auto ptr = blob->buffer().as<float*>();

            bool passed = true;
            float maxDiff = 0;
            for (int n = 0; n < N; n++)
                for (int s = 0; s < S; s++) {

                    if (!cmp_near(*ptr, val, eq_threshold)) {
                        printf("float eq %dx%d fail: %f : %f\n", n, s, *ptr, val);
                        passed = false;
                    }
                    maxDiff = std::max(std::abs(*ptr - val), maxDiff);
                    ptr++;
                }
            if (eq_threshold != 1e-5) {
                printf("max diff= %.6f\n", maxDiff);
            }
            return passed;
        };

        /************ Test Body  *****************************/

        LSTMCellNet topology(N, S, D);
        auto net = topology.net(w_filler, b_filler);

        Core ie;
        auto execNet = ie.LoadNetwork(net, device_name, config);
        auto req = execNet.CreateInferRequest();

        auto in_data    = req.GetBlob("Input0");
        auto in_h_state = req.GetBlob("Input1");
        auto in_c_state = req.GetBlob("Input2");

        data_filler(in_data);
        stat_filler(in_h_state, H0);
        stat_filler(in_c_state, C0);

        req.Infer();

        auto out_h_state = req.GetBlob("LSTMCell3.0");
        auto out_c_state = req.GetBlob("LSTMCell3.1");

        EXPECT_TRUE(stat_checker(out_h_state, H[T]));
        EXPECT_TRUE(stat_checker(out_c_state, C[T]));
    }
};

using LSTMCellTest  = LSTMCellTestBase;

TEST_P(LSTMCellTest, SingleLSTM) {
    runSingleLSTMTest();
}


static const lstm_cell_param workload[] = {{1, StateSize, DataSize}, {2, StateSize, DataSize}};
