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

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

#include "hark-lib/source.h"
#include "hark-lib/matrix.h"

#include <math.h>

namespace hark
{
namespace node
{
    void
    VADZC::ZeroCrossCounter::reset(int size, float threshold)
    {
        // std::cout << "ZeroCross: reset: size=" << size << ", threshold=" << threshold << std::endl;
        //this->buffer.clear();
        //this->hasPower.clear();
        //this->hasZeroCross.clear();
        this->buffer = std::vector<float>(size, 0);
        this->hasPower = std::vector<bool>(size, 0);
        this->hasZeroCross = std::vector<bool>(size, 0);
        this->size = size;
        this->offset = 0;
        this->threshold = threshold;
        this->zeroCrossCount = 0;
    }

    int
    VADZC::ZeroCrossCounter::update(const float *data, int count)
    {
        // std::cout << "ZeroCross: update: ";
        // for(int i = 0; i < count; i++){ std::cout << data[i] << ", "; }
        // std::cout << std::endl;

        int i;
        for(i = 0; i < count; i++){
            int po = (this->offset == 0)? (this->size - 1): (this->offset - 1);
            this->buffer[this->offset] = data[i];
            this->hasPower[this->offset] = (fabs(this->buffer[this->offset]) > this->threshold);

            if(this->hasZeroCross[this->offset]){
                // std::cout << "zeroCrossCount--" << std::endl;
                this->zeroCrossCount--;
            }
            this->hasZeroCross[this->offset] = (this->hasPower[po]
                                                && (this->buffer[po] * this->buffer[this->offset]) < 0);
            if(this->hasZeroCross[this->offset]){
                // std::cout << "zeroCrossCount++" << std::endl;
                this->zeroCrossCount++;
            }
            if(++this->offset >= this->size){
                this->offset = 0;
            }
        }

        return this->zeroCrossCount;
    }

    int
    VADZC::ZeroCrossCounter::getZeroCrossCount() const
    {
        return this->zeroCrossCount;
    }

    VADZC::VADZC():
        NodeDef("VADZC"),
        _channel_count(0),
        _advance(0),
        _length(0),
        _split_period(0),
        _sampling_rate(0),
        _strip_zero(false),
        _zero_mean(false),
        _level_threshold(0),
        _zerocross_threshold(0),
        _head_margin(0),
        _tail_margin(0),
        _zerocross_threshold_count(0),
        _source_count(0),
        _head_margin_frames(0),
        _tail_margin_frames(0),
        _last_frame_index(0),
        _channel_buffer_length(0),
        _channel_buffer_offset(0),
        _channel_sources(),
        _count(0)
    {
        *this
            << Terminal::in("INPUT")
            << Terminal::in("ADVANCE").def(160)
            << Terminal::in("LENGTH").def(512)
            << Terminal::in("CHANNEL_COUNT").def(0)
            << Terminal::in("SAMPLING_RATE").def(16000)
            << Terminal::in("STRIP_ZERO").def(true)
            << Terminal::in("ZERO_MEAN").def(false)
            << Terminal::in("LEVEL_THRESHOLD").def(2000)
            << Terminal::in("ZEROCROSS_THRESHOLD").def(60)
            << Terminal::in("HEAD_MARGIN").def(300)
            << Terminal::in("TAIL_MARGIN").def(400)
            << Terminal::in("SPLIT_PERIOD").def(0)
            >> Terminal::out("OUTPUT") /* Map<int,ObjectRef> */
            >> Terminal::out("SOURCES"); /* Vector<ObjectRef> */
    }

    void
    VADZC::reset()
    {
        this->_channel_count = 0;
        this->_advance = 0;
        this->_length = 0;
        this->_split_period = 0;
        this->_sampling_rate = 0;
        this->_strip_zero = false;
        this->_zero_mean = false;
        this->_level_threshold = 0;
        this->_zerocross_threshold = 0;
        this->_head_margin = 0;
        this->_tail_margin = 0;
        this->_zerocross_threshold_count = 0;
        this->_source_count = 0;
        this->_split_period_frames = 0;
        this->_head_margin_frames = 0;
        this->_tail_margin_frames = 0;
        this->_last_frame_index = 0;
        this->_channel_buffer_length = 0;
        this->_channel_buffer_offset = 0;
        this->_channel_sources.clear();
        this->_count = 0;
        this->_channel_worksets.clear();
    }

