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

#if !defined(_HARKLIB_NODE_SAVE_HTK_FEATURES_H_INCLUDED_)
#define _HARKLIB_NODE_SAVE_HTK_FEATURES_H_INCLUDED_ 1

#include <fstream>

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

#include "hark-lib/byteswap.h"

namespace hark
{
namespace node
{
    class SaveHTKFeatures: public NodeDef
    {
    public:
        SaveHTKFeatures();
        ~SaveHTKFeatures();

        void initialize(InputDataBundle &input);
        virtual void reset();
        virtual bool operator () (size_t count, InputDataBundle &input, OutputDataBundle &output);

    private:
        struct alignas(1) HTKHeader {
            uint32_t nSamples;
            uint32_t Period;
            uint16_t BytesPerSample;
            uint16_t FeatureType;
        };
        class HTKFileIO {
        public:
            static constexpr uint16_t BASEMASK= 0x003F;
            //----------------------------------
            static constexpr uint16_t MFCC    = 0x0006;
            static constexpr uint16_t FBANK   = 0x0007;
            static constexpr uint16_t USER    = 0x0009;
            static constexpr uint16_t PLP     = 0x000B;
            static constexpr uint16_t FSTREAM = 0x000C;
            //----------------------------------
            static constexpr uint16_t _E      = 0x0040; // log energy included
            static constexpr uint16_t _N      = 0x0080; // absolute energy suppressed
            static constexpr uint16_t _D      = 0x0100; // delta coef appended
            static constexpr uint16_t _A      = 0x0200; // acceleration (delta-delta) coefs appended
            static constexpr uint16_t _C      = 0x0400; // is compressed
            static constexpr uint16_t _Z      = 0x0800; // zero meaned
            static constexpr uint16_t _K      = 0x1000; // has CRC check
            static constexpr uint16_t _0      = 0x2000; // 0'th Cepstra included
            static constexpr uint16_t _V      = 0x4000; // has VQ index attached
            static constexpr uint16_t _T      = 0x8000; // has jerk (delta-delta-delta) coefs appended
            //----------------------------------
            static constexpr uint16_t MFCC_E_D   = MFCC + _E + _D;
            static constexpr uint16_t MFCC_E_D_Z = MFCC + _E + _D + _Z;

            HTKFileIO() : filename(""), write_mode(false) {
                this->header.nSamples = 0;
                this->header.Period = 0;
                this->header.BytesPerSample = 0;
                this->header.FeatureType = 0;
            }
            ~HTKFileIO(){
                if(this->stream.is_open()){ this->closefile(); }
            }

            void set_header(HTKHeader &header){
                this->header = header;
            }
            void set_header(uint32_t samples, uint32_t period, uint16_t bps, uint16_t type){
                this->header.nSamples = hark::byteswap(samples);
                this->header.Period = hark::byteswap(period);
                this->header.BytesPerSample = hark::byteswap(bps);
                this->header.FeatureType = hark::byteswap(type);
            }

            HTKHeader get_header(){
                return this->header;
            }

            void set_samples(uint32_t samples){
                this->header.nSamples = hark::byteswap(samples);
            }
            void set_period(uint32_t period){
                this->header.Period = hark::byteswap(period);
            }
            void set_bps(uint16_t bps){
                this->header.BytesPerSample = hark::byteswap(bps);
            }
            void set_type(uint16_t type){
                this->header.FeatureType = hark::byteswap(type);
            }

            uint32_t get_samples(){
                return hark::byteswap(this->header.nSamples);
            }
            uint32_t get_period(){
                return hark::byteswap(this->header.Period);
            }
            uint16_t get_bps(){
                return hark::byteswap(this->header.BytesPerSample);
            }
            uint16_t get_type(){
                return hark::byteswap(this->header.FeatureType);
            }

            void openfile(std::string &filename, bool write_mode=false){
                if(filename == this->filename &&
                   write_mode == this->write_mode &&
                   this->stream.is_open()){
                    std::cout << "already opened " << this->filename << "." << std::endl;
                    return;
                }
                if(this->stream.is_open()){ this->closefile(); }
                this->filename = filename;
                this->write_mode = write_mode;
                if(this->write_mode){
                    this->stream.open(this->filename, std::ios::out | std::ios::binary);
                    this->header_begin_pos = this->stream.tellp();
                    this->current_data_pos = this->stream.tellp();
                }
                else{
                    this->stream.open(this->filename, std::ios::in | std::ios::binary);
                    this->header_begin_pos = this->stream.tellg();
                    this->current_data_pos = this->stream.tellg();
                }
                if(!this->stream.is_open()){
                    std::cout << "cannot open file " << this->filename << std::endl;
                    this->filename = "";
                }
            }

