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

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

#include <fmt/core.h>

#include <math.h>
#include <fstream>

#include <iomanip>
#include <chrono>

namespace hark
{
namespace node
{
    static std::string
    format_basename_by_location(const std::string &basename, float x, float y, float z, int id, const std::chrono::system_clock::time_point *tv, bool format, const std::string &suffix)
    {
        float azimuth_deg = atan2(y, x) * 180 / M_PI;
        float elevation_deg = atan2(z, hypot(x, y)) * 180 / M_PI;

        int azimuth_deg_int = (int)round(azimuth_deg);
        int elevation_deg_int = (int)round(elevation_deg);

        // float azimuth_f = 180.0f / M_PI * atan2(x1, x0);
        // int azimuth = (azimuth_f > 0.0f) ? (int)(azimuth_f + 0.5f) : (int)(-1.0f * (fabsf(azimuth_f) + 0.5f));
        // float elevation_f = 180.0f / M_PI * atan2(x2, sqrt(x0 * x0 + x1 * x1));
        // int elevation = (elevation_f > 0.0f) ? (int)(elevation_f + 0.5f) : (int)(-1.0f * (fabsf(elevation_f) + 0.5f));

        // current date/time based on current system
        std::string date;
        {
            auto now = std::chrono::system_clock::now();
            auto now_tt = std::chrono::system_clock::to_time_t(now);

            std::stringstream ss;
            ss << std::put_time(std::localtime(&now_tt), "%Y%m%d-%H%M%S");
            date = ss.str();

//            time_t now = time(0);
//            struct tm tm;
//#if _WIN32
//            tm = *localtime(&now);
//#else
//            localtime_r(&now, &tm);
//#endif
//
//            char buffer[30];
//            strftime(buffer, 30, "%Y%m%d-%H%M%S", &tm);
//            date = buffer;
        }

        std::string timestamp;
        if(tv){
            auto now_tt = std::chrono::system_clock::to_time_t(*tv);
            auto usecs = std::chrono::time_point_cast<std::chrono::milliseconds>(*tv);
            long us = usecs.time_since_epoch().count() % 1000000;

            std::stringstream ss;
            ss << std::put_time(std::localtime(&now_tt), "%Y%m%d.%H%M%S.") << std::setw(6) << std::setfill('0') << us;
            timestamp = ss.str();

//
//            struct tm tm;
//#if _WIN32
//            time_t t = tv.tv_sec;
//            tm = *localtime(&t);
//#else
//            localtime_r(&tv.tv_sec, &tm);
//#endif
//
//            char buf[128];
//            snprintf(buf, 128, ".%04d%02d%02d.%02d%02d%02d.%06ld", tm.tm_year + 1900, tm.tm_mon + 1, tm.tm_mday,
//                    tm.tm_hour, tm.tm_min, tm.tm_sec, tv.tv_usec);
//
//            timestamp = buf;
        }

        std::ostringstream basename_stream;
        if(format){
            basename_stream
                    << fmt::format(basename, fmt::arg("azimuth", azimuth_deg_int), fmt::arg("elevation", elevation_deg_int), fmt::arg("srcid", id), fmt::arg("date", date))
                    << timestamp
                    << suffix; //format basename
        }else{ //use original format if basename is not formatted
            basename_stream
                    << basename << id << timestamp << suffix;
        }

        return basename_stream.str();
    }

    static std::string
    format_basename_by_source(const std::string &basename, const Source &source, const std::chrono::system_clock::time_point *tv, bool format, const std::string &suffix)
    {
        return format_basename_by_location(basename, source.x[0], source.x[1], source.x[2], source.id, tv, format, suffix);
    }