    void
    VADZC::update_parameters(InputDataBundle &input)
    {
        int channel_count = input["CHANNEL_COUNT"].last().as<int>();
        this->_advance = input["ADVANCE"].last().as<int>();
        this->_length = input["LENGTH"].last().as<int>();
        this->_split_period = input["SPLIT_PERIOD"].last().as<float>();
        this->_sampling_rate = input["SAMPLING_RATE"].last().as<int>();
        this->_strip_zero = input["STRIP_ZERO"].last().as<bool>();
        this->_zero_mean = input["ZERO_MEAN"].last().as<bool>();
        this->_level_threshold = input["LEVEL_THRESHOLD"].last().as<float>();
        this->_zerocross_threshold = input["ZEROCROSS_THRESHOLD"].last().as<int>();
        this->_head_margin = input["HEAD_MARGIN"].last().as<float>();
        this->_tail_margin = input["TAIL_MARGIN"].last().as<float>();

        this->_split_period_frames = (int)(this->_split_period * this->_sampling_rate / this->_advance / 1000);
        this->_head_margin_frames = (int)(this->_head_margin * this->_sampling_rate / this->_advance / 1000);
        this->_tail_margin_frames = (int)(this->_tail_margin * this->_sampling_rate / this->_advance / 1000);
        if(this->_head_margin_frames == 0){
            this->_head_margin_frames = 1;
        }

        this->_channel_buffer_length = this->_advance * this->_head_margin_frames;
        this->_zerocross_threshold_count = this->_zerocross_threshold * this->_channel_buffer_length / this->_sampling_rate;

        if(this->_channel_count == 0){
            this->_channel_count = channel_count;
            this->_channel_worksets.resize(this->_channel_count);
            for(int i = 0; i < this->_channel_count; i++){
                this->_channel_worksets[i].buffer.resize(this->_head_margin_frames);
                this->_channel_worksets[i].zerocross_counter.reset(this->_channel_buffer_length, this->_level_threshold);
                for(int j = 0; j < this->_head_margin_frames; j++){
                    this->_channel_worksets[i].buffer[j].resize(this->_length);
                }
                this->_channel_worksets[i].last_active_frame_index = -this->_tail_margin_frames;
                this->_channel_worksets[i].dc = 0;
                this->_channel_worksets[i].dc_prepared = false;
            }

        }else if(channel_count != this->_channel_count){
            std::stringstream ss;
            ss << "VADZC: Channel count mismatch (input=" << channel_count << ", expected=" << this->_channel_count << ")";
            throw RuntimeError(ss.str());
        }

        input["INPUT"].set_look_window(this->_head_margin_frames, this->_head_margin_frames);
    }

    void
    VADZC::receive(InputDataBundle &input, int count)
    {
        while(this->_last_frame_index < count){
            // std::cout << "receive: " << this->_last_frame_index << "/" << count << std::endl;
            /* cursor in channel_buffers */
            int buffer_cursor = this->_last_frame_index % this->_head_margin_frames;

            /* receive input object */
//            data in = getInput(this->inputID, this->last_frame_index);

            /* input entity is Matrix<float> type. */
            data in = input["INPUT"].at(this->_last_frame_index);
            const Matrix<float> &x = in.cast<Matrix<float>>();
            // std::cout << "VADZC: count=" << this->_last_frame_index << ", input=" << in << std::endl;

            /* check inter-frame cosistency */
            if(this->_channel_count != x.rows()){
                std::stringstream ss;
                ss << "VADZC: Channel count mismatch (input=" << x.rows() << ", expected=" << this->_channel_count << ")";
                throw RuntimeError(ss.str());
            }
            if(this->_length > 0 && this->_length != x.cols()){
                std::stringstream ss;
                ss << "VADZC: Buffer length mismatch (input=" << x.cols() << ", expected=" << this->_length << ")";
                throw RuntimeError(ss.str());
            }
            else if(this->_length == 0){
                this->_length = x.cols();
            }

            /* append to buffer */
            int channel_index;
            for(channel_index = 0; channel_index < this->_channel_count; channel_index++){
                ChannelWorkset &channel_workset = this->_channel_worksets[channel_index];

                if(count == 0 && this->_zero_mean && !channel_workset.dc_prepared){
                    float dc = 0;
                    for(int k = 0; k < this->_length; k++){
                        dc += x(channel_index, k);
                    }
                    channel_workset.dc = dc / this->_length;
                    channel_workset.dc_prepared = true;
                }

                /* update buffer */
                for(int i = 0; i < this->_length; i++){
                    channel_workset.buffer[buffer_cursor][i] = x(channel_index, i) - channel_workset.dc;
                }

                /* update zero-cross counter */
                channel_workset.zerocross_counter.update(x.row(channel_index).data() + (this->_length - this->_advance),
                                                         this->_advance);
                // std::cout << "VADZC: channel-workset: dc=" << channel_workset.dc << ", buffer=" << channel_workset.buffer << ", zc=" << channel_workset.zerocross_counter.hasZeroCross << std::endl;
                // std::cout << "VADZC: channel-workset: dc=" << channel_workset.dc << ", zc=" << channel_workset.zerocross_counter.hasZeroCross << std::endl;
            }

            /* update cursor */
            this->_last_frame_index++;
        }
    }

