/********************************************************
 *  ██████╗  ██████╗████████╗██╗
 * ██╔════╝ ██╔════╝╚══██╔══╝██║
 * ██║  ███╗██║        ██║   ██║
 * ██║   ██║██║        ██║   ██║
 * ╚██████╔╝╚██████╗   ██║   ███████╗
 *  ╚═════╝  ╚═════╝   ╚═╝   ╚══════╝
 * Geophysical Computational Tools & Library (GCTL)
 *
 * Copyright (c) 2022  Yi Zhang (yizhang-geo@zju.edu.cn)
 *
 * GCTL is distributed under a dual licensing scheme. You can redistribute 
 * it and/or modify it under the terms of the GNU Lesser General Public 
 * License as published by the Free Software Foundation, either version 2 
 * of the License, or (at your option) any later version. You should have 
 * received a copy of the GNU Lesser General Public License along with this 
 * program. If not, see <http://www.gnu.org/licenses/>.
 * 
 * If the terms and conditions of the LGPL v.2. would prevent you from using 
 * the GCTL, please consider the option to obtain a commercial license for a 
 * fee. These licenses are offered by the GCTL's original author. As a rule, 
 * licenses are provided "as-is", unlimited in time for a one time fee. Please 
 * send corresponding requests to: yizhang-geo@zju.edu.cn. Please do not forget 
 * to include some description of your company and the realm of its activities. 
 * Also add information on how to contact you by electronic and paper mail.
 ******************************************************/

#include "hlayer_avgpooling.h"

gctl::avgpooling::avgpooling(){}

gctl::avgpooling::avgpooling(int in_rows, int in_cols, int pool_rows, int pool_cols, int stride_rows, int stride_cols, 
    pad_type_e pl_type, activation_type_e acti_type)
{
    init_avgpooling(in_rows, in_cols, pool_rows, pool_cols, stride_rows, stride_cols, pl_type, acti_type);
}

gctl::avgpooling::~avgpooling(){}

void gctl::avgpooling::init_avgpooling(int in_rows, int in_cols, int pool_rows, int pool_cols, int stride_rows, 
    int stride_cols, pad_type_e pl_type, activation_type_e acti_type)
{
    i_rows_ = in_rows; i_cols_ = in_cols;
    p_rows_ = pool_rows; p_cols_ = pool_cols;
    s_rows_ = stride_rows; s_cols_ = stride_cols;
    p_type_ = pl_type;

    if (acti_type == Identity) activator_ = new identity;
    else if (acti_type == Mish) activator_ = new mish;
    else if (acti_type == ReLU) activator_ = new relu;
    else if (acti_type == PReLU) activator_ = new prelu;
    else if (acti_type == Sigmoid) activator_ = new sigmoid;
    else if (acti_type == SoftMax) activator_ = new softmax;
    else if (acti_type == Tanh) activator_ = new tanh;
    else throw std::invalid_argument("[gctl::avgpooling] Invalid activation type.");

    // 计算输出大小
    o_rows_ = (in_rows + s_rows_ - p_rows_)/s_rows_;
    o_cols_ = (in_cols + s_cols_ - p_cols_)/s_cols_;
    if (pl_type == Same && (in_rows + s_rows_ - p_rows_)%s_rows_ != 0)
    {
        o_rows_++;
        u_pad_ = ((in_rows + s_rows_ - p_rows_)%s_rows_)/2;
    }
    else u_pad_ = 0;

    if (pl_type == Same && (in_cols + s_cols_ - p_cols_)%s_cols_ != 0)
    {
        o_cols_++;
        l_pad_ = ((in_cols + s_cols_ - p_cols_)%s_cols_)/2;
    }
    else l_pad_ = 0;

    w_is_ = i_rows_*i_cols_;
    w_outs_ = o_rows_*o_cols_;
    p_idx_.resize(p_rows_*p_cols_);
    return;
}

void gctl::avgpooling::forward_propagation(const array<double> &all_weights, const matrix<double> &prev_layer_data)
{
    // z_: out_size x nobs
    // a_: out_size x nobs
    o_is_ = prev_layer_data.col_size();
    // Forward linear terms
    z_.resize(w_outs_, o_is_);
    a_.resize(w_outs_, o_is_);

    double p_val;
    for (size_t i = 0; i < o_rows_; i++)
    {
        for (size_t j = 0; j < o_cols_; j++)
        {
            if (p_type_ == Valid) cal_valid_padding_idx(p_idx_, i, j, i_rows_, i_cols_, p_rows_, p_cols_, s_rows_, s_cols_);
            if (p_type_ == Same) cal_same_padding_idx(p_idx_, i, j, i_rows_, i_cols_, p_rows_, p_cols_, s_rows_, s_cols_, u_pad_, l_pad_);

            for (size_t o = 0; o < o_is_; o++)
            {
                p_val = 0.0;
                for (size_t p = 0; p < p_idx_.size(); p++)
                {
                    if (p_idx_[p] != -1) p_val += prev_layer_data[p_idx_[p]][o];
                }

                z_[i*o_cols_ + j][o] = p_val/p_idx_.size();
            }
        }
    }

    // Apply activation function
    activator_->activate(z_, a_);
    return;
}

