/*******************************************************************************
* Copyright 2018 Intel Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* 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.
*******************************************************************************/

#include <mkldnn_types.h>
#include "gtest/gtest.h"
#include "mkldnn_test_common.hpp"
#include "mkldnn.hpp"

namespace mkldnn {

template <typename T> inline T scale_shift_fwd(T s_val, T w_val, T b_val) {
    return s_val*w_val + b_val;
}

template <typename T> inline T prelu_fwd(T s_val, T w_val) {
    return s_val >= 0 ? s_val : w_val*s_val;
}

template <typename data_t>
struct depthwise_test_params {
    engine::kind engine_kind;
    algorithm alg_kind;
    memory::format data_format;
    memory::dims dims;
};

template <typename data_t>
void check_depthwise_fwd(const depthwise_test_params<data_t> &p,
        const memory::desc &md, const memory &src, const memory &weights,
        const memory &bias, bool with_bias, const memory &dst)
{
    data_t *src_data     = (data_t *)src.get_data_handle();
    data_t *weights_data = (data_t *)weights.get_data_handle();
    data_t *bias_data    = with_bias ? (data_t *)bias.get_data_handle() : nullptr;
    data_t *dst_data     = (data_t *)dst.get_data_handle();

    const memory::desc src_d = src.get_primitive_desc().desc();
    const memory::desc weights_d = weights.get_primitive_desc().desc();
    const memory::desc dst_d = dst.get_primitive_desc().desc();

    ASSERT_EQ(md.data.data_type, memory::data_type::f32); // TODO: type assert

    int N = md.data.ndims > 0 ? md.data.dims[0] : 1;
    int C = md.data.ndims > 1 ? md.data.dims[1] : 1;
    int H = md.data.ndims > 2 ? md.data.dims[2] : 1;
    int W = md.data.ndims > 3 ? md.data.dims[3] : 1;

    for (int n = 0; n < N; ++n) {
        for (int c = 0; c < C; ++c) {
            for (int h = 0; h < H; ++h) {
                for (int w = 0; w < W; ++w) {
                    int idx = n*C*H*W + c*H*W + h*W + w;

                    data_t s_val = src_data[map_index(src_d, idx)];
                    data_t w_val = weights_data[map_index(weights_d, c)];
                    data_t b_val = with_bias ? bias_data[map_index(bias.get_primitive_desc().desc(), c)] : 0;

                    data_t ref_d = 0;
                    switch (p.alg_kind) {
                        case depthwise_scale_shift:
                            ref_d = scale_shift_fwd(s_val, w_val, b_val);
                            break;
                        case depthwise_prelu:
                            ref_d = prelu_fwd(s_val, w_val);
                            break;
                        default:
                            assert(!"unknown alg_kind");
                    }

                    EXPECT_NEAR(dst_data[map_index(dst_d, idx)], ref_d, 1.e-6);
                }
            }
        }
    }
}

template <typename data_t>
class depthwise_test : public ::testing::TestWithParam<depthwise_test_params<data_t>> {
private:
    std::shared_ptr<memory> src;
    std::shared_ptr<memory> weights;
    std::shared_ptr<memory> bias;
    std::shared_ptr<memory> dst;
    std::shared_ptr<memory> workspace;
    std::shared_ptr<memory::desc> src_desc;
    std::shared_ptr<memory::desc> dst_desc;
    std::shared_ptr<memory::desc> weights_desc;
    std::shared_ptr<memory::desc> bias_desc;
    std::shared_ptr<depthwise_forward::primitive_desc> depthwise_prim_desc;
    depthwise_test_params<data_t> p;
    std::shared_ptr<engine> eng;
    memory::data_type data_type;
    int data_size;
    int weights_size;

protected:
    virtual void SetUp() {
        p = ::testing::TestWithParam<depthwise_test_params<data_t>>::GetParam();

        ASSERT_TRUE(p.engine_kind == engine::kind::cpu);
        eng.reset(new engine(p.engine_kind, 0));

        data_type = data_traits<data_t>::data_type;
        ASSERT_EQ(data_type, mkldnn::memory::data_type::f32);

        data_size = p.dims[0] * p.dims[1] * p.dims[2] * p.dims[3];
        weights_size = p.dims[1];

        Forward();
    }