    static void
    write_riff_header(std::ostream &os, size_t channel_count, size_t frames_per_second, size_t bits_per_sample)
    {
        const char    *riff_marker = "RIFF";
        const uint32_t tsize = 44;
        const char    *wave_marker = "WAVE";
        const char    *fmt_marker = "fmt ";
        const uint32_t fmt_size = 16;
        const uint16_t wFormatTag = 1; // PCM
        const uint16_t nChannels = channel_count;
        const uint32_t nSamplesPerSec = frames_per_second;
        const uint32_t nAvgBytesPerSec = channel_count * frames_per_second * bits_per_sample / 8;
        const uint16_t nBlockAlign = channel_count * bits_per_sample / 8;
        const uint16_t wBitsPerSample = bits_per_sample;
        const uint16_t pcm_type = 1;
        const char    *data_marker = "data";
        const uint32_t dsize = 0;

        os.seekp(0, std::ios::beg);
        os.write(riff_marker, 4);
        os.write(reinterpret_cast<const char *>(&tsize), 4);
        os.write(wave_marker, 4);

        os.write(fmt_marker, 4);
        os.write(reinterpret_cast<const char *>(&fmt_size), 4);
        os.write(reinterpret_cast<const char *>(&wFormatTag), 2);
        os.write(reinterpret_cast<const char *>(&nChannels), 2);
        os.write(reinterpret_cast<const char *>(&nSamplesPerSec), 4);
        os.write(reinterpret_cast<const char *>(&nAvgBytesPerSec), 4);
        os.write(reinterpret_cast<const char *>(&nBlockAlign), 2);
        os.write(reinterpret_cast<const char *>(&wBitsPerSample), 2);

        os.write(data_marker, 4);
        os.write(reinterpret_cast<const char *>(&dsize), 4);
    }

    static void
    update_riff_header(std::ostream &os)
    {
        os.flush();
        const uint32_t cpos = os.tellp();
        
        // retrieve actual file size
        os.seekp(0, std::ios::end);
        uint32_t fsize = os.tellp();
        os.seekp(0, std::ios::beg);
        fsize -= os.tellp();

        const uint32_t tsize = fsize - 8;
        const uint32_t dsize = fsize - 44;

        os.seekp(4, std::ios::beg);
        os.write(reinterpret_cast<const char *>(&tsize), 4);
        os.seekp(40, std::ios::beg);
        os.write(reinterpret_cast<const char *>(&dsize), 4);
        os.seekp(cpos, std::ios::beg);

    }

    static std::vector<char>
    pack_pcm(const Matrix<float> &x, size_t t_first, size_t t_last, size_t input_bits, size_t output_bits)
    {
#if _WIN32
        typedef int ssize_t;
#endif
        const size_t bytes_per_sample = (output_bits / 8);

        std::vector<char> buffer(x.rows() * (t_last - t_first) * bytes_per_sample);
        auto it = buffer.begin();

        float clipmax, clipmin;
        ssize_t diff_bits = (output_bits - input_bits);

        switch(bytes_per_sample){
        case 2:
            clipmax = std::numeric_limits<int16_t>::max();
            clipmin = std::numeric_limits<int16_t>::min();
            break;
        case 4:
            clipmax = std::numeric_limits<int32_t>::max();
            clipmin = std::numeric_limits<int32_t>::min();
            break;
        default:
            std::stringstream ss;
            ss << "SaveWavePCM: Unexpected output_bits: " << output_bits;
            throw RuntimeError(ss.str());
        }

        // std::cout << "pack_pcm: rows=" << x.rows() << ", cols=" << x.cols() << ", t_first=" << t_first << ", t_last=" << t_last << ", clipmax=" << clipmax << ", clipmin=" << clipmin << ", diff_bits=" << diff_bits << ", bytes_per_sample=" << bytes_per_sample << std::endl;

        bool has_clipped = false;
        for(size_t j = t_first; j < t_last; j++){
            for(size_t i = 0; i < x.rows(); i++){
                float v = x(i, j);

                if(diff_bits > 0){
                    v *= (1 << diff_bits);
                }else if(diff_bits < 0){
                    v /= (1 << -diff_bits);
                }

                if(v > clipmax){
                    v = clipmax;
                    has_clipped = true;
                }else if(v < clipmin){
                    v = clipmin;
                    has_clipped = true;
                }

                int32_t iv = (int32_t)v;
                char *p = reinterpret_cast<char *>(&iv);

                it = std::copy(p, p + bytes_per_sample, it);
            }
        }
        // if(clipped){
        //     cout << "Clip has been detected !!" << endl;
        // }

        return buffer;
    }

    static std::vector<char>
    pack_pcm(const std::vector<float> &x, size_t t_first, size_t t_last, size_t input_bits, size_t output_bits)
    {
        const auto &matrix = Eigen::Map<const Eigen::MatrixXf, Eigen::Unaligned>(x.data(), 1, x.size());
        return pack_pcm(matrix, t_first, t_last, input_bits, output_bits);
    }