            void closefile(){
                if(!this->filename.empty()){
                    if(this->stream.is_open()){ this->stream.close(); }
                }
                else{
                    std::cout << "already closed " << this->filename << "." << std::endl;
                }
                this->filename = "";
            }

            void read_header(){
                if(!this->stream.is_open()){
                    std::cout << this->filename << " is not opened.\n@HTKFileIO::read_header" << std::endl;
                    return;
                }
                if(this->write_mode){
                    std::cout << this->filename << " is opened in write mode.\n@HTKFileIO::read_header" << std::endl;
                    return;
                }
                this->stream.seekg(this->header_begin_pos);
                this->stream.read(reinterpret_cast<char*>(&this->header.nSamples), sizeof(uint32_t));
                this->stream.read(reinterpret_cast<char*>(&this->header.Period), sizeof(uint32_t));
                this->stream.read(reinterpret_cast<char*>(&this->header.BytesPerSample), sizeof(uint16_t));
                this->stream.read(reinterpret_cast<char*>(&this->header.FeatureType), sizeof(uint16_t));
                if(this->header_begin_pos == this->current_data_pos){
                    this->current_data_pos = this->stream.tellg();
                }
                this->stream.seekg(this->current_data_pos);
            }

            void write_header(){
                if(!this->stream.is_open()){
                    std::cout << this->filename << " is not opened.\n@HTKFileIO::write_header" << std::endl;
                    return;
                }
                if(!this->write_mode){
                    std::cout << this->filename << " is opened in read mode.\n@HTKFileIO::write_header" << std::endl;
                    return;
                }
                this->stream.seekp(this->header_begin_pos);
                this->stream.write(reinterpret_cast<char*>(&this->header.nSamples), sizeof(uint32_t));
                this->stream.write(reinterpret_cast<char*>(&this->header.Period), sizeof(uint32_t));
                this->stream.write(reinterpret_cast<char*>(&this->header.BytesPerSample), sizeof(uint16_t));
                this->stream.write(reinterpret_cast<char*>(&this->header.FeatureType), sizeof(uint16_t));
                if(this->header_begin_pos == this->current_data_pos){
                    this->current_data_pos = this->stream.tellp();
                }
                this->stream.seekp(this->current_data_pos);
            }

            void read_features(std::vector<float> &features){
                if(!this->stream.is_open()){
                    std::cout << this->filename << " is not opened.\n@HTKFileIO::read_features" << std::endl;
                    return;
                }
                if(this->write_mode){
                    std::cout << this->filename << " is opened in write mode.\n@HTKFileIO::read_features" << std::endl;
                    return;
                }
                if(this->header_begin_pos == this->current_data_pos){ this->read_header(); }
                if(features.empty()){ return; }
                this->stream.read(reinterpret_cast<char*>(features.data()), sizeof(float) * features.size());
                hark::byteswap(features.begin(), features.end());
                this->current_data_pos = this->stream.tellg();
            }

            void write_features(std::vector<float> features){
                // Rewrite the vector of arguments in this method.
                // Since it does not change the original vector,
                // it requires a copy on purpose, so please do not
                // rewrite the argument to pass by referebce.
                if(!this->stream.is_open()){
                    std::cout << this->filename << " is not opened.\n@HTKFileIO::write_features" << std::endl;
                    return;
                }
                if(!this->write_mode){
                    std::cout << this->filename << " is opened in read mode.\n@HTKFileIO::write_features" << std::endl;
                    return;
                }
                if(this->header_begin_pos == this->current_data_pos){ this->write_header(); }
                hark::byteswap(features.begin(), features.end());
                this->stream.write(reinterpret_cast<char*>(features.data()), sizeof(float) * features.size());
                this->current_data_pos = this->stream.tellp();
            }

        private:
            HTKHeader header;
            std::string filename;
            bool write_mode;
            std::fstream stream;
            std::fstream::pos_type header_begin_pos;
            std::fstream::pos_type current_data_pos;
        };

        //std::map<int, std::shared_ptr<std::ostream> >_streams;
        std::map<int, std::shared_ptr<HTKFileIO> > htkfileio;
        std::string basename;
        HTKHeader init_header;
    };
}
}

#endif /* #if !defined(_HARKLIB_NODE_SAVE_HTK_FEATURES_H_INCLUDED_) */

/* end of file */
