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

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

#if defined(HARKLIB_USES_LIBTORCH)

#undef HARKLIB_VAD_PYANNOTE_DUMP_SCORE
//	#define HARKLIB_VAD_PYANNOTE_DUMP_SCORE

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

#include <math.h>
#include <torch/script.h>

#include "../../models/pyannote_jit_script.stats"

namespace hark
{
namespace node
{
    class Impl: public VADPyAnnote::BaseImpl
    {
    public:
        Impl()
        {
            this->reset();
        }

        void reset()
        {
            this->_score_buffer_offset = 0;
            this->_score_active = false;
            this->_score_duration = 0;
            this->_source_active = false;
            this->_source_head = 0;
            this->_source_tail = 0;
            this->_source_frame_active = 0;
            this->_source_frame_count = 0;
            // this->_source_sequence = 0;
            this->_audio_segments.clear();
            this->_audio_segment_offset = 0;
            this->_last_active_frame = 0;
            this->_maximum_base_scores_count = 0;
            this->_base_scores.clear();
        }

    public:
        torch::jit::script::Module _module;

        size_t _score_buffer_offset;
        bool   _score_active;
        size_t _score_duration;

        bool   _source_active;
        size_t _source_head, _source_tail;

        bool   _source_frame_active;
        size_t _source_frame_count;

        std::deque<data> _audio_segments;
        size_t _audio_segment_offset;
        size_t _last_active_frame;
        
        size_t _maximum_base_scores_count;
        std::deque<float> _base_scores;
    };

    VADPyAnnote::VADPyAnnote():
        NodeDef("VADPyAnnote")
    {
        *this
            << Terminal::in("INPUT")
            << Terminal::in("MODEL_PATH")
            << Terminal::in("SAMPLING_RATE").def(16000)
            << Terminal::in("LENGTH").def(512)
            << Terminal::in("ADVANCE").def(160)
            << Terminal::in("CHANNEL_COUNT").def(1)
            << Terminal::in("SINCNET_KERNEL_SIZE").def(27)
            << Terminal::in("SINCNET_PADDING").def(148)
            << Terminal::in("ONSET_SCORE_THRESOLD").def(0.1)
            << Terminal::in("OFFSET_SCORE_THRESOLD").def(0.0)
            << Terminal::in("HEAD_MARGIN").def(300)
            << Terminal::in("TAIL_MARGIN").def(400)
            << Terminal::in("SPLIT_PERIOD").def(0)
            << Terminal::in("INITIAL_NORMALIZE_SAMPLE_COUNT").def("")
            << Terminal::in("INITIAL_NORMALIZE_MEAN").def("")
            << Terminal::in("INITIAL_NORMALIZE_STDDEV").def("")
            // << Terminal::in("INITIAL_NORMALIZE_SAMPLE_COUNT").def("96000 31470 10482 3486")
            // << Terminal::in("INITIAL_NORMALIZE_MEAN").def("0.5 2.1 0.2 0.1")
            // << Terminal::in("INITIAL_NORMALIZE_STDDEV").def("0.3 1.2 1.0 0.6")
            >> Terminal::out("OUTPUT") /* Map<int,ObjectRef> */
            >> Terminal::out("SOURCES"); /* Vector<ObjectRef> */
    }

    void
    VADPyAnnote::reset()
    {
        this->_channel_impls.clear();
        this->_model_path.clear();
        this->_sampling_rate = 0;
        this->_advance = 0;
        this->_length = 0;
        this->_split_period = 0;
        this->_sincnet_kernel_size = 0;
        this->_sincnet_padding = 0;
        this->_onset_score_threshold = 0;
        this->_offset_score_threshold = 0;
        this->_split_period_frames = 0;
        this->_head_margin_frames = 0;
        this->_tail_margin_frames = 0;
        this->_head_margin_kernels = 0;
        this->_tail_margin_kernels = 0;
        this->_initial_normalize_counts.clear();
        this->_initial_normalize_means.clear();
        this->_initial_normalize_sigmas.clear();
        // this->_score_buffer_offset = 0;
        // this->_score_active = false;
        // this->_score_duration = 0;
        // this->_source_active = false;
        // this->_source_head = 0;
        // this->_source_tail = 0;
        // this->_source_frame_active = 0;
        // this->_source_frame_count = 0;
        this->_source_sequence = 0;
        // this->_audio_segments.clear();
        // this->_audio_segment_offset = 0;
        // this->_last_active_frame = 0;
        // this->_maximum_base_scores_count = 0;
        // this->_base_scores.clear();
    }