    SaveWavePCM::SaveWavePCM():
        NodeDef("SaveWavePCM")
    {
        *this
            << Terminal::in("INPUT")
            << Terminal::in("SOURCES").def(nullptr)
            << Terminal::in("ENABLE").def(true)
            << Terminal::in("BASENAME").def("sep_")
            << Terminal::in("ADVANCE").def(160)
            << Terminal::in("SAMPLING_RATE").def(16000)
            << Terminal::in("BITS").def("int16")
            << Terminal::in("INPUT_BITS").def("as_BITS")
            >> Terminal::out("OUTPUT")
            >> Terminal::out("OUTPUT_FILE_NAME")
            >> Terminal::out("OUTPUT_FRAME_COUNT");
    }

    void
    SaveWavePCM::reset()
    {
    }

    void
    SaveWavePCM::initialize(InputDataBundle &input)
    {
        this->_basename      = input["BASENAME"].last().cast<std::string>();
        this->_advance       = input["ADVANCE"].last().as<size_t>();
        this->_sampling_rate = input["SAMPLING_RATE"].last().as<size_t>();

        this->_basename_has_azimuth   = (this->_basename.find("{azimuth", 0) != std::string::npos);
        this->_basename_has_elevation = (this->_basename.find("{elevation", 0) != std::string::npos);
        this->_basename_has_format = 
            this->_basename_has_azimuth || this->_basename_has_elevation
            || (this->_basename.find("{srcid", 0) != std::string::npos)
            || (this->_basename.find("{date", 0) != std::string::npos);

        auto isequal = [](const std::string &lhs, const std::string &rhs){
            return ((lhs.size() == rhs.size()) && std::equal(lhs.begin(), lhs.end(), rhs.begin(), [](const char & c1, const char & c2){
                return (c1 == c2 || std::toupper(c1) == std::toupper(c2));
            }));
        };

        std::string output_bits_name = input["BITS"].last().cast<std::string>();
        if(isequal(output_bits_name, "int16")){
            this->_output_bits = 16;
        }else if(isequal(output_bits_name, "int24")){
            this->_output_bits = 24;
        }else if(isequal(output_bits_name, "int32")){
            this->_output_bits = 32;
        }else{
            std::stringstream ss;
            ss << "SaveWavePCM: Unexpected BITS value '" << output_bits_name << "', expected int16, int24 or int32";
            throw RuntimeError(ss.str());
        }

        std::string input_bits_name = input["INPUT_BITS"].last().cast<std::string>();
        if(isequal(input_bits_name, "int16")){
            this->_input_bits = 16;
        }else if(isequal(input_bits_name, "int24")){
            this->_input_bits = 24;
        }else if(isequal(input_bits_name, "as_BITS")){
            this->_input_bits = this->_output_bits;
        }else{
            std::stringstream ss;
            ss << "SaveWavePCM: Unexpected INPUT_BITS value '" << input_bits_name << "', expected int16, int24 or as_BITS";
            throw RuntimeError(ss.str());
        }
    }

