/**
 * Wave file writing, with or without splitting by length.
 *
 * Author: Jan Schlüter <jan.schluter@lis-lab.fr>
 * Author: Maxence Ferrari <maxence.ferrari@lis-lab.fr>
 */

#include "filewriter.h"
#include <stdexcept>
#include <vector>
#include <iostream>
#include <chrono>
#include <sstream>
#include <iomanip>

FileWriter::FileWriter(std::string &filename_template, size_t num_channels, size_t sample_rate, size_t depth) :
        filename_template(filename_template), num_channels(num_channels), sample_rate(sample_rate), depth(depth) {
    // nothing
}

FileWriter::~FileWriter() {
    // nothing
}

std::string FileWriter::generate_filename() {
    // this has of course nothing to do with file writing and could be pulled
    // out, but I doubt anybody will ever care
    using namespace std::chrono;
    auto miliseconds = duration_cast<milliseconds>(std::chrono::system_clock::now().time_since_epoch()).count();
    long seconds = miliseconds/1000;
    struct tm *timeinfo = localtime(&seconds);

    size_t found = filename_template.find("%z");
    while(found != std::string::npos){
      std::stringstream ss;
      ss << std::setw(3) << std::setfill('0') << miliseconds%1000;
      std::string s = ss.str();
      filename_template.replace(found, 2, s);
      found = filename_template.find("%z", found+3);
    }

    size_t length = 0;
    size_t space = 0;
    std::vector<char> buffer;
    while (!length) {
        space += 100;
        buffer.resize(filename_template.size() + space);
        length = strftime(buffer.data(), buffer.size(),
                filename_template.c_str(), timeinfo);
    }
    return std::string(buffer.begin(), buffer.begin() + length);
}

void FileWriter::write(std::vector<uint8_t> &samples, std::vector<uint8_t> &imu_data) {
    write(samples.data(), samples.size(), imu_data.data());
}


void store_little_endian(uint8_t (&target)[2], uint16_t value) {
    target[0] = value & 0xFF;
    target[1] = (value >> 8) & 0xFF;
}

void store_little_endian(uint8_t (&target)[4], uint32_t value) {
    target[0] = value & 0xFF;
    target[1] = (value >> 8) & 0xFF;
    target[2] = (value >> 16) & 0xFF;
    target[3] = (value >> 24) & 0xFF;
}

uint32_t read_little_endian(uint8_t (&source)[4]) {
    return (source[0] + source[1] << 8 + source[2] << 16 + source[3] << 24);
}

WavFileWriter::WavFileWriter(std::string &filename_template, size_t num_channels, size_t sample_rate, size_t depth) :
        WavFileWriter(filename_template, num_channels, sample_rate, depth, 0) {
    // nothing
}

WavFileWriter::WavFileWriter(std::string &filename_template, size_t num_channels, size_t sample_rate, size_t depth,
                             size_t expected_num_samples) :
        FileWriter(filename_template, num_channels, sample_rate, depth),
        outfile(generate_filename(), std::ios::out | std::ios::binary | std::ios::trunc),
        samples_written(0) {
    // check if we could open the file
    if (!outfile.is_open()) {
        throw std::runtime_error("could not create output file");
    }
    // write header
    store_little_endian(header.fmt_channels, num_channels);
    store_little_endian(header.fmt_sample_rate, sample_rate);
    store_little_endian(header.fmt_bits_per_sample, 8 * depth);
    store_little_endian(header.fmt_byte_rate, num_channels * sample_rate * depth);
    store_little_endian(header.fmt_frame_size, num_channels * depth);
    if (expected_num_samples) {
        size_t expected_data_size = expected_num_samples * num_channels * depth;
        store_little_endian(header.data_size, expected_data_size);
        store_little_endian(header.chunk_size, expected_data_size + 36);
        // TODO: on linux, we could use fallocate to reserve the final size
    }
    // TODO: on posix, we could use posix_fadvice to indicate sequential access
    outfile.seekp(0);
    outfile.write((char*) &header, sizeof(header));
}