void gctl::avgpooling::backward_propagation(const array<double> &all_weights, const array<double> &all_ders, 
    const matrix<double> &prev_layer_data, const matrix<double> &next_layer_data)
{
    der_z_.resize(w_outs_, o_is_);
    der_in_.resize(w_is_, o_is_, 0.0);
    // prev_layer_data: in_size x nobs
    // next_layer_data: out_size x nobs
    // After forward stage, m_z contains z = max_avgpooling(in)
    // Now we need to calculate d(L) / d(z) = [d(a) / d(z)] * [d(L) / d(a)]
    // d(L) / d(z) is computed in the next layer, contained in next_layer_data
    // The Jacobian matrix J = d(a) / d(z) is determined by the activation function
    // der_z_: out_size x nobs   
    activator_->apply_jacobian(z_, a_, next_layer_data, der_z_);

    for (size_t j = 0; j < o_is_; j++)
    {
        for (size_t r = 0; r < o_rows_; r++)
        {
            for (size_t c = 0; c < o_cols_; c++)
            {
                if (p_type_ == Valid) cal_valid_padding_idx(p_idx_, r, c, i_rows_, i_cols_, p_rows_, p_cols_, s_rows_, s_cols_);
                if (p_type_ == Same) cal_same_padding_idx(p_idx_, r, c, i_rows_, i_cols_, p_rows_, p_cols_, s_rows_, s_cols_, u_pad_, l_pad_);

                for (size_t p = 0; p < p_rows_*p_cols_; p++)
                {
                    if (p_idx_[p] != -1) der_in_[p_idx_[p]][j] += der_z_[r*o_cols_ + c][j]/(p_rows_*p_cols_);
                }
            }
        }
    }
    return;
}

gctl::hlayer_type_e gctl::avgpooling::get_layer_type() const
{
    return AvgPooling;
}

std::string gctl::avgpooling::get_layer_name() const
{
    return "AvgPooling";
}

std::string gctl::avgpooling::layer_info() const
{
    std::string info = std::to_string(w_is_) + "x" + std::to_string(w_outs_)
        + " (" + std::to_string(i_rows_) + "," + std::to_string(i_cols_) + ")x"
        + "(" + std::to_string(o_rows_) + "," + std::to_string(o_cols_) + "), "
        + std::to_string(p_rows_) + "x" + std::to_string(p_cols_)
        + ", " + std::to_string(s_rows_) + "x" + std::to_string(s_cols_)
        + ", AvgPooling ";

    if (p_type_ == Same) info += "(Same), ";
    if (p_type_ == Valid) info += "(Valid), ";

    info += activator_->activation_name();
    return info;
}

void gctl::avgpooling::save_layer_setup(std::ofstream &os) const
{
    hlayer_type_e h_type = get_layer_type();
    activation_type_e a_type = get_activation_type();

    os.write((char*)&i_rows_, sizeof(int));
    os.write((char*)&i_cols_, sizeof(int));
    os.write((char*)&p_rows_, sizeof(int));
    os.write((char*)&p_cols_, sizeof(int));
    os.write((char*)&s_rows_, sizeof(int));
    os.write((char*)&s_cols_, sizeof(int));
    os.write((char*)&h_type,  sizeof(hlayer_type_e));
    os.write((char*)&p_type_, sizeof(pad_type_e));
    os.write((char*)&a_type,  sizeof(activation_type_e));
    return;
}

void gctl::avgpooling::load_layer_setup(std::ifstream &is)
{
    int in_rows, in_cols, pool_rows, pool_cols, stride_rows, stride_cols;
    hlayer_type_e h_type;
    pad_type_e p_type;
    activation_type_e a_type;

    is.read((char*)&in_rows, sizeof(int));
    is.read((char*)&in_cols, sizeof(int));
    is.read((char*)&pool_rows, sizeof(int));
    is.read((char*)&pool_cols, sizeof(int));
    is.read((char*)&stride_rows, sizeof(int));
    is.read((char*)&stride_cols, sizeof(int));
    is.read((char*)&h_type, sizeof(hlayer_type_e));
    is.read((char*)&p_type, sizeof(pad_type_e));
    is.read((char*)&a_type, sizeof(activation_type_e));

    if (h_type != AvgPooling)
    {
        throw std::invalid_argument("[gctl::avgpooling] Invalid hind layer type.");
    }

    init_avgpooling(in_rows, in_cols, pool_rows, pool_cols, stride_rows, stride_cols, p_type, a_type);
    return;
}

void gctl::avgpooling::save_weights2text(const array<double> &all_weights, std::ofstream &os) const
{
    return;
}