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

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

#include "hark-lib/extlib/fft.hpp"

namespace hark
{
namespace node
{

    EstimateLeak::EstimateLeak():
        NodeDef("EstimateLeak")
    {
        *this
            << Terminal::in("INPUT_POWER_SPEC")
            << Terminal::in("LEAK_FACTOR").def(0.25f)
            << Terminal::in("OVER_CANCEL_FACTOR").def(1.0f)
            >> Terminal::out("LEAK_POWER_SPEC");
    }
 
    void
    EstimateLeak::reset()
    {
        this->leak_factor = 0.25;
        this->over_cancel_factor = 1.0f;
        this->alpha = 0.25f;
    }
 
    void
    EstimateLeak::initialize(InputDataBundle &input)
    {
        this->leak_factor = input["LEAK_FACTOR"].last().as<float>();
        this->over_cancel_factor = input["OVER_CANCEL_FACTOR"].last().as<float>();

        this->alpha = this->leak_factor * this->over_cancel_factor;
    }

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

        const std::map<int, data> &inputMap = input["INPUT_POWER_SPEC"].last().cast<std::map<int, data> >();
        data outputMapData = data::create<std::map<int, data> >();
        std::map<int, data> &outputMap = outputMapData.cast<std::map<int, data> >();

        if(!inputMap.empty()){
            std::vector<float> vec_sum_spec;
            std::map<int, std::vector<float> > leak;
            for(auto it: inputMap){
                const std::vector<float> &spec = it.second.cast<std::vector<float> >();

                if(vec_sum_spec.size() < spec.size()){
                    vec_sum_spec.resize(spec.size());
                }
                for(size_t i = 0; i < spec.size(); i++){
                    vec_sum_spec[i] += spec[i];
                }

                std::vector<float> vec_leak(spec);
                leak[it.first] = std::move(vec_leak);
            }
            for(auto it: leak){
                for(size_t i = 0; i < it.second.size(); i++){
                    it.second[i] *= -1.0f;
                    it.second[i] += vec_sum_spec[i];
                }
            }
            float _beta = (-1.0f * this->alpha / (1.0f - this->alpha * this->alpha + this->alpha *  (1.0f - this->alpha) * static_cast<float>(inputMap.size() - 2)));
            float _alpha = 1.0f - static_cast<float>(inputMap.size() - 1) * this->alpha * _beta;
            std::map<int, std::vector<float> > z;
            std::vector<float> vec_sum_z;
            for(auto it: inputMap){
                const std::vector<float> &spec = it.second.cast<std::vector<float> >();

                std::map<int, std::vector<float> >::iterator leak_it = leak.find(it.first);
                if(leak_it != leak.end()){
                    std::vector<float> vec_z(spec.size());
                    for(size_t i = 0; i < spec.size(); i++){
                        vec_z[i] = _alpha * spec[i] + _beta * leak_it->second[i];
                        if(vec_z[i] < 1.0f){ vec_z[i] = 1.0f; }
                    }

                    if(vec_sum_z.size() < spec.size()){
                        vec_sum_z.resize(spec.size());
                    }
                    for(size_t i = 0; i < spec.size(); i++){
                        vec_sum_z[i] += vec_z[i];
                    }

                    z[it.first] = std::move(vec_z);
                }
            }
            for(auto it: z){
                const std::vector<float> &vec_z = it.second;

                data outData = data::create<std::vector<float> >(vec_z.size());
                std::vector<float> &out = outData.cast<std::vector<float> >();

                for(size_t i = 0; i < vec_z.size(); i++){
                    out[i] = this->alpha * (vec_sum_z[i] - vec_z[i]);
                }
                outputMap[it.first] = outData;
            }
        }

        output["LEAK_POWER_SPEC"].push(outputMapData);
        return true;
    }
}
}

/* end of file */
