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

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

namespace hark
{
namespace node
{
    EstimateTF::EstimateTF():
        NodeDef("EstimateTF")
    {
        *this
            << Terminal::in("INPUT")
            << Terminal::in("SOURCES")
            << Terminal::in("INPUT_INPROGRESS")
            << Terminal::in("SEND_DATA")
            << Terminal::in("RESET")
            << Terminal::in("SEND_ALL")
            << Terminal::in("TARGET_COUNT")
            << Terminal::in("TRANSFER_FUNCTION")
            << Terminal::in("EST_TF_FILENAME").def("")
            << Terminal::in("LENGTH").def(512)
            << Terminal::in("SAMPLING_RATE").def(16000)
            << Terminal::in("INPUT_TYPE").def("TF")
            << Terminal::in("BASE_TF_FILENAME").def("")
            << Terminal::in("IMPORT_INTERMEDIATE_DATA").def("")
            << Terminal::in("EXPORT_INTERMEDIATE_DATA").def("")
            << Terminal::in("EXPORT_HISTORY_FILE").def(false)
            << Terminal::in("HISTORY_FILE_FRAMES").def(10000)
            << Terminal::in("HISTORY_FILE_PERIOD").def(1)
            << Terminal::in("INITIAL_TF_OUTPUT").def("ALL")
            << Terminal::in("LOOK_BACK_COUNT").def(10)
            << Terminal::in("DEBUG").def(false)
            >> Terminal::out("OUTPUT_INPROGRESS")
            >> Terminal::out("TRANSFER_FUNCTION");
    }

    void
    EstimateTF::reset()
    {
    }

    void
    EstimateTF::initialize(InputDataBundle &input)
    {
        this->filename = input["EST_TF_FILENAME"].last().cast<std::string>();

        int length = input["LENGTH"].last().as<int>();
        int sampling_rate = input["SAMPLING_RATE"].last().as<int>();
        std::string input_type_str = input["INPUT_TYPE"].last().cast<std::string>();
        std::string filename_base_tf = input["BASE_TF_FILENAME"].last().cast<std::string>();
        std::string dirnameImport = input["IMPORT_INTERMEDIATE_DATA"].last().cast<std::string>();
        std::string dirnameExport = input["EXPORT_INTERMEDIATE_DATA"].last().cast<std::string>();
        bool is_export = input["EXPORT_HISTORY_FILE"].last().as<bool>();
        int history_frames = input["HISTORY_FILE_FRAMES"].last().as<int>();
        int history_period = input["HISTORY_FILE_PERIOD"].last().as<int>();
        std::string initial_tf_str = input["INITIAL_TF_OUTPUT"].last().cast<std::string>();

        this->lookback_count = input["LOOK_BACK_COUNT"].last().as<int>();

        bool is_debug = input["DEBUG"].last().as<bool>();


        if(this->filename.empty()){
            std::stringstream ss;
            ss << "EstimateTF: EST_TF_FILENAME is not set.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
            throw RuntimeError(ss.str());
        }

        if(length <= 0){
            std::stringstream ss;
            ss << "EstimateTF: LENGTH is invalid value.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
            throw RuntimeError(ss.str());
        }

        if(sampling_rate <= 0){
            std::stringstream ss;
            ss << "EstimateTF: SAMPLING_RATE is invalid value.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
            throw RuntimeError(ss.str());
        }

        INPUT_TYPE input_type;
        if(input_type_str == "TF"){
            input_type = INPUT_TYPE::INPUT_TF;
        }
        else if(input_type_str == "INTERMEDIATE"){
            input_type = INPUT_TYPE::INPUT_INTERMEDIATE;
        }
        else{
            std::stringstream ss;
            ss << "EstimateTF: Unexpected INPUT_TYPE is set.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
            throw RuntimeError(ss.str());
        }

        if(input_type == INPUT_TYPE::TF && filename_base_tf.empty()){
            std::stringstream ss;
            ss << "EstimateTF: BASE_TF_FILENAME is not set.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
            throw RuntimeError(ss.str());
        }
        if(input_type == INPUT_TYPE::INPUT_INTERMEDIATE && dirnameImport.empty()){
            std::stringstream ss;
            ss << "EstimateTF: IMPORT_INTERMEDIATE_DATA is not set.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
            throw RuntimeError(ss.str());
        }

        if(history_period <= 0){
            std::stringstream ss;
            ss << "EstimateTF: HISTORY_FILE_PERIOD is invalid value.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
            throw RuntimeError(ss.str());
        }

        //INTIAL_TF initial_tf;
        if(initial_tf_str == "ALL"){
            this->initial_tf = INITIAL_TF::OUTPUT_ALL;
        }
        else if(initial_tf_str == "PARTIAL"){
            this->initial_tf = INITIAL_TF::OUTPUT_PARTIAL;
        }
        else{
            std::stringstream ss;
            ss << "EstimateTF: Unexpected INITIAL_TF_OUTPUT is set.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
            throw RuntimeError(ss.str());
        }

        if(this->lookback_count <= 0){
            std::stringstream ss;
            ss << "EstimateTF: LOOK_BACK_COUNT is invalid value.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
            throw RuntimeError(ss.str());
        }

        this->estimater->reset(new EstimationTransferFunction(sampling_rate, length, is_debug, is_export, history_frames, history_period));

        if(input_type == INPUT_TYPE::INPUT_TF){
            if(!this->estimater->getBasefromFile(filename_base_tf, this->num_frequency, this->num_channel, this->num_source)){
                std::stringstream ss;
                ss << "EstimateTF: Failed to get base data from the file.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
                throw RuntimeError(ss.str());
            }
        }
        else if(input_type == INPUT_TYPE::INPUT_INTERMEDIATE){
            if(!this->estimater->getBaseinDirectory(dirnameImport, this->num_frequency, this->num_channel, this->num_source)){
                std::stringstream ss;
                ss << "EstimateTF: Failed to get base data in the directory.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
                throw RuntimeError(ss.str());
            }
        }
        else{
        }

        if(is_export){
            this->estimater->initHistoryData();
        }

        if(!dirnameExport.empty()){
            if(!this->estimater->exportTransferFunction(dirnameExport, true)){
                std::stringstream ss;
                ss << "EstimateTF: Failed to export estimated TF file.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
                throw RuntimeError(ss.str());
            }
            if(!this->estimater->exportIntermData(dirnameExport, -1)){
                std::stringstream ss;
                ss << "EstimateTF: Failed to export intermediate data.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
                throw RuntimeError(ss.str());
            }
            if(is_export){
                if(!this->estimater->exportHistoryData(dirnameExport, -1)){
                    std::stringstream ss;
                    ss << "EstimateTF: Failed to export history file.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
                    throw RuntimeError(ss.str());
                }
            }
        }

        this->estimater->copyKeepTF();
        this->estimater->copyInitTFdatabase();
    }