    template<typename T>
    std::vector<T>
    _parse_array_string(const std::string &input)
    {
        std::istringstream ss(input);	

        std::vector<T> values;
        T value;
        while(ss >> value){
            values.push_back(value);
        }

        return values;
    }

    bool
    VADPyAnnote::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        // std::cout << "VADPyAnnote::operator(): count=" << count << std::endl;
        if(count == 0){
            this->_model_path = input["MODEL_PATH"].last().cast<std::string>();
            this->_sampling_rate = input["SAMPLING_RATE"].last().as<size_t>();
            this->_length = input["LENGTH"].last().as<size_t>();
            this->_advance = input["ADVANCE"].last().as<size_t>();
            this->_channel_count = input["CHANNEL_COUNT"].last().as<size_t>();
            this->_split_period = input["SPLIT_PERIOD"].last().as<float>();
            this->_sincnet_kernel_size = input["SINCNET_KERNEL_SIZE"].last().as<size_t>();
            this->_sincnet_padding = input["SINCNET_PADDING"].last().as<size_t>();
            this->_onset_score_threshold = input["ONSET_SCORE_THRESOLD"].last().as<float>();
            this->_offset_score_threshold = input["OFFSET_SCORE_THRESOLD"].last().as<float>();

            this->_initial_normalize_counts = _parse_array_string<int>(input["INITIAL_NORMALIZE_SAMPLE_COUNT"].last().cast<std::string>());
            this->_initial_normalize_means = _parse_array_string<float>(input["INITIAL_NORMALIZE_MEAN"].last().cast<std::string>());
            this->_initial_normalize_sigmas = _parse_array_string<float>(input["INITIAL_NORMALIZE_STDDEV"].last().cast<std::string>());

            this->_split_period_frames = (int)(this->_split_period * this->_sampling_rate / 1000 / this->_advance);
            this->_head_margin_frames = input["HEAD_MARGIN"].last().as<int>() * this->_sampling_rate / 1000 / this->_advance;
            this->_tail_margin_frames = input["TAIL_MARGIN"].last().as<int>() * this->_sampling_rate / 1000 / this->_advance;
            this->_head_margin_kernels = input["HEAD_MARGIN"].last().as<int>() * this->_sampling_rate / 1000 / this->_sincnet_kernel_size;
            this->_tail_margin_kernels = input["TAIL_MARGIN"].last().as<int>() * this->_sampling_rate / 1000 / this->_sincnet_kernel_size;

            // this->_window_count_processed = 0;
            // this->_audio_buffer_offset = 0;
            // this->_audio_buffer.clear();

            this->_channel_impls.resize(this->_channel_count);
            for(size_t channel_index = 0; channel_index < this->_channel_count; channel_index++){
                this->_channel_impls[channel_index] = std::make_unique<Impl>();
            }
            // this->_score_buffer_offset = 0;
            // this->_score_active = false;
            // this->_score_duration = 0;

            // this->_source_head = 0;
            // this->_source_tail = 0;
            // this->_source_frame_active = false;
            // this->_source_frame_count = 0;
            this->_source_sequence = 0;

            // this->_audio_segments.clear();
            // this->_last_active_frame = 0;
            
            // this->_base_scores.clear();
            // this->_maximum_base_scores_count = 10 * this->_sampling_rate / this->_sincnet_kernel_size;

            // this->_score_count_processed = 0;

            // std::cout << "VADPyAnnote: initialize: model-path=" << this->_model_path << std::endl;
            for(auto &_impl: this->_channel_impls){
                Impl *impl = dynamic_cast<Impl *>(_impl.get());

                impl->reset();
                try{
                    impl->_module = torch::jit::script::Module(torch::jit::load(this->_model_path));
                }catch(const c10::Error& e){
                    std::stringstream ss;
                    ss << "Failed to load pyannote model: " << e.what();
                    throw hark::RuntimeError(ss.str());
                }
                try{
                    if(this->_initial_normalize_counts.empty() && this->_initial_normalize_means.empty() && this->_initial_normalize_sigmas.empty()){
                        // impl._module.get_method("reset")({});
                        std::vector<at::Tensor> means = {
                            torch::from_blob(STATS_MEANS[0].data(), {(int)STATS_MEANS[0].size()}, at::kFloat),
                            torch::from_blob(STATS_MEANS[1].data(), {(int)STATS_MEANS[1].size()}, at::kFloat),
                            torch::from_blob(STATS_MEANS[2].data(), {(int)STATS_MEANS[2].size()}, at::kFloat),
                            torch::from_blob(STATS_MEANS[3].data(), {(int)STATS_MEANS[3].size()}, at::kFloat),
                        };
                        std::vector<at::Tensor> sigmas = {
                            torch::from_blob(STATS_SIGMAS[0].data(), {(int)STATS_SIGMAS[0].size()}, at::kFloat),
                            torch::from_blob(STATS_SIGMAS[1].data(), {(int)STATS_SIGMAS[1].size()}, at::kFloat),
                            torch::from_blob(STATS_SIGMAS[2].data(), {(int)STATS_SIGMAS[2].size()}, at::kFloat),
                            torch::from_blob(STATS_SIGMAS[3].data(), {(int)STATS_SIGMAS[3].size()}, at::kFloat),
                        };
                        std::vector<at::Tensor> counts = {
                            torch::from_blob(STATS_COUNTS.data() + 0, {1}, at::kInt),
                            torch::from_blob(STATS_COUNTS.data() + 1, {1}, at::kInt),
                            torch::from_blob(STATS_COUNTS.data() + 2, {1}, at::kInt),
                            torch::from_blob(STATS_COUNTS.data() + 3, {1}, at::kInt),
                        };
                        impl->_module.get_method("reset")({at::TensorList(means), at::TensorList(sigmas), at::TensorList(counts)});
                    }else if(this->_initial_normalize_counts.size() == 4 && this->_initial_normalize_means.size() == 4 && this->_initial_normalize_sigmas.size() == 4){
                        std::vector<at::Tensor> means = {
                            torch::from_blob(this->_initial_normalize_means.data() + 0, {1}, at::kFloat),
                            torch::from_blob(this->_initial_normalize_means.data() + 1, {1}, at::kFloat),
                            torch::from_blob(this->_initial_normalize_means.data() + 2, {1}, at::kFloat),
                            torch::from_blob(this->_initial_normalize_means.data() + 3, {1}, at::kFloat),
                        };
                        std::vector<at::Tensor> sigmas = {
                            torch::from_blob(this->_initial_normalize_sigmas.data() + 0, {1}, at::kFloat),
                            torch::from_blob(this->_initial_normalize_sigmas.data() + 1, {1}, at::kFloat),
                            torch::from_blob(this->_initial_normalize_sigmas.data() + 2, {1}, at::kFloat),
                            torch::from_blob(this->_initial_normalize_sigmas.data() + 3, {1}, at::kFloat),
                        };
                        std::vector<at::Tensor> counts = {
                            torch::from_blob(this->_initial_normalize_counts.data() + 0, {1}, at::kInt),
                            torch::from_blob(this->_initial_normalize_counts.data() + 1, {1}, at::kInt),
                            torch::from_blob(this->_initial_normalize_counts.data() + 2, {1}, at::kInt),
                            torch::from_blob(this->_initial_normalize_counts.data() + 3, {1}, at::kInt),
                        };
                        impl->_module.get_method("reset")({at::TensorList(means), at::TensorList(sigmas), at::TensorList(counts)});
                    }else{
                        throw hark::RuntimeError("Unexpected size of INITIAL_NORMALIZE_SAMPLE_COUNT/MEAN/STDDEV.  Expected (0, 0, 0) or (4, 4, 4)");
                    }
                
                }catch(const c10::Error& e){
                    std::stringstream ss;
                    ss << "Failed to initialize pyannote model: " << e.what();
                    throw hark::RuntimeError(ss.str());
                }
                impl->_module.eval();
            }
        }

