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

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

namespace hark
{
namespace node
{
    enum WINDOW_TYPE {
        CONJ_WINDOW, RECTANGLE_WINDOW, HAMMING_WINDOW, HANNING_WINDOW
    };

    Synthesize::Synthesize():
        NodeDef("Synthesize")
    {
        *this
            << Terminal::in("INPUT")
            << Terminal::in("LENGTH").def(512)
            << Terminal::in("ADVANCE").def(160)
            << Terminal::in("SAMPLING_RATE").def(16000)
            << Terminal::in("MIN_FREQUENCY").def(125)
            << Terminal::in("MAX_FREQUENCY").def(7900)
            << Terminal::in("WINDOW").def("HAMMING")
            << Terminal::in("OUTPUT_GAIN").def(1.0)
            >> Terminal::out("OUTPUT");
    }

    void
    Synthesize::reset()
    {
    }

    void
    Synthesize::initialize(InputDataBundle &input)
    {
        this->_length = input["LENGTH"].last().as<int>();
        this->_advance = input["ADVANCE"].last().as<int>();
        this->_sampling_rate = input["SAMPLING_RATE"].last().as<int>();
        this->_output_gain = input["OUTPUT_GAIN"].last().as<float>();
        this->_min_freq = input["MIN_FREQUENCY"].last().as<int>();
        this->_max_freq = input["MAX_FREQUENCY"].last().as<int>();
        if(this->_min_freq < 0 || this->_max_freq > this->_sampling_rate / 2 || this->_min_freq > this->_max_freq){
            throw RuntimeError("SAMPLING_RATE, MIN_FREQUENCY or MAX_FREQUENCY is wrong");
        }

        this->_min_freq_ind = this->_length * this->_min_freq / this->_sampling_rate;
        this->_max_freq_ind = this->_length * this->_max_freq / this->_sampling_rate;
        this->_overlap_frame = this->_length / this->_advance - (this->_length % this->_advance ? 0 : 1);

        this->_wn_factor = sqrtf(this->_length) / 2.0f;

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

        this->_rfft1.initialize(this->_length);
        this->_wave_buffer_offset = 0;

        std::string window_name = input["WINDOW"].last().cast<std::string>();
        if(window_name == "CONJ"){
            this->_window_type = CONJ_WINDOW;
        }
        else if(window_name == "RECTANGLE"){
            this->_window_type = RECTANGLE_WINDOW;
        }
        else if(window_name == "HAMMING"){
            this->_window_type = HAMMING_WINDOW;
        }
        else if(window_name == "HANNING"){
            this->_window_type = HANNING_WINDOW;
        }
        else{
            throw RuntimeError("MultiFFT: WINDOW should be RECTANGLE, CONJ or HAMMING");
        }

        // this->_window.resize(this->_length);
        this->_wave_buffer_offset = 0;
        this->_window.resize(this->_length);
        switch(this->_window_type){
        case CONJ_WINDOW:
            ConjWindow(this->_window);
            break;

        case RECTANGLE_WINDOW:
            for (int i = 0; i < this->_window.size(); i++) {
                this->_window[i] = 1.0f;
            }
            break;

        case HAMMING_WINDOW:
            Hamming(this->_window);
            break;

        case HANNING_WINDOW:
            Hanning(this->_window);
            break;
        }
    }

    bool
    Synthesize::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        // std::cout << "(): count=" << count << std::endl;
        if(count == 0){
            this->initialize(input);
        }
        if(!input["INPUT"].is_source_ready()){
            return false;
        }

        //data inputData = input["INPUT"].last();
        data inputData = input["INPUT"].at(count);
        data outputData;

