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

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

#include "hark-lib/matrix.h"
#include "hark-lib/data.h"
#include <map>

namespace hark
{
namespace node
{
    MultiFFT::MultiFFT():
        NodeDef("MultiFFT")
    {
        *this
            << Terminal::in("INPUT")
            << Terminal::in("LENGTH").def(512)
            << Terminal::in("WINDOW").def("CONJ")
            << Terminal::in("WINDOW_LENGTH").def(512)
            >> Terminal::out("OUTPUT");
    }

    void
    MultiFFT::reset()
    {
        this->_length = 0;
        this->_window_type = UNKNOWN_WINDOW;
        this->_pslength = 0;
        this->_window.clear();
        this->_rfft.initialize(0);
    }

    void
    MultiFFT::initialize_parameters(InputDataBundle &input)
    {
        // std::cout << "MultiFFT: initialize_parameters(): start" << std::endl;

        // retrive valuesAudioStreamFromMemory
        int length = input["LENGTH"].last().cast<int>();
        int window_length = input["WINDOW_LENGTH"].last().cast<int>();
        std::string window_name = input["WINDOW"].last().cast<std::string>();

        // std::cout << "MultiFFT: initialize_parameters(): length=" << length << ", window_length=" << window_length << ", window_name=" << window_name << std::endl;

        // validate/convert values
        //// window-type
        WINDOW_TYPE window_type = UNKNOWN_WINDOW;
        if(window_name == "CONJ"){
            window_type = CONJ_WINDOW;
        }
        else if(window_name == "RECTANGLE"){
            window_type = RECTANGLE_WINDOW;
        }
        else if(window_name == "HAMMING"){
            window_type = HAMMING_WINDOW;
        }
        else if(window_name == "HANNING"){
            window_type = HANNING_WINDOW;
        }
        else{
            throw RuntimeError("MultiFFT: WINDOW should be \"CONJ\", \"RECTANGLE\" or \"HAMMING\"");
        }

        //// window-length
        if(window_length > length){
            std::stringstream ss;
            ss << "MultiFFT: WINDOW_LENGTH(" << window_length << ") should be <= LENGTH(" << length << ")";
            throw RuntimeError(ss.str());
        }

        // set fields
        //// length
        this->_length = length;
        this->_pslength = length / 2 + 1;
        // int r = (int)(log2(length) + 0.5);

        //// window
        this->_window.resize(window_length);
        switch(window_type){
        case CONJ_WINDOW:
            ConjWindow(this->_window);
            break;
        case RECTANGLE_WINDOW:
            std::fill(this->_window.begin(), this->_window.end(), 1.0f);
            break;
        case HAMMING_WINDOW:
            Hamming(this->_window);
            break;
        case HANNING_WINDOW:
            Hanning(this->_window);
            break;
        }
        this->_window.resize(length);

        //// FFT engine
        this->_rfft.initialize(length);

    }

    bool
    MultiFFT::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        // std::cout << "MultiFFT: operator(): start: count=" << count << std::endl;
        if(count == 0){
            // std::cout << "MultiFFT: operator(): initialize" << std::endl;
            this->initialize_parameters(input);
            // std::cout << "MultiFFT: operator(): initialize: end" << std::endl;
        }

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

        // std::cout << "MultiFFT: operator(): input" << inputData << std::endl;

        if(inputData.is_type<Matrix<float>>()){
            // std::cout << "MultiFFT: operator(): input: Matrix<float>" << std::endl;

            const auto &inputMatrix = inputData.cast<Matrix<float>>();
            // std::cout << "MultiFFT: operator(): input: Matrix<float>: rows=" << inputMatrix.rows() << ", cols=" << inputMatrix.cols() << std::endl;

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

            size_t channel_count = inputMatrix.rows();
            if(inputMatrix.cols() != this->_length){
                std::stringstream ss;
                ss << "MultiFFT: FFT length (" << this->_length << ") and column number of input matrix (" << inputMatrix.cols() << ") mismatch";
                throw RuntimeError(ss.str());
            }

            auto &fft = this->_rfft;
            for(int m = 0; m < inputMatrix.rows(); m++){
                // std::cout << "MultiFFT: operator(): fft: row=" << m << std::endl;
                for(int i = 0; i < inputMatrix.cols(); i++){
                    fft.a[i] = inputMatrix(m, i) * this->_window[i];
                }
                fft.analyze();
                for(int k = 0; k < outputMatrix.cols(); k++){
                    outputMatrix(m, k) = fft.A[k];
                }
            }

        }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>>();

            // std::cout << "MultiFFT: count=" << count << ", input=" << inputMap << std::endl;

            auto &fft = this->_rfft;
            for(auto it: inputMap){
                const auto &x = it.second.cast<std::vector<float>>();
                data yData = data::create<std::vector<std::complex<float>>>(fft.flength(), .0f);
                auto &y = yData.cast<std::vector<std::complex<float>>>();

                for(int i = 0; i < x.size(); i++){
                    fft.a[i] = x[i] * this->_window[i];
                }
                fft.analyze();
                for(int k = 0; k < y.size(); k++){
                    y[k] = fft.A[k];
                }

                outputMap[it.first] = yData;
            }

        }else{
            std::stringstream ss;
            ss << "MultiFFT: Unexpected input data type";
            throw RuntimeError(ss.str());

        }

        // std::cout << "MultiFFT: operator(): output" << outputData << std::endl;
        output["OUTPUT"].push(outputData);

        // std::cout << "MultiFFT: operator(): end, count=" << count << std::endl;
        return true;
    }
}
}

/* end of file */