    bool
    SaveWavePCM::operator() (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        if(count == 0){
            this->initialize(input);
        }

        data input_data = input["INPUT"].last();
        bool enable = input["ENABLE"].last().as<bool>();

        //struct timeval tv;
        //gettimeofday(&tv, NULL);
        auto now = std::chrono::system_clock::now();

        // unsigned long long timestamp = tv.tv_sec * 1000000ull + tv.tv_usec;

        // インプットがMAP形式のときの処理
        if(input_data.is_type<std::map<int, data>>()){
            const auto &input_map = input_data.cast<std::map<int, data>>();

            data sources_data = input["SOURCES"].last();

            // if(sources_data.isNil() || enable){
            if(sources_data.is_empty() || enable){
                for(auto &it: input_map){
                    const auto &x = it.second.cast<std::vector<float>>();

                    std::shared_ptr<std::ostream> os;
                    auto sit = this->_streams.find(it.first);
                    if(sit == this->_streams.end()){

                        // // new ID, increment counter
                        // if(file_count < it->first){ file_count = it->first; }

                        // //src is null in this case
                        // //throw exception if azimuth or elevation is included in basename
                        // if(is_elevation || is_azimuth){
                        //     throw new NodeException(this,
                        //             string("Azimuth or elevation is indicated in BASENAME but no input terminal SOURCES."),
                        //             __FILE__, __LINE__);
                        // }
                        if(this->_basename_has_azimuth || this->_basename_has_elevation){
                            throw RuntimeError("SaveWavePCM: Basename with azimuth/elevation embedded requires SOURCES input.");
                        }

                        std::string path = format_basename_by_location(this->_basename, 0.0f, 0.0f, 0.0f, it.first, &now, this->_basename_has_format, ".wav");

                        os.reset(new std::ofstream(path.c_str(), std::ios::binary));
                        write_riff_header(*os, 1, this->_sampling_rate, this->_output_bits);

                        this->_streams[it.first] = os;
                        
                        std::vector<char> buffer = pack_pcm(x, 0, x.size(), this->_input_bits, this->_output_bits);
                        os->write(buffer.data(), buffer.size());
                        // std::cout << "output[" << it.first << "]: " << buffer.size() << " bytes" << std::endl;

                    }else{
                        os = sit->second;
                        std::vector<char> buffer = pack_pcm(x, x.size() - this->_advance, x.size(), this->_input_bits, this->_output_bits);
                        os->write(buffer.data(), buffer.size());
                        // std::cout << "output[" << it.first << "]: " << buffer.size() << " bytes, x.size()=" << x.size() << ", advance=" << this->_advance << std::endl;
                    }

                    update_riff_header(*os);
                    os->flush();
                }

            }else if(enable){
                const auto &sources = sources_data.cast<std::vector<Source>>();

                for(int i = 0; i < sources.size(); i++){
                    const auto &source = sources[i];

                    auto it = input_map.find(source.id);
                    if(it == input_map.end()){
                        continue;
                    }

                    const auto &x = it->second.cast<std::vector<float>>();

                    std::shared_ptr<std::ostream> os;
                    auto sit = this->_streams.find(source.id);
                    if(sit == this->_streams.end()){
                        std::string path = format_basename_by_source(this->_basename, source, &now, this->_basename_has_format, ".wav");

                        os.reset(new std::ofstream(path.c_str(), std::ios::binary));
                        write_riff_header(*os, 1, this->_sampling_rate, this->_output_bits);

                        this->_streams[source.id] = os;
                        std::vector<char> buffer = pack_pcm(x, 0, x.size(), this->_input_bits, this->_output_bits);
                        os->write(buffer.data(), buffer.size());

                    }else{
                        os = sit->second;
                        std::vector<char> buffer = pack_pcm(x, x.size() - this->_advance, x.size(), this->_input_bits, this->_output_bits);
                        os->write(buffer.data(), buffer.size());
                    }

                    update_riff_header(*os);
                    os->flush();
                }

            }

            if(!enable){
                // std::cout << "flush" << std::endl;
                for(auto it: this->_streams){
                    it.second->flush();
                }
                this->_streams.clear();
            }

        }else if(input_data.is_type<Matrix<float>>()){
            const auto &x = input_data.cast<Matrix<float>>();

            const size_t num_chan = x.rows();
            const size_t length = x.cols();
            if(enable){
                // foutが何も指してない＝準備が必要
                if(!this->_all_channel_stream){
                    std::string path = format_basename_by_location(this->_basename, x(0, 0), x(0, 1), x(0, 2), 0, &now, this->_basename_has_format, ".wav");
                    std::shared_ptr<std::ostream> os(new std::ofstream(path.c_str(), std::ios::binary));

                    this->_last_file_name = path;
                    this->_last_frame_count = count;

                    write_riff_header(*os, num_chan, this->_sampling_rate, this->_output_bits);

                    // 準備が終わったから、データを書き込む
                    std::vector<char> buffer = pack_pcm(x, 0, length, this->_input_bits, this->_output_bits);
                    os->write(buffer.data(), buffer.size());

                    update_riff_header(*os);

                    this->_all_channel_stream = os;

                }else{
                    std::shared_ptr<std::ostream> os = this->_all_channel_stream;

                    std::vector<char> buffer = pack_pcm(x, length - this->_advance, length, this->_input_bits, this->_output_bits);
                    os->write(buffer.data(), buffer.size());

                    update_riff_header(*os);
                }

            }else if(this->_all_channel_stream){
                this->_all_channel_stream->flush();
                this->_all_channel_stream.reset();
            }

        }else{
            throw RuntimeError("SaveWavePCM: Only map<int, data> and Matrix<float> are supported.");

        }

        output["OUTPUT"].push(input_data);
        output["OUTPUT_FILE_NAME"].push(data::create<std::string>(this->_last_file_name));
        output["OUTPUT_FRAME_COUNT"].push(data::create<int>(this->_last_frame_count));

        return true;
    }
}
}

/* end of file */