        if(inputData.is_type<Matrix<std::complex<float> > >()){
            const auto &inputMatrix = inputData.cast<Matrix<std::complex<float> > >();

            outputData = data::create<Matrix<float> >(inputMatrix.rows(), this->_length);
            auto &outputMatrix = outputData.cast<Matrix<float> >();

            if(count == 0){
                this->_wave_buffer_map.resize(0);
                this->_wave_buffer_matrix.resize(this->_overlap_frame + 1);
                this->_last_frame_matrix.resize(inputMatrix.rows(), this->_length);
                this->_last_frame_matrix.setZero();

                // Perform overlap add for first frame.
                for(size_t frame = 0; frame <= this->_overlap_frame; frame++){
                    this->SynthesizeFrameMatrix(frame, input);

                    size_t add_cols = this->_length - this->_advance * frame;
                    this->_last_frame_matrix.rightCols(add_cols) = this->_last_frame_matrix.rightCols(add_cols) +
                                                                   this->_wave_buffer_matrix.at(frame - this->_wave_buffer_offset).leftCols(add_cols);
                    // Is this->_wave_buffer_offset alway 0 in this for-loop?
                    //this->_last_frame_matrix.rightCols(add_cols) = this->_last_frame_matrix.rightCols(add_cols) +
                    //                                               this->_wave_buffer_matrix.at(frame).leftCols(add_cols);
                }
            }
            else{
                this->SynthesizeFrameMatrix(count + this->_overlap_frame, input);

                // Perform overlap add for last advance samples.
                for(size_t frame = 0; frame <= this->_overlap_frame; frame++){
                    int col_offset = this->_length - this->_advance * static_cast<int>(frame + 1);
                    if(col_offset < 0 || col_offset > (this->_length - this->_advance)){
                        // skip
                    }
                    else{
                        this->_last_frame_matrix.rightCols(this->_advance) = this->_last_frame_matrix.rightCols(this->_advance) +
                                this->_wave_buffer_matrix.at(count + frame - this->_wave_buffer_offset).middleCols(col_offset, this->_advance);
                        // Is (count - this->_wave_buffer_offset) alway 0 in this-loop?
                        //this->_last_frame_matrix.rightCols(this->_advance) = this->_last_frame_matrix.rightCols(this->_advance) +
                        //                                                     this->_wave_buffer_matrix.at(frame).middleCols(col_offset, this->_advance);
                    }
                }
            }

            outputMatrix = this->_last_frame_matrix;
            outputMatrix *= this->_output_gain;

            // Perform advance shift.
            this->_last_frame_matrix.leftCols(this->_length - this->_advance) = this->_last_frame_matrix.rightCols(this->_length - this->_advance);
            this->_last_frame_matrix.rightCols(this->_advance).setZero();
        }
        else if(inputData.is_type<std::map<int, data> >()){
            const auto &inputMap = inputData.cast<std::map<int, data> >();

            outputData = data::create<std::map<int, data> >();
            auto &outputMap = outputData.cast<std::map<int, data> >();

            if(count == 0){
                this->_wave_buffer_matrix.resize(0);
                this->_wave_buffer_map.resize(this->_overlap_frame + 1);
                for(size_t frame = 0; frame < this->_overlap_frame; frame++){
                    this->SynthesizeFrameMap(frame, input);

                    const auto &wave_map = this->_wave_buffer_map.at(frame - this->_wave_buffer_offset);
                    size_t begin = 0;
                    size_t end = this->_length - this->_advance * (frame + 1);
                    size_t lf_begin = this->_advance * frame;
                    size_t lf_end = this->_length - this->_advance;

                    for(auto &it: wave_map){
                        const auto &w = it.second.cast<std::vector<float> >();
                        std::vector<float> &lf = this->_last_frame[it.first];
                        lf.resize(this->_length, 0.0f);

                        for(size_t i = 0; i < end; i++) {
                            lf[lf_begin + i] += w[i];
                        }
                    }
                }
            }
            SynthesizeFrameMap(count + this->_overlap_frame, input);

            // Perform overlap add.
            for(size_t frame = 0; frame <= this->_overlap_frame; frame++) {
                const auto &wave_map = this->_wave_buffer_map[count + frame - this->_wave_buffer_offset];

                size_t begin = static_cast<int>(std::max(static_cast<int>(this->_length - this->_advance * (frame + 1)), static_cast<int>(0)));
                size_t end = this->_length - this->_advance * frame;
                size_t lf_begin = this->_length - (end - begin);

                for(auto it: wave_map){
                    const auto &w = it.second.cast<std::vector<float>>();
                    std::vector<float> &lf = this->_last_frame[it.first];
                    lf.resize(this->_length, 0.0f);

                    for(size_t i = 0; i < end - begin; i++) {
                        lf[lf_begin + i] += w[begin + i];
                    }
                }
            }

            // Remove sound sources which have gone from the last_frame map.
            std::map<int, std::vector<float> >::iterator it = this->_last_frame.begin();
            while(it != this->_last_frame.end()) {
                bool srcid_exists = false;
                for(size_t frame = 0; frame <= this->_overlap_frame; frame++){
                    const auto &wave_map = this->_wave_buffer_map[count + frame - this->_wave_buffer_offset];
                    if(wave_map.find(it->first) != wave_map.end()) {
                        srcid_exists = true;
                        break;
                    }
                }
                if(!srcid_exists){
                    it = this->_last_frame.erase(it);
                }
                else{
                    ++it;
                }
            }
    
            // for(std::map<int, std::vector<float>>::iterator it_lf = last_frame.begin(); it_lf != last_frame.end(); ++it_lf) {
            for(auto &it_lf: this->_last_frame){
                bool to_output = true;
                auto &spec_map = input["INPUT"].last().cast<std::map<int, data> >();
                if(spec_map.find(it_lf.first) == spec_map.end()) {
                    to_output = false;
                }

                if(to_output){
                    data xdata = data::create<std::vector<float> >(it_lf.second.size());
                    auto &x = xdata.cast<std::vector<float> >();
                    std::copy(it_lf.second.begin(), it_lf.second.end(), x.begin());
                    for(size_t i = 0; i < x.size(); i++) {
                        x[i] *= this->_output_gain;
                    }
                    outputMap[it_lf.first] = xdata;
                }

                for(size_t i = 0; i < this->_length - this->_advance; i++) {
                    it_lf.second[i] = it_lf.second[i + this->_advance];
                }
                for(size_t i = this->_length - this->_advance; i < this->_length; i++) {
                    it_lf.second[i] = 0.0f;
                }
            }
        }
        else{
            std::stringstream ss;
            ss << "Synthesize: Unexpected input data type.";
            throw RuntimeError(ss.str());
        }

        // std::cout << "output: " << outputData;
        output["OUTPUT"].push(outputData);
        // std::cout << "Synthesize: output: count=" << count << std::endl;

        return true;
    }

    void
    Synthesize::SynthesizeFrameMatrix(int count, InputDataBundle &input)
    {
        data inputData = input["INPUT"].at(count);
        const auto &spec_matrix = inputData.cast<Matrix<std::complex<float> > >();
        Matrix<float> wave_matrix(spec_matrix.rows(), this->_length);

        for(size_t r = 0; r < spec_matrix.rows(); r++){
            const auto &spec = spec_matrix.row(r);
            std::vector<float> wave(this->_length);

            for(size_t k = 0; k < this->_rfft1.flength(); k++) {
                if (k < this->_min_freq_ind || k > this->_max_freq_ind) {
                    //fft.complex_data[k] = 0.0f;
                    this->_rfft1.A[k] = 0.0f;
                }
                else {
                    //fft.complex_data[k] = spec[k];
                    this->_rfft1.A[k] = spec[k];
                }
            }
            //fft.CalcBackward();
            this->_rfft1.synthesize();

            for(size_t i = 0; i < this->_length; i++) {
                //wave[i] = fft.real_data[i] * (1.0f / fft.length) * window[i];
                //wave[i] = fft.real_data[i] * (1.0f / fft.length) * hamming[i];
                wave_matrix(r, i) = this->_rfft1.a[i] * (1.0f / this->_rfft1.tlength()) * this->_window[i];
            }
        }

        if(this->_wave_buffer_offset + this->_wave_buffer_matrix.size() <= count){
            this->_wave_buffer_matrix.emplace_back(wave_matrix);
            this->_wave_buffer_matrix.pop_front();
            this->_wave_buffer_offset++;
        }
        else{
            this->_wave_buffer_matrix[count - this->_wave_buffer_offset] = wave_matrix;
        }
    }

    void
    Synthesize::SynthesizeFrameMap(int count, InputDataBundle &input)
    {
        data inputData = input["INPUT"].at(count);
        const auto &spec_map = inputData.cast<std::map<int, data> >();
        std::map<int, data> wave_map;

        for(auto &it: spec_map){
            const auto &spec = it.second.cast<std::vector<std::complex<float> > >();
            data wave_data = data::create<std::vector<float> >(this->_length);
            auto &wave = wave_data.cast<std::vector<float> >();

            for(size_t k = 0; k < this->_rfft1.flength(); k++) {
                if (k < this->_min_freq_ind || k > this->_max_freq_ind) {
                    //fft.complex_data[k] = 0.0f;
                    this->_rfft1.A[k] = 0.0f;
                }

                else {
                    //fft.complex_data[k] = spec[k];
                    this->_rfft1.A[k] = spec[k];
                }
            }
            //fft.CalcBackward();
            this->_rfft1.synthesize();

            for(size_t i = 0; i < this->_length; i++) {
                //wave[i] = fft.real_data[i] * (1.0f / fft.length) * window[i];
                //wave[i] = fft.real_data[i] * (1.0f / fft.length) * hamming[i];
                wave[i] = this->_rfft1.a[i] * (1.0f / this->_rfft1.tlength()) * this->_window[i];
            }
            wave_map[it.first] = wave_data;
        }

        if(this->_wave_buffer_offset + this->_wave_buffer_map.size() <= count){
            // std::cout << "wave_buffer: count=" << count << ", emplace_back" << std::endl;
            this->_wave_buffer_map.emplace_back(wave_map);
            this->_wave_buffer_map.pop_front();
            this->_wave_buffer_offset++;
        }
        else{
            // std::cout << "wave_buffer: count=" << count << ", replace: " << (count - this->_wave_buffer_offset) << std::endl;
            this->_wave_buffer_map[count - this->_wave_buffer_offset] = wave_map;
        }
        // while(this->_wave_buffer_map.size() > this->_length){
        //     std::cout << "wave_buffer: count=" << count << ", rotate" << std::endl;
        // }
    }

}
}

/* end of file */
