// Copyright 2016, Tobias Hermann.
// https://github.com/Dobiasd/frugally-deep
// Distributed under the MIT License.
// (See accompanying LICENSE file or at
//  https://opensource.org/licenses/MIT)

#pragma once

#include "fdeep/layers/layer.hpp"

#include <string>

namespace fdeep {
namespace internal {

    class prelu_layer : public layer {
    public:
        explicit prelu_layer(const std::string& name, const float_vec& alpha,
            std::vector<std::size_t> shared_axes)
            : layer(name)
            , alpha_(fplus::make_shared_ref<float_vec>(alpha))
            , shared_axes_(shared_axes)
        {
        }

    protected:
        fdeep::shared_float_vec alpha_;
        std::vector<std::size_t> shared_axes_;
        tensors apply_impl(const tensors& input) const override
        {
            // We need to shift shared_axes if the original Keras tensor
            // was one or two dimensional.
            // We detect this by checking if the axes indicated in shared_axes
            // has length 1.
            // For this to work we need to remove axes with length 1
            // from shared axes in Python.
            std::vector<std::size_t> shared_axes_shifted;
            std::size_t shift = 0;
            for (std::size_t i = 0; i < shared_axes_.size(); ++i) {
                if ((shared_axes_[i] == 1 && input[0].shape().height_ == 1) || (shared_axes_[i] == 2 && input[0].shape().width_ == 1)) {
                    shift++;
                }
                shared_axes_shifted.push_back(shared_axes_[i] + shift);
            }

            const bool height_shared = fplus::is_elem_of(1, shared_axes_shifted);
            const bool width_shared = fplus::is_elem_of(2, shared_axes_shifted);
            const bool channels_shared = fplus::is_elem_of(3, shared_axes_shifted);
            const size_t width = width_shared ? 1 : input[0].shape().width_;
            const size_t depth = channels_shared ? 1 : input[0].shape().depth_;

            fdeep::tensor out(input[0].shape(), 1.0f);
            for (std::size_t y = 0; y < out.shape().height_; ++y) {
                for (std::size_t x = 0; x < out.shape().width_; ++x) {
                    for (std::size_t z = 0; z < out.shape().depth_; ++z) {
                        if (input[0].get_ignore_rank(tensor_pos(y, x, z)) > 0) {
                            out.set_ignore_rank(tensor_pos(y, x, z),
                                input[0].get_ignore_rank(tensor_pos(y, x, z)));
                        } else {
                            const size_t y_temp = height_shared ? 0 : y;
                            const size_t x_temp = width_shared ? 0 : x;
                            const size_t z_temp = channels_shared ? 0 : z;
                            const size_t pos = y_temp * width * depth + x_temp * depth + z_temp;
                            out.set_ignore_rank(tensor_pos(y, x, z), (*alpha_)[pos] * input[0].get_ignore_rank(tensor_pos(y, x, z)));
                        }
                    }
                }
            }
            return { out };
        }
    };

}
}