WavFileWriter::~WavFileWriter() {
    // finalize header, if needed
    size_t data_size = samples_written * num_channels * depth;
    if (data_size != read_little_endian(header.data_size)) {
        store_little_endian(header.data_size, data_size);
        store_little_endian(header.chunk_size, data_size + 36);
        outfile.seekp(0);
        outfile.write((char*) &header, sizeof(header));
    }
}

void WavFileWriter::write(uint8_t *samples, size_t num_samples, uint8_t *imu_data) {
    outfile.write((char*) samples, num_samples * sizeof(*samples));
    samples_written += num_samples /(num_channels * depth);
}

IMUFileWriter::IMUFileWriter(std::string &filename_template, size_t num_channels, size_t sample_rate,size_t depth, size_t timestamp) :
        FileWriter(filename_template, num_channels, sample_rate, depth),
        outfile(generate_filename(), std::ios::out | std::ios::trunc),
        last_timestamp(0) {
        outfile << header;
}

inline float le16tof(uint8_t *array){
    return static_cast<float>(static_cast<int16_t>(__builtin_bswap16(*reinterpret_cast<uint16_t*>(array))));
}

void IMUFileWriter::write(uint8_t *sample, size_t size, uint8_t *imu_data) {
    uint8_t *imu_data_cur(imu_data);
    while(imu_data_cur + frame_size + 5 < imu_data + additional_data_size){
        if(!(imu_data_cur[0]==0xFE && imu_data_cur[1]==0x0A && imu_data_cur[2]==0x0A && imu_data_cur[5]==0x08)){
          // skip trame if header is incorrect
          imu_data_cur += frame_size + 5;
          continue;
        }
        imu_data_cur += 5; // skip frame header
        auto timestamp = static_cast<int32_t>(__builtin_bswap32(*reinterpret_cast<uint32_t*>(imu_data_cur + 9)));
        if (timestamp > last_timestamp) {
            last_timestamp = timestamp;
            outfile << timestamp;
            outfile << "," << le16tof(imu_data_cur + 13) / 32756 * 19.62;   // ax resolution +- 2g
            outfile << "," << le16tof(imu_data_cur + 15) / 32756 * 19.62;   // ay resolution +- 2g
            outfile << "," << le16tof(imu_data_cur + 17) / 32756 * 19.62;   // az resolution +- 2g
            outfile << "," << le16tof(imu_data_cur + 19) / 32756 * 250;     // gx resolution +- 255deg/sec
            outfile << "," << le16tof(imu_data_cur + 21) / 32756 * 250;     // gy resolution +- 255deg/sec
            outfile << "," << le16tof(imu_data_cur + 23) / 32756 * 250;     // gz resolution +- 255deg/sec
            outfile << "," << le16tof(imu_data_cur + 25) / 32756 * 4900.;    // mx +- 4900µTesla
            outfile << "," << le16tof(imu_data_cur + 27) / 32756 * (-4900.); // my +- 4900µTesla
            outfile << "," << le16tof(imu_data_cur + 29) / 32756 * (-4900.); // mz +- 4900µTesla
            outfile << std::endl;
        }
        imu_data_cur += frame_size;
    }
}

size_t IMUFileWriter::get_last_timestamp(){
    return last_timestamp;
}


SplitWavFileWriter::SplitWavFileWriter(std::string &filename_template, size_t num_channels, size_t sample_rate,
                                       size_t depth, size_t samples_per_file) :
        FileWriter(filename_template, num_channels, sample_rate, depth),
        samples_per_file(samples_per_file),
        current_file(nullptr),
        current_file_samples_written(0) {
    // nothing
}

SplitWavFileWriter::~SplitWavFileWriter() {
    if (current_file) {
        delete current_file;
    }
}

void SplitWavFileWriter::write(uint8_t *samples, size_t num_samples, uint8_t *imu_data) {
    size_t available = num_samples /(num_channels * depth);
    // start as many new files as required to write out the samples
    while (current_file_samples_written + available >= samples_per_file) {
        if (!current_file) {
            current_file = new WavFileWriter(filename_template, num_channels, sample_rate, depth, samples_per_file);
        }
        // write out as much as fits into the current file and move the pointer
        size_t missing = samples_per_file - current_file_samples_written;
        current_file->write(samples, missing * num_channels * depth, imu_data);
        samples += missing * num_channels * depth;
        available -= missing;
        // start a new file
        delete current_file;
        current_file = nullptr;
        current_file_samples_written = 0;
    }
    // if there are samples left, write them to the current file
    if (available) {
        if (!current_file) {
            current_file = new WavFileWriter(filename_template, num_channels, sample_rate, depth, samples_per_file);
        }
        current_file->write(samples, available * num_channels * depth, imu_data);
        current_file_samples_written += available;
    }
}


SplitIMUWavFileWriter::SplitIMUWavFileWriter(std::string &filename_template, std::string &imu_name_template,
                                             size_t num_channels, size_t sample_rate, size_t depth, size_t samples_per_file):
        FileWriter(filename_template, num_channels, sample_rate, depth),
        imu_name_template(imu_name_template),
        samples_per_file(samples_per_file),
        current_file(nullptr),
        imu_file(nullptr),
        current_file_samples_written(0),
        max_timestamp(0) {
    // nothing
}

SplitIMUWavFileWriter::~SplitIMUWavFileWriter() {
    if (current_file) {
        delete current_file;
    }
    if (imu_file) {
        delete imu_file;
    }
}

void SplitIMUWavFileWriter::write(uint8_t *samples, size_t num_samples, uint8_t* imu_data) {
    size_t available = num_samples /(num_channels * depth);
    // start as many new files as required to write out the samples
    while (current_file_samples_written + available >= samples_per_file) {
        if (!current_file) {
            current_file = new WavFileWriter(filename_template, num_channels, sample_rate, depth, samples_per_file);
            if (imu_file) {
                max_timestamp = imu_file->get_last_timestamp();
                delete imu_file;
                imu_file = new IMUFileWriter(imu_name_template, num_channels, sample_rate, depth, max_timestamp);
            }
        }
        if (!imu_file) imu_file = new IMUFileWriter(imu_name_template, num_channels, sample_rate, depth, max_timestamp);
        // write out as much as fits into the current file and move the pointer
        size_t missing = samples_per_file - current_file_samples_written;
        current_file->write(samples, missing * num_channels * depth, imu_data);
        if (imu_data) imu_file->write(samples, missing * num_channels * depth, imu_data);
        imu_data = nullptr; // prevent multiple writing
        samples += missing * num_channels * depth;
        available -= missing;
        // start a new file
        delete current_file;
        max_timestamp = imu_file->get_last_timestamp();
        delete imu_file;
        current_file = nullptr;
        imu_file = nullptr;
        current_file_samples_written = 0;
    }
    // if there are samples left, write them to the current file
    if (available) {
        if (!current_file) {
            current_file = new WavFileWriter(filename_template, num_channels, sample_rate, depth, samples_per_file);
            if (imu_file) {
                max_timestamp = imu_file->get_last_timestamp();
                delete imu_file;
                imu_file = new IMUFileWriter(imu_name_template, num_channels, sample_rate, depth, max_timestamp);
            }
        }
        if (!imu_file) imu_file = new IMUFileWriter(imu_name_template, num_channels, sample_rate, depth, max_timestamp);
        current_file->write(samples, available * num_channels * depth, imu_data);
        if (imu_data) imu_file->write(samples, available * num_channels * depth, imu_data);
        imu_data = nullptr;
        current_file_samples_written += available;
    }

}