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

#include "transformations/batch_norm_decomposition.hpp"

#include <memory>
#include <vector>

#include <ngraph/opsets/opset1.hpp>


void ngraph::pass::BatchNormDecomposition::batch_norm_decomposition() {
    Shape shape{2, 2, 1, 1};
    auto input = make_shared<pattern::op::Label>(element::f32, shape);
    auto mean_shape = Shape{2};
    auto mean = make_shared<pattern::op::Label>(element::f32, mean_shape);
    auto var_shape = Shape{2};
    auto var = make_shared<pattern::op::Label>(element::f32, var_shape);
    auto gamma_shape = Shape{2};
    auto gamma = make_shared<pattern::op::Label>(element::f32, gamma_shape);
    auto beta_shape = Shape{2};
    auto beta = make_shared<pattern::op::Label>(element::f32, beta_shape);
    auto bn = make_shared<opset1::BatchNormInference>(input, gamma, beta, mean, var, 0.001);

    ngraph::graph_rewrite_callback callback = [input, gamma, beta, mean, var](ngraph::pattern::Matcher &m) {
        auto pattern_map = m.get_pattern_map();

        auto m_input = pattern_map[input];
        auto m_gamma = pattern_map[gamma];
        auto m_beta = pattern_map[beta];
        auto m_mean = pattern_map[mean];
        auto m_var = pattern_map[var];

        // TODO: check that all input shapes are static

        auto m_bn = dynamic_pointer_cast<opset1::BatchNormInference>(m.get_match_root());
        if (!m_bn) {
            return false;
        }

        // The code above represents this formulas
        //  scale = 1. / np.sqrt(variance + eps)
        //  shift = (mean * (-1)) * scale
        auto input_type = m_input->get_element_type();
        auto scale = make_shared<ngraph::opset1::Divide>(
                opset1::Constant::create(input_type, Shape{}, {1}),
                make_shared<opset1::Power>(
                        make_shared<opset1::Add>(
                                m_var,
                                opset1::Constant::create(input_type, Shape{}, {m_bn->get_eps_value()})),
                        opset1::Constant::create(input_type, Shape{}, {0.5})));

        auto shift = make_shared<opset1::Multiply>(
                scale,
                make_shared<opset1::Multiply>(
                        m_mean,
                        opset1::Constant::create(m_input->get_element_type(), Shape{}, {-1})));

        // Expand Scale, Shift, Gamma and Beta to be aligned with layout
        size_t dims_to_add = m_input->get_shape().size() - 2;
        Shape gamma_shape = m_gamma->get_shape();
        for (size_t i = 0; i < dims_to_add; ++i) gamma_shape.push_back(1);
        auto gamma_aligned = make_shared<opset1::Reshape>(m_gamma, opset1::Constant::create(element::i64, Shape{gamma_shape.size()}, gamma_shape), true);

        Shape beta_shape = m_beta->get_shape();
        for (size_t i = 0; i < dims_to_add; ++i) beta_shape.push_back(1);
        auto beta_aligned = make_shared<opset1::Reshape>(m_beta, opset1::Constant::create(element::i64, Shape{beta_shape.size()}, beta_shape), true);

        Shape scale_shape = scale->get_shape();
        for (size_t i = 0; i < dims_to_add; ++i) scale_shape.push_back(1);
        auto scale_aligned = make_shared<opset1::Reshape>(scale, opset1::Constant::create(element::i64, Shape{scale_shape.size()}, scale_shape), true);

        Shape shift_shape = scale->get_shape();
        for (size_t i = 0; i < dims_to_add; ++i) shift_shape.push_back(1);
        auto shift_aligned = make_shared<opset1::Reshape>(shift, opset1::Constant::create(element::i64, Shape{shift_shape.size()}, shift_shape), true);

        // Connect: Mul(input, scale)->Add(mul, shift)->Mul(add, gamma)->Add(mul, beta)
        auto result = make_shared<opset1::Add>(
                make_shared<opset1::Multiply>(
                        make_shared<opset1::Add>(
                                make_shared<opset1::Multiply>(
                                        m_input,
                                        scale_aligned),
                                shift_aligned),
                        gamma_aligned),
                beta_aligned);

        result->set_friendly_name(m_bn->get_friendly_name());
        replace_node(m_bn, result);

        return true;
    };

    auto m = std::make_shared<ngraph::pattern::Matcher>(bn, "BatchNormDecomposition");
    this->add_matcher(m, callback, PassProperty::CHANGE_DYNAMIC_STATE);
}