
#include <cmath>
#include <complex>
#include <algorithm>
#include <iostream>

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

namespace hark
{
namespace node
{
    CMMakerFromFFT::CMMakerFromFFT():
        NodeDef("CMMakerFromFFT")
    {
        *this
            << Terminal::in("INPUT")
            << Terminal::in("WINDOW").def(50)
            << Terminal::in("PERIOD").def(50)
            << Terminal::in("WINDOW_TYPE").def("FUTURE")
            << Terminal::in("ENABLE_DEBUG").def(false)
            >> Terminal::out("OPERATION_FLAG")
            >> Terminal::out("OUTPUT");
    }

    void
    CMMakerFromFFT::reset()
    {
    }

    void
    CMMakerFromFFT::initialize(InputDataBundle &input)
    {
        const auto &spectrum = input["INPUT"].last().cast<Matrix<std::complex<float> > >();

        this->window = input["WINDOW"].last().as<int>();
        this->period = input["PERIOD"].last().as<int>();
        this->window_type = input["WINDOW_TYPE"].last().cast<std::string>();
        this->enable_debug = input["ENABLE_DEBUG"].last().as<bool>();
        if(this->window_type == "FUTURE"){
            this->look_ahead = this->window -1;
            this->look_back = 0;
        }
        else if(this->window_type == "MIDDLE"){
            this->look_ahead = this->window / 2;
            this->look_back = this->window / 2 + (this->window % 2) - 1;
        }
        else if(this->window_type == "PAST"){
            this->look_ahead = 0;
            this->look_back = 0;
        }
        else{
            std::stringstream ss;
            ss << "CMMakerFromFFT: WINDOW_TYPE should be \"FUTURE\", \"MIDDLE\" or \"PAST\".\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
            throw RuntimeError(ss.str());
        }

        this->pslength = spectrum.cols();
        this->nb_channels = spectrum.rows();

        Rxx.resize(this->pslength);
        RxxT.resize(this->pslength);
        x.resize(this->nb_channels, this->pslength);
        for(size_t k = 0; k < this->pslength; k++){
            Matrix<std::complex<float> > tmp_rxx(this->nb_channels, this->nb_channels);
            Matrix<std::complex<float> > tmp_rxxt(this->nb_channels, this->nb_channels);
            Rxx[k].noalias() = tmp_rxx.setIdentity();
            RxxT[k].noalias() = tmp_rxxt.setZero();
        }
    }

    bool
    CMMakerFromFFT::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        ////const auto &operation_flag = input["OPERATION_FLAG"].last().as<bool>();

#ifdef DEBUG
        ptime start_time;
        ptime end_time;
        start_time = microsec_clock::local_time();
#endif

        const auto &spectrum = input["INPUT"].last().cast<Matrix<std::complex<float> > >();
        //const auto &operation_flag = input["OPERATION_FLAG"].last().as<bool>();

        if(count == 0){ this->initialize(input); }

        auto out = data::create<Matrix<std::complex<float> > >(this->pslength, this->nb_channels * this->nb_channels);
        auto &outputcm = out.cast<Matrix<std::complex<float> > >();

        auto out_flag = data::create<bool>();
        auto &operation_flag = out_flag.cast<bool>();
        operation_flag = false;

        if(this->window_type == "PAST"){
            //Matrix<std::complex<float> > x(this->nb_channels, this->pslength);
            for(size_t ch = 0; ch < this->nb_channels; ch++){
                for(size_t k = 0; k < this->pslength; k){
                    x(ch, k) = spectrum(ch, k);
                }
            }
            if(xbuf.size() < this->window){
                xbuf.push_back(x);
            }
            else{
                xbuf.push_back(x);
                xbuf.pop_front();
            }
            if((count % this->period == 0) && ((count != 0) || (this->period == 1))){
                for(size_t i = 0; i < xbuf.size(); i++){
                    for(size_t ch = 0; ch < this->nb_channels; ch++){
                        for(size_t k = 0; k < this->pslength; k++){
                            x(ch, k) = xbuf[i](ch, k);
                        }
                    }
                    //AddToCorrelationAllFreq(RxxT, this->pslength, this->x);
                    for(size_t k = 0; k < RxxT.size(); k++){
                        RxxT[k] += x.col(k) * x.col(k).adjoint();
                    }
                    //this->sum_count++; // this->sum_count == xbuf.size()
                }
                //NormalizeCorrelationAllFreq(RxxT, this->pslength, this->sum_count);
                const float divN = 1.0f / this->xbuf.size();
                for(size_t k = 0; k < RxxT.size(); k++){
                    RxxT[k] *= divN;
                }
                for(size_t k = 0; k < Rxx.size(); k++){
                    for(size_t i = 0; i < Rxx[k].rows(); i++){
                        for(size_t j = 0; j < Rxx[k].cols(); j++){
                            Rxx[k](i, j) = RxxT[k](i, j);
                        }
                    }
                }
                if(this->enable_debug){
                    std::cout << "Count[" << count << "] : Summation Counter[" << xbuf.size() << "] : Period[" << this->period << "] : CM Created." << std::endl;
                }
                operation_flag = true;
                for(size_t k = 0; k < Rxx.size(); k++){
                    for(size_t i = 0; i < Rxx[k].rows(); i++){
                        for(size_t j = 0; j < Rxx[k].cols(); j++){
                            outputcm(k, i * Rxx[k].cols() + j) = Rxx[k](i, j);
                        }
                    }
                }
            }
            else{
                for(size_t k = 0; k < Rxx.size(); k++){
                    for(size_t i = 0; i < Rxx[k].rows(); i++){
                        for(size_t j = 0; j < Rxx[k].cols(); j++){
                            outputcm(k, i * Rxx[k].cols() + j) = Rxx[k](i, j);
                        }
                    }
                }
                //return;
            }
        }
        else{
            if((count % this->period == 0) && (count >= this->look_back)){
                for(size_t i = count - this->look_back; i <= count + this->look_ahead; i++){
                    const auto spec_i = input["INPUT"].at(i).cast<Matrix<std::complex<float> > >();
                    for(size_t ch = 0; ch < this->nb_channels; ch++){
                        for(size_t k = 0; k < this->pslength; k++){
                            x(ch, k) = spec_i(ch, k);
                        }
                    }
                    //AddToCorrelationAllFreq(RxxT, this->pslength, this->x);
                    for(size_t k = 0; k < RxxT.size(); k++){
                        RxxT[k] += x.col(k) * x.col(k).adjoint();
                    }
                    //this->sum_count++; // this->sum_count == this->look_back + 1 + this->look_ahead == this->window
                }
                //NormalizeCorrelationAllFreq(RxxT, this->pslength, this->sum_count);
                const float divN = 1.0f / this->window;
                for(size_t k = 0; k < RxxT.size(); k++){
                    RxxT[k] *= divN;
                }
                for(size_t k = 0; k < Rxx.size(); k++){
                    for(size_t i = 0; i < Rxx[k].rows(); i++){
                        for(size_t j = 0; j < Rxx[k].cols(); j++){
                            Rxx[k](i, j) = RxxT[k](i, j);
                        }
                    }
                }
                if(this->enable_debug){
                    std::cout << "Count[" << count << "] : Summation Counter[" << this->window << "] : Period[" << this->period << "] : CM Created." << std::endl;
                }
                operation_flag = true;
                for(size_t k = 0; k < Rxx.size(); k++){
                    for(size_t i = 0; i < Rxx[k].rows(); i++){
                        for(size_t j = 0; j < Rxx[k].cols(); j++){
                            outputcm(k, i * Rxx[k].cols() + j) = Rxx[k](i, j);
                        }
                    }
                }
            }
            else{
                for(size_t k = 0; k < Rxx.size(); k++){
                    for(size_t i = 0; i < Rxx[k].rows(); i++){
                        for(size_t j = 0; j < Rxx[k].cols(); j++){
                            outputcm(k, i * Rxx[k].cols() + j) = Rxx[k](i, j);
                        }
                    }
                }
                //return;
            }
        }
        //ClearCorrelationAllFreq(RxxT, this->pslength);
        for(size_t k = 0; k < RxxT.size(); k++){
            RxxT[k].setZero();
        }

        output["OUTPUTCM"].push(out);
        output["OPERATION_FLAG"].push(out_flag);

#ifdef DEBUG
        end_time = microsec_clock::local_time();
        std::cout << end_time - start_time << std::endl;
#endif

        return true;
    }
}
}

/* end of file */