    bool
    VADZC::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        /* update parameters */
        this->update_parameters(input);
        if(!input["INPUT"].is_source_ready()){
            return false; // false;
        }

        {// for debug
            data in = input["INPUT"].at(this->_last_frame_index);
            const Matrix<float> &x = in.cast<Matrix<float>>();

//            std::cout << "VADZC: norm=" << x.norm() << std::endl;
        }

        /* 先読みが必要なので、head_margin_frames フレーム追加で読み込んでおく */
        this->receive(input, count + this->_head_margin_frames);

        /* prepare output data buffer */
        std::vector<data>sources;
        std::map<int, data>outputs;

        /* cache parameter */
        size_t channel_count = this->_channel_count;
        size_t length = this->_length;

        size_t channel_index;
        for(channel_index = 0; channel_index < channel_count; channel_index++){
            ChannelWorkset &channel_workset = this->_channel_worksets[channel_index];

            int zc = 0; // debug
            bool source_exists = false;

            // std::cout << "VADZC: channel_index=" << channel_index << ", zc=" << channel_workset.zerocross_counter.getZeroCrossCount() << std::endl;
            if((zc = channel_workset.zerocross_counter.getZeroCrossCount()) > this->_zerocross_threshold_count){
                // output frames
                channel_workset.last_active_frame_index = count;
                source_exists = true;

            }
            else if(count < channel_workset.last_active_frame_index + this->_tail_margin_frames){
                // output frames
                source_exists = true;

            }
            else{
                // no need to output
            }

            // std::cout << "source-exists[" << channel_index << "]: " << source_exists << ", zc=" << zc << "/" << this->zerocross_threshold_count << std::endl; //debug

            std::map<int, Source>::iterator it = this->_channel_sources.find(channel_index);
            if(source_exists){
                Source new_source;

                if(it != this->_channel_sources.end()){
                    new_source = it->second;
                }
                else{
                    int source_id = ++this->_source_count;
                    new_source.x[0] = 0;
                    new_source.x[1] = 0;
                    new_source.x[2] = 0;
                    new_source.power = 0;
                    new_source.remaining = 0;
                    new_source.id = source_id;
                    new_source.mic_id = channel_index;
                    new_source.sub_id = 0;
                    new_source.first_found_count = count;
                    this->_channel_sources[channel_index] = new_source;
                }
                new_source.count = count;
                new_source.last_found_count = count;
                if(this->_split_period > 0){
                    new_source.sub_id = (count - new_source.first_found_count) / this->_split_period_frames;
                }

                /* target data to output */
                sources.push_back(data(new_source));
                outputs.emplace(new_source.id, channel_workset.buffer[count % this->_head_margin_frames]);
            }
            else{
                if(it != this->_channel_sources.end()){
                    this->_channel_sources.erase(it);
                }
            }
        }

        output["SOURCES"].push(data(std::move(sources)));
        output["OUTPUT"].push(data(std::move(outputs)));
        return true;
    }
}
}

/* end of file */
