/*
   HARK Library
   Copyright (C) 2021 Honda Research Institute Japan Inc.
 */

#include "hark-lib/node/BGNEstimator.h"

namespace hark
{
namespace node
{

    BGNEstimator::BGNEstimator():
        NodeDef("BGNEstimator")
    {
        *this
            << Terminal::in("INPUT_POWER")
            << Terminal::in("DELTA").def(3.0f)
            << Terminal::in("L").def(150)
            << Terminal::in("ALPHA_S").def(0.7f)
            << Terminal::in("NOISE_COMPENSE").def(1.0f)
            << Terminal::in("ALPHA_D_MIN").def(0.05f)
            << Terminal::in("NUM_INIT_FRAME").def(100)
            >> Terminal::out("NOISE_POWER");
    }
 
    void
    BGNEstimator::reset()
    {
    }
 
    void
    BGNEstimator::initialize(InputDataBundle &input)
    {
        this->delta = input["DELTA"].last().as<float>();
        this->l = input["L"].last().as<int>();
        this->alpha_s = input["ALPHA_S"].last().as<float>();
        this->noise_compense = input["NOISE_COMPENSE"].last().as<float>();
        this->alpha_d_min = input["ALPHA_D_MIN"].last().as<float>();
        this->num_init_frame = input["NUM_INIT_FRAME"].last().as<int>();

        this->num_proc = 0;
    }

    bool
    BGNEstimator::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        if(count == 0){ this->initialize(input); }

        const auto &input_spec = input["INPUT_POWER"].last().cast<Matrix<float> >();

        if(count == 0){
            this->S = Matrix<float>(input_spec.rows(), input_spec.cols());
            this->Sf = Matrix<float>(input_spec.rows(), input_spec.cols());
            this->Smin = Matrix<float>(input_spec.rows(), input_spec.cols());
            this->Stmp = Matrix<float>(input_spec.rows(), input_spec.cols());
            this->Lambda = Matrix<float>(input_spec.rows(), input_spec.cols());

            //this->S.resize(input_spec.rows(), input_spec.cols());
            //this->S.setZero();
            //this->Sf.resize(input_spec.rows(), input_spec.cols());
            //this->Sf.setZero();
            //this->Smin.resize(input_spec.rows(), input_spec.cols());
            //this->Smin.setZero();
            //this->Stmp.resize(input_spec.rows(), input_spec.cols());
            //this->Stmp.setZero();
            //this->Lambda.resize(input_spec.rows(), input_spec.cols());
            //this->Lambda.setZero();
        }

        data out = data::create<Matrix<float> >(input_spec.rows(), input_spec.cols());
        Matrix<float> &output_spec = out.cast<Matrix<float> >();

        // Smoothing along the time axis
        for(size_t i = 0; i < input_spec.rows(); i++){
            for(size_t j = 0; j < input_spec.cols(); j++){
                float tmp = this->alpha_s * this->S(i, j)
                          + (1.0f - this->alpha_s) * input_spec(i, j);
                this->S(i, j) = tmp;
            }
        }

        // Update coef.
        float alpha_d = 1.0f / static_cast<float>(this->num_proc + 1);
        if(alpha_d < this->alpha_d_min){ alpha_d = this->alpha_d_min; }

        // Smoothing along the freq. axis
        size_t jmax = input_spec.cols() - 1;
        for(size_t i = 0; i < input_spec.rows(); i++){
            this->Sf(i, 0) = 0.5f * this->S(i, 0) + 0.25f * this->S(i, 1);
            for(size_t j = 1; j < jmax; j++){
                this->Sf(i, j) = 0.25f * this->S(i, j - 1) + 0.5f * this->S(i, j) + 0.25f * this->S(i, j + 1);
            }
            this->Sf(i, jmax) = 0.25f * this->S(i, jmax - 1) + 0.5f * this->S(i, jmax);
        }

        // Update minimum value
        if((this->num_proc % this->l) == 0){
            for(size_t i = 0; i < input_spec.rows(); i++){
                for(size_t j = 0; j < input_spec.cols(); j++){
                    this->Smin(i, j) = std::min(this->Stmp(i, j), this->Sf(i, j));
                    this->Stmp(i, j) = this->Sf(i, j);
                }
            }
        }
        else{
            for(size_t i = 0; i < input_spec.rows(); i++){
                for(size_t j = 0; j < input_spec.cols(); j++){
                    this->Smin(i, j) = std::min(this->Smin(i, j), this->Sf(i, j));
                    this->Stmp(i, j) = std::min(this->Stmp(i, j), this->Sf(i, j));
                }
            }
        }

        // Noise estimation
        for(size_t i = 0; i < input_spec.rows(); i++){
            for(size_t j = 0; j < input_spec.cols(); j++){
                float alpha_dc = 0.0f;
                // 
                if((this->Sf(i, j) < this->Smin(i, j) * this->delta) ||
                   (this->num_proc < this->num_init_frame) ||
                   (this->Lambda(i, j) > this->S(i, j))){
                    alpha_dc = alpha_d;
                }
                output_spec(i, j) = (1.0f - alpha_dc) * this->Lambda(i, j)
                                  + alpha_dc * this->noise_compense * input_spec(i, j);
            }
        }
        this->num_proc++;

        output["NOISE_POWER"].push(out);
        return true;
    }
}
}

/* end of file */