    bool
    EstimateTF::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        ////const auto &operation_flag = input["OPERATION_FLAG"].last().as<bool>();
        const auto &sources_input = input["SOURCES"].last().cast<std::vector<Source> >();
        const auto &inprogress_input = input["INPUT_INPROGRESS"].last().as<bool>();
        const auto &senddata_input = input["SEND_DATA"].last().as<bool>();
        const auto &reset_input = input["RESET"].last().as<bool>();
        const auto &sendall_input = input["SEND_ALL"].last().as<bool>();
        const auto &targetcount_input = input["TARGET_COUNT"].last().as<int>();

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

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

        std::vector<BF::SourceLocation> source_localized;
        for(size_t i = 0; i < sources_input.size(); i++){
            if(sources_input[i].compare_mode == Source::TFINDEX){
                source_localized.push_back(BF::SourceLocation(sources_input[i].x[0], sources_input[i].x[1], sources_input[i].x[2], sources_input[i].tfindex[0], sources_input[i].id, true));
            }
            else{ // sources[i].compare_mode == Source::DEG
                source_localized.push_back(BF::SourceLocation(sources_input[i].x[0], sources_input[i].x[1], sources_input[i].x[2], -1, sources_input[i].id, true));
            }
        }

        if(targetcount_input >= 0){
            if(targetcount_input < count - this->lookback_count){
                std::cout << "EstimateTF: [ERROR] TARGET_COUNT < count - LOOK_BACK_COUNT" << std::endl;
            }
        }

        int input_index;
        if(targetcount_input < 0 || (targetcount_input >= 0 && (targetcount_input > count || targetcount_input < count - this->lookback_count))){
            input_index = count;
        }
        else{
            input_index = targetcount_input;
        }
        const auto &spectrum_input = input["INPUT"].at(input_index).cast<Matrix<std::complex<float> > >();
        if(count == 0){
            if(this->num_channel != spectrum_input.rows()){
                std::stringstream ss;
                ss << "EstimateTF: Input channel size is different from one in TF.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
                throw RuntimeError(ss.str());
            }
            if(this->num_frequency != spectrum_input.cols()){
                std::stringstream ss;
                ss << "EstimateTF: Input frame size is different from one in TF.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
                throw RuntimeError(ss.str());
            }
        }
        //Matrix<std::complex<float> > x(this->num_channel, this->num_frequency);
        Matrix<std::complex<double> > x(this->num_channel, this->num_frequency);
        for(size_t j = 0; j < this->num_channel; j++){
            for(size_t k = 0; k < this->num_frequency; k++){
                x(j, k) = spectrum_input(j, k);
            }
        }

        auto out_inprogress = data::create<bool>();
        auto &outprogress = out_inprogress.cast<bool>();
        outprogress = this->in_progress ? true : false;

        auto out_transferfunction = data::create<TransferFunction>();
        auto &transferfunction = out.cast<TransferFunction>();
        if(!(this->in_progress)){
            if(!this->estimater->exportTransferFunction(this->filename, false)){
                std::stringstream ss;
                ss << "EstimateTF: Failed to export estimated TF file.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
                throw RuntimeError(ss.str());
            }
            //return;
        }

        if(this->num_loc_source != 0){
            double max_val = x.array().abs().max().real();
            if(max_val != 0.0){
                this->estimater->estimateTransferFunction(count, x, source_localized);
            }
        }

        if(this->is_senddata){
            if(this->iocnt < 0){
                this->is_sendall = (this->initial_tf == INITIAL_TF::OUTPUT_ALL);
                this->iocnt++;
            }
            if(this->is_reset){
                this->estimater->resetTFdatabase();
                this->is_sendall = true;
            }
            this->estimater->outputTF(this->is_sendall, transferfunction, source_localized);
        }

        output["OUTPUT_INPROGRESS"].push(out_inprogress);
        output["TRANSFER_FUNCTION"].push(out_transferfunction);

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

        return true;
    }
}
}

/* end of file */
