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

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

namespace hark
{
namespace node
{
    Delta::Delta():
        NodeDef("Delta")
    {
        *this
            << Terminal::in("INPUT")
            << Terminal::in("FBANK_COUNT").def(41)
            >> Terminal::out("OUTPUT");
    }

    void
    Delta::reset()
    {
        //// retrieve parameters
        int fbank_count = 0;
        int delta_window = 2;

        //// setup values
        this->_fbank_count = fbank_count;
        this->_delta_window = delta_window;
        this->_A1.resize(fbank_count);
        this->_A2.resize(fbank_count);
        this->_sum.resize(fbank_count);
        this->_B = delta_window * (delta_window + 1) * (2 * delta_window + 1) / 6; // sum of 0^2, 1^2, 2^2, ..., delta_window^2
    }

    void
    Delta::initialize(InputDataBundle &input)
    {
        //// retrieve parameters
        int fbank_count = input["FBANK_COUNT"].last().cast<int>();
        int delta_window = 2;

        //// setup values
        this->_fbank_count = fbank_count;
        this->_delta_window = delta_window;
        this->_A1.resize(fbank_count);
        this->_A2.resize(fbank_count);
        this->_sum.resize(fbank_count);
        this->_B = delta_window * (delta_window + 1) * (2 * delta_window + 1) / 6; // sum of 0^2, 1^2, 2^2, ..., delta_window^2

        input["INPUT"].set_look_window(delta_window, delta_window, 0);
    }

    bool
    Delta::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        if(count == 0){
            this->initialize(input);
        }
        const int w = this->_delta_window;

        if(!input["INPUT"].is_source_ready()){
            return false;
        }

        // std::cout << "Delta: input: count=" << count << ", input=" << input["INPUT"].at(count) << std::endl;

        std::vector<std::map<int, data>>input_maps(2 * w + 1);
        input_maps[w] = input["INPUT"].at(count).cast<std::map<int, data>>();
        for(int i = 1; i <= w; i++){
            // std::cout << "c=" << (count + i) << std::endl;
            input_maps[w + i] = input["INPUT"].at(count + i).cast<std::map<int, data>>();

            if(count >= i){
                // std::cout << "c=" << (count - i) << std::endl;
                input_maps[w - i] = input["INPUT"].at(count - i).cast<std::map<int, data>>();
            }else{
                input_maps[w - i] = input_maps[w - i + 1];
            }
        }

        std::vector<float> &a1 = this->_A1;
        std::vector<float> &a2 = this->_A2;
        std::vector<float> &sum = this->_sum;
        for(auto it: input_maps[w]){
            int id = it.first;
            std::vector<float> &fbank = it.second.cast<std::vector<float>>();

            a1 = a2 = fbank;
            sum.assign(sum.size(), .0f);

            for(int i = 1; i <= w; i++){
                {
                    std::map<int, data>::const_iterator it2;
                    for(int t = w - i; t <= w; t++){
                        if((it2 = input_maps[t].find(id)) != input_maps[t].end()){
                            break;
                        }
                    }

                    a1 = it2->second.cast<std::vector<float>>();
                }

                {
                    std::map<int, data>::const_iterator it2;
                    for(int t = w + i; t >= w; t--){
                        if((it2 = input_maps[t].find(id)) != input_maps[t].end()){
                            break;
                        }
                    }

                    a2 = it2->second.cast<std::vector<float>>();
                }

                for(int n = 0; n < this->_fbank_count; n++){
                    sum[n] += i * (a2[n] - a1[n]);
                }
            }

            for(int n = 0; n < this->_fbank_count; n++){
                fbank[this->_fbank_count + n] = sum[n] / (2 * this->_B);
            }
        }

        output["OUTPUT"].push(data(input_maps[w]));
        // std::cout << "Delta: count=" << count << ", output=" << input_maps[w] << std::endl;
        return true;
    }
}
}

/* end of file */