    void Forward() {
        bool with_bias = p.alg_kind == depthwise_scale_shift;

        memory::dims dims = p.data_format == mkldnn_nc ? memory::dims({p.dims[0], p.dims[1]}) : p.dims;

        src_desc.reset(new memory::desc(dims, data_type, p.data_format));
        dst_desc.reset(new memory::desc(dims, data_type, p.data_format));
        src.reset(new memory({*src_desc, *eng}));
        dst.reset(new memory({*dst_desc, *eng}));
        fill_data<data_t>(data_size, (data_t *)src->get_data_handle(),
                          data_t(0), data_t(1));

        weights_desc.reset(new memory::desc({dims[1]}, data_type, memory::format::x));
        weights.reset(new memory({*weights_desc, *eng}));
        fill_data<data_t>(weights_size, (data_t *)weights->get_data_handle(),
                          data_t(0), data_t(1));

        if (with_bias) {
            bias_desc.reset(new memory::desc({dims[1]}, data_type, memory::format::x));
            bias.reset(new memory({*bias_desc, *eng}));
            fill_data<data_t>(weights_size, (data_t *) bias->get_data_handle(),
                              data_t(0), data_t(1));
        }

        std::vector<primitive> pipeline;
        auto depthwise_desc = with_bias
                              ? depthwise_forward::desc(prop_kind::forward_training, p.alg_kind, *src_desc, *dst_desc, *weights_desc, *bias_desc)
                              : depthwise_forward::desc(prop_kind::forward_training, p.alg_kind, *src_desc, *dst_desc, *weights_desc);
        depthwise_prim_desc.reset(new depthwise_forward::primitive_desc(depthwise_desc, *eng));

        auto depthwise = with_bias
                         ? depthwise_forward(*depthwise_prim_desc, *src, *weights, *bias, *dst)
                         : depthwise_forward(*depthwise_prim_desc, *src, *weights, *dst);

        pipeline.push_back(depthwise);
        auto s = stream(stream::kind::lazy);
        s.submit(pipeline).wait();

        check_depthwise_fwd(p, *src_desc, *src, *weights, *bias, with_bias, *dst);
    }
};

using depthwise_test_float = depthwise_test<float>;
using depthwise_test_params_float = depthwise_test_params<float>;

TEST_P(depthwise_test_float, TestsDepthwise)
{
}

#define EXPAND(args) args

#define EXPAND_FORMATS(data) memory::format::data

#define ENGINE engine::kind::cpu

#define PARAMS(alg, data, mb, c, h, w) \
    depthwise_test_params_float { ENGINE, algorithm::alg, \
    EXPAND_FORMATS(data), {mb, c, h, w} }

#define PARAMS_ALL_ALG(...) \
    EXPAND(PARAMS(depthwise_scale_shift, __VA_ARGS__)), \
    EXPAND(PARAMS(depthwise_prelu, __VA_ARGS__))

#define INST_TEST_CASE(str, ...) INSTANTIATE_TEST_CASE_P( \
        str, depthwise_test_float, ::testing::Values(__VA_ARGS__))

INST_TEST_CASE(Simple_NC,
    PARAMS_ALL_ALG(nc, 2, 8, 1, 1),
    PARAMS_ALL_ALG(nc, 2, 16, 1, 1),
    PARAMS_ALL_ALG(nc, 2, 16, 1, 1),
    PARAMS_ALL_ALG(nc, 2, 16, 1, 1),
    PARAMS_ALL_ALG(nc, 2, 16, 1, 1),
    PARAMS_ALL_ALG(nc, 10, 10, 1, 1),
    PARAMS_ALL_ALG(nc, 256, 64, 1, 1),
    PARAMS_ALL_ALG(nc, 1, 1, 1, 1),
    PARAMS_ALL_ALG(nc, 3, 5, 1, 1)
);

INST_TEST_CASE(Simple_NCHW,
    PARAMS_ALL_ALG(nchw, 2, 8, 4, 4),
    PARAMS_ALL_ALG(nchw, 2, 16, 4, 4),
    PARAMS_ALL_ALG(nchw, 2, 16, 8, 8),
    PARAMS_ALL_ALG(nchw, 2, 16, 16, 8),
    PARAMS_ALL_ALG(nchw, 2, 16, 10, 8),
    PARAMS_ALL_ALG(nchw, 10, 10, 10, 10),
    PARAMS_ALL_ALG(nchw, 256, 64, 8, 16),
    PARAMS_ALL_ALG(nchw, 1, 1, 1, 1),
    PARAMS_ALL_ALG(nchw, 3, 5, 7, 11)
);

INST_TEST_CASE(Simple_Blocked,
    PARAMS_ALL_ALG(nChw8c, 2, 8, 4, 4),
    PARAMS_ALL_ALG(nChw8c, 2, 16, 4, 4),
    PARAMS_ALL_ALG(nChw8c, 2, 16, 8, 8),
    PARAMS_ALL_ALG(nChw8c, 2, 16, 16, 8),
    PARAMS_ALL_ALG(nChw8c, 2, 32, 10, 8),
    PARAMS_ALL_ALG(nChw8c, 256, 64, 8, 16)
);

INST_TEST_CASE(Simple_Blocked16,
    PARAMS_ALL_ALG(nChw16c, 2, 16, 4, 4),
    PARAMS_ALL_ALG(nChw16c, 2, 16, 8, 8),
    PARAMS_ALL_ALG(nChw16c, 2, 16, 16, 8),
    PARAMS_ALL_ALG(nChw16c, 2, 32, 10, 8),
    PARAMS_ALL_ALG(nChw16c, 256, 64, 8, 16)
);

}