        const Matrix<float> &audio_matrix = input["INPUT"].last().cast<Matrix<float>>();
        if(audio_matrix.rows() != this->_channel_count){
            std::stringstream ss;
            ss << "VADPyAnnote: Channel count mismatch (input=" << audio_matrix.rows() << ", expected=" << this->_channel_count << ")";
            throw RuntimeError(ss.str());
        }

        std::vector<data>   sources;
        std::map<int, data> outputs;
        for(size_t channel_index = 0; channel_index < this->_channel_count; channel_index++){
            auto &_impl = this->_channel_impls[channel_index];
            Impl *impl = dynamic_cast<Impl *>(_impl.get());

            // do pyannote process
            std::vector<float> audio(audio_matrix.data(), audio_matrix.data() + audio_matrix.size());
            impl->_audio_segments.emplace_back(audio);

            at::Tensor inputTensor = torch::from_blob(
                (void *)(count == 0 ? audio.data() : audio.data() + (audio.size() - this->_advance)),
                {1, (int)(count == 0 ? audio.size() : this-> _advance), 1},
                at::kFloat);

            impl->_module.get_method("push_status")({});
    //        at::Tensor dummyInputTensor = torch::zeros({1, (int)this->_sampling_rate, 1}, at::kFloat);
            at::Tensor dummyInputTensor = torch::zeros_like(inputTensor);
            at::Tensor dummyOutputTensor;
            {
                torch::autograd::AutoGradMode guard(false);
                at::AutoNonVariableTypeMode non_var_type_mode(true);
                dummyOutputTensor = impl->_module.forward({dummyInputTensor}).toTensor();
            }
            impl->_module.get_method("pop_status")({});

            at::Tensor outputTensor;
            {
                torch::autograd::AutoGradMode guard(false);
                at::AutoNonVariableTypeMode non_var_type_mode(true);
                outputTensor = impl->_module.forward({inputTensor}).toTensor();
            }

            // retrieve scores
            size_t dummy_score_count = dummyOutputTensor.size(1);
            const float *dummy_scores = dummyOutputTensor.data_ptr<float>();
            for(size_t i = 0; i < dummy_score_count; i++){
                impl->_base_scores.push_front(dummy_scores[i * 2 + 0]);
            }
            if(impl->_base_scores.size() > impl->_maximum_base_scores_count){
                impl->_base_scores.resize(impl->_maximum_base_scores_count);
            }
            float base_score_mean = std::accumulate(impl->_base_scores.begin(), impl->_base_scores.end(), 0.0f) / impl->_base_scores.size();

            size_t score_count = outputTensor.size(1);
            size_t feature_count = outputTensor.size(2);
            const float *scores = outputTensor.data_ptr<float>();

            // store scores
            for(size_t i = 0; i < score_count; i++){
                size_t index = impl->_score_buffer_offset + i;

                float s1 = scores[i * 2 + 0];
                float s2 = scores[i * 2 + 1];
    //            float bs = dummy_scores[i * 2 + 0];
                float sr = s1 - base_score_mean; // s1 - bs; // s1; // s1 - s2; //s1 / (s2 + 1e-6);

                bool   active = impl->_score_active;
                size_t duration = impl->_score_duration;

                if(active){
                    if(sr < this->_offset_score_threshold){
                        active = false;
                        duration = 0;
                    }else{
                        duration++;
                    }
                }else{
                    if(sr > this->_onset_score_threshold){
                        active = true;
                        duration = 0;
                    }else{
                        duration++;
                    }
                }
                impl->_score_active = active;
                impl->_score_duration = duration;

                // Score score = {s1, s2, sr, active, duration};

                // std::cout << "index=" << index << ", sr=" << sr << ", th=(" << this->_onset_score_threshold << ", " << this->_offset_score_threshold << "), active=" << active << ", source-active-prev=" << this->_source_active;
    #if defined(HARKLIB_VAD_PYANNOTE_DUMP_SCORE)
                std::cout.setf(std::ios::fixed | std::ios::showpos);
                std::cout << "index=" << index << std::setprecision(3) << ", s1=" << s1 << ", s2=" << s2 << ", sr=" << sr << ", bs=" << base_score_mean << ", th=(" << this->_onset_score_threshold << ", " << this->_offset_score_threshold << "), active=" << active << ", source-active-prev=" << impl->_source_active;
    #endif /* #if defined(HARKLIB_VAD_PYANNOTE_DUMP_SCORE) */
                if(impl->_source_active){
                    if(active){
                        // continue activity
                        impl->_source_tail = index;
                    }else if(index < impl->_source_tail + std::max(this->_head_margin_kernels, this->_tail_margin_kernels)){
                        // do nothing
                    }else{
                        // close activity
                        impl->_source_active = false;
                    }
                }else{
                    if(active){
                        // start activity
                        impl->_source_head = impl->_source_tail = index;
                        impl->_source_active = true;
                    }else{
                        // do nothing
                    }
                }
                // std::cout << ", source-active-next=" << this->_source_active << std::endl;
    #if defined(HARKLIB_VAD_PYANNOTE_DUMP_SCORE)
                std::cout << ", source-active-next=" << impl->_source_active << std::endl;
    #endif /* #if defined(HARKLIB_VAD_PYANNOTE_DUMP_SCORE) */

                // std::cout << "count=" << count << ", kernel=" << index << ", time=" << ((float)index * this->_sincnet_kernel_size / this->_sampling_rate) << ", s1=" << s1 << ", s2=" << s2 << ", sr=" << sr << ", active=" << active << ", score-active=" << this->_score_active << ", score-head=" << this->_source_head << ", score-tail=" << this->_source_tail << std::endl;
                if(index * this->_sincnet_kernel_size < this->_length){
                    continue;
                }

                size_t frame_index = (index * this->_sincnet_kernel_size - this->_length) / this->_advance;
                // std::cout << "frame-index: " << frame_index << ", index: " << index << ", kernel-size: " << this->_sincnet_kernel_size << ", advance: " << this->_advance << std::endl;
                if(impl->_source_frame_count < frame_index &&
                impl->_audio_segment_offset + impl->_audio_segments.size() >= impl->_source_frame_count + 1 + this->_head_margin_frames){
                    // if((this->_source_active && this->_source_frame_count <= frame_index) || (!this->_source_active && this->_source_frame_count < frame_index)){
                    // std::cout << "frame-index: " << frame_index << ", source-active=" << this->_source_active << std::endl;
                    
                    // std::cout << "source-active=" << this->_source_active << ", frame-index=" << frame_index << ", last-active-frame=" << this->_last_active_frame << std::endl;
                    if(impl->_source_active){
                        impl->_last_active_frame = count;
                    }

                    bool output_frame = (impl->_source_active || (this->_source_sequence > 0 && frame_index <= impl->_last_active_frame + this->_tail_margin_frames));

                    // std::cout << "output-frame=" << output_frame << ", sequence=" << this->_source_sequence << std::endl;
                    if(output_frame){//this->_source_frame_count + this->_head_margin_frames >= this->_last_active_frame){
                        hark::Source new_source;

                        size_t source_id;
                        if(impl->_source_frame_active){
                            new_source.id = this->_source_sequence;
                        }else{
                            new_source.id = ++this->_source_sequence;
                        }
                        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.mic_id = channel_index;
                        new_source.sub_id = 0;
                        new_source.count = count;
                        new_source.first_found_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;
                        }

                        sources.push_back(data(new_source));
                        outputs.emplace(new_source.id, impl->_audio_segments[this->_head_margin_frames]);
                    }

                    // std::cout << "push: " << this->_source_frame_count << ", active=" << this->_source_frame_active << ", sequence=" << this->_source_sequence << std::endl;
                    impl->_source_frame_count++;
                    impl->_source_frame_active = output_frame;

                    if(impl->_audio_segments.size() > 1 + this->_head_margin_frames + this->_tail_margin_frames){
                        impl->_audio_segments.pop_front();
                        impl->_audio_segment_offset++;
                    }
                }
            }

            impl->_score_buffer_offset += score_count;
        }
        output["SOURCES"].push(data(std::move(sources)));
        output["OUTPUT"].push(data(std::move(outputs)));

        return true;
    }
}
}

#endif /* #if defined(HARKLIB_USES_LIBTORCH) */

/* end of file */
