vpIoTools_npy.cpp

/*
 * ViSP, open source Visual Servoing Platform software.
 * Copyright (C) 2005 - 2024 by Inria. All rights reserved.
 *
 * This software is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * See the file LICENSE.txt at the root directory of this source
 * distribution for additional information about the GNU GPL.
 *
 * For using ViSP with software that can not be combined with the GNU
 * GPL, please contact Inria about acquiring a ViSP Professional
 * Edition License.
 *
 * See https://visp.inria.fr for more information.
 *
 * This software was developed at:
 * Inria Rennes - Bretagne Atlantique
 * Campus Universitaire de Beaulieu
 * 35042 Rennes Cedex
 * France
 *
 * If you have questions regarding the use of this file, please contact
 * Inria at visp@inria.fr
 *
 * This file is provided AS IS with NO WARRANTY OF ANY KIND, INCLUDING THE
 * WARRANTY OF DESIGN, MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 *
 * Description:
 * Io tools dedicated to npy.
 */
 
#include <visp3/core/vpConfig.h>
#include <visp3/core/vpIoTools.h>
 
#if defined(VISP_HAVE_MINIZ) && defined(VISP_HAVE_WORKING_REGEX)
#define USE_ZLIB_API 0
 
#if !USE_ZLIB_API
// See: https://github.com/BinomialLLC/basis_universal/blob/master/encoder/basisu_miniz.h
// Apache License, Version 2.0
#include "basisu_miniz.h"
 
using namespace buminiz;
#else
#include <zlib.h>
#endif
 
// To avoid warnings such as: warning: unused variable ‘littleEndian’ [-Wunused-variable]
#define UNUSED(x) ((void)(x)) // see: https://stackoverflow.com/a/777359
 
// Copyright (C) 2011  Carl Rogers
// Released under MIT License
// license available in LICENSE file, or at http://www.opensource.org/licenses/mit-license.php
 
 
// anonymous namespace
namespace
{
#if defined(ENABLE_VISP_NAMESPACE)
using namespace VISP_NAMESPACE_NAME; // for vpEndian calls
#endif
 
void reverse_data(std::shared_ptr<std::vector<char> > &data_holder, const std::vector<size_t> &shape,
  size_t word_size, char data_type)
{
  if (!shape.empty()) {
    size_t total_size = shape[0];
    for (size_t i = 1; i < shape.size(); i++) {
      total_size *= shape[i];
    }
 
    // std::complex
    if (data_type == 'c') {
      const size_t half_word_size = word_size / 2;
      for (size_t i = 0; i < total_size; i++) {
        // real
        std::reverse(data_holder->begin() + i*word_size, data_holder->begin() + i*word_size + half_word_size);
        // imag
        std::reverse(data_holder->begin() + i*word_size + half_word_size, data_holder->begin() + (i+1)*word_size);
      }
    }
    else if (data_type != 'U') {
      for (size_t i = 0; i < total_size; i++) {
        std::reverse(data_holder->begin() + i*word_size, data_holder->begin() + (i+1)*word_size);
      }
    }
  }
  else {
    // 0d-array (number) with NumPy
    if (data_type == 'c') {
      const size_t half_word_size = word_size / 2;
      // real
      std::reverse(data_holder->begin(), data_holder->begin() + half_word_size);
      // imag
      std::reverse(data_holder->begin() + half_word_size, data_holder->begin() + word_size);
    }
    else if (data_type != 'U') {
      std::reverse(data_holder->begin(), data_holder->begin() + word_size);
    }
  }
 
  if (data_type == 'U') { // special case to handle UTF-32 string data
    size_t utf32_size = 4;
    for (size_t i = 0; i < data_holder->size(); i += utf32_size) {
      std::reverse(data_holder->begin() + i, data_holder->begin() + (i+utf32_size)); // NumPy saves string in UTF-32
    }
  }
}
 
uint16_t swap16bits_if(uint16_t val, bool swap)
{
  if (swap) {
    return vpEndian::swap16bits(val);
  }
 
  return val;
}
 
uint32_t swap32bits_if(uint32_t val, bool swap)
{
  if (swap) {
    return vpEndian::swap32bits(val);
  }
 
  return val;
}
 
// https://github.com/rogersce/cnpy/pull/78/files
// https://github.com/nmcclatchey/Bugfix-for-cnpy/blob/e148a5ce5db80fa3e28ce6d551343cfc8ebdc832/cnpy.cpp#L285
struct AutoCloser
{
  FILE *fp;
  AutoCloser() : fp(nullptr) { }
  ~AutoCloser()
  {
    if (fp != nullptr) {
      fclose(fp);
    }
  }
};
} // anonymous namespace
 
char visp::cnpy::BigEndianTest()
{
  int x = 1;
  return (((reinterpret_cast<char *>(&x))[0]) ? '<' : '>');
}
 
char visp::cnpy::map_type(const std::type_info &t)
{
  if (t == typeid(float)) { return 'f'; }
  if (t == typeid(double)) { return 'f'; }
  if (t == typeid(long double)) { return 'f'; }
 
  if (t == typeid(int)) { return 'i'; }
  if (t == typeid(char)) { return 'i'; }
  if (t == typeid(short)) { return 'i'; }
  if (t == typeid(long)) { return 'i'; }
  if (t == typeid(long long)) { return 'i'; }
 
  if (t == typeid(unsigned char)) { return 'u'; }
  if (t == typeid(unsigned short)) { return 'u'; }
  if (t == typeid(unsigned long)) { return 'u'; }
  if (t == typeid(unsigned long long)) { return 'u'; }
  if (t == typeid(unsigned int)) { return 'u'; }
 
  if (t == typeid(bool)) { return 'b'; }
 
  if (t == typeid(std::complex<float>)) { return 'c'; }
  if (t == typeid(std::complex<double>)) { return 'c'; }
  if (t == typeid(std::complex<long double>)) { return 'c'; }
 
  if (t == typeid(std::string)) { return 'U'; }
 
  else { return '?'; }
}
 
void visp::cnpy::parse_npy_header(unsigned char *buffer, size_t &word_size, std::vector<size_t> &shape,
  bool &fortran_order, bool &little_endian, char &data_type)
{
  uint16_t header_len = *reinterpret_cast<uint16_t *>(buffer+8);
  std::string header(reinterpret_cast<char *>(buffer+9), header_len);
 
  //fortran order
  size_t loc1 = header.find("fortran_order")+16;
  fortran_order = (header.substr(loc1, 4) == "True" ? true : false);
 
  //shape
  loc1 = header.find("(");
  size_t loc2 = header.find(")");
 
  std::regex num_regex("[0-9][0-9]*");
  std::smatch sm;
  shape.clear();
 
  std::string str_shape = header.substr(loc1+1, loc2-loc1-1);
  while (std::regex_search(str_shape, sm, num_regex)) {
    // https://github.com/rogersce/cnpy/commit/ca6c0ce5bed57e3b5b64aede4f39aa07a9e71f5e
    shape.push_back(std::stoll(sm[0].str()));
    str_shape = sm.suffix().str();
  }
 
  //endian, word size, data type
  //byte order code | stands for not applicable.
  //not sure when this applies except for byte array
  loc1 = header.find("descr")+9;
  little_endian = ((header[loc1] == '<') || (header[loc1] == '|') ? true : false);
  data_type = header[loc1+1];
 
  std::string str_ws = header.substr(loc1+2);
  loc2 = str_ws.find("'");
  word_size = atoll(str_ws.substr(0, loc2).c_str());
}
 
void visp::cnpy::parse_npy_header(FILE *fp, size_t &word_size, std::vector<size_t> &shape,
  bool &fortran_order, bool &little_endian, char &data_type)
{
  char buffer[256];
  size_t res = fread(buffer, sizeof(char), 11, fp);
  if (res != 11) {
    std::ostringstream oss;
    oss << "parse_npy_header: failed fread, res=" << res;
    throw std::runtime_error(oss.str());
  }
  std::string header = fgets(buffer, 256, fp);
  assert(header[header.size()-1] == '\n');
 
  size_t loc1, loc2;
 
  //fortran order
  loc1 = header.find("fortran_order");
  if (loc1 == std::string::npos) {
    throw std::runtime_error("parse_npy_header: failed to find header keyword: 'fortran_order'");
  }
  loc1 += 16;
  fortran_order = (header.substr(loc1, 4) == "True" ? true : false);
 
  //shape
  loc1 = header.find("(");
  loc2 = header.find(")");
  if ((loc1 == std::string::npos) || (loc2 == std::string::npos)) {
    throw std::runtime_error("parse_npy_header: failed to find header keyword: '(' or ')'");
  }
 
  std::regex num_regex("[0-9][0-9]*");
  std::smatch sm;
  shape.clear();
 
  std::string str_shape = header.substr(loc1+1, loc2-loc1-1);
  while (std::regex_search(str_shape, sm, num_regex)) {
    // https://github.com/rogersce/cnpy/commit/ca6c0ce5bed57e3b5b64aede4f39aa07a9e71f5e
    shape.push_back(std::stoll(sm[0].str()));
    str_shape = sm.suffix().str();
  }
 
  //endian, word size, data type
  //byte order code | stands for not applicable.
  //not sure when this applies except for byte array
  loc1 = header.find("descr");
  if (loc1 == std::string::npos) {
    throw std::runtime_error("parse_npy_header: failed to find header keyword: 'descr'");
  }
  loc1 += 9;
  little_endian = ((header[loc1] == '<') || (header[loc1] == '|') ? true : false);
  data_type = header[loc1+1];
 
  std::string str_ws = header.substr(loc1+2);
  loc2 = str_ws.find("'");
  word_size = atoll(str_ws.substr(0, loc2).c_str());
  if (data_type == 'U') {
    word_size *= 4; // UTF-32 with NumPy
  }
}
 
void visp::cnpy::parse_zip_footer(FILE *fp, uint16_t &nrecs, size_t &global_header_size, size_t &global_header_offset)
{
  std::vector<char> footer(22);
  fseek(fp, -22, SEEK_END);
  size_t res = fread(&footer[0], sizeof(char), 22, fp);
  if (res != 22) {
    std::ostringstream oss;
    oss << "parse_zip_footer: failed fread, res=" << res;
    throw std::runtime_error(oss.str());
  }
 
  uint16_t disk_no, disk_start, nrecs_on_disk, comment_len;
#ifdef VISP_BIG_ENDIAN
  disk_no = vpEndian::swap16bits(*(uint16_t *)&footer[4]);
  disk_start = vpEndian::swap16bits(*(uint16_t *)&footer[6]);
  nrecs_on_disk = vpEndian::swap16bits(*(uint16_t *)&footer[8]);
  nrecs = vpEndian::swap16bits(*(uint16_t *)&footer[10]);
  global_header_size = vpEndian::swap32bits(*(uint32_t *)&footer[12]);
  global_header_offset = vpEndian::swap32bits(*(uint32_t *)&footer[16]);
  comment_len = vpEndian::swap16bits(*(uint16_t *)&footer[20]);
#else
  disk_no = *(uint16_t *)&footer[4];
  disk_start = *(uint16_t *)&footer[6];
  nrecs_on_disk = *(uint16_t *)&footer[8];
  nrecs = *(uint16_t *)&footer[10];
  global_header_size = *(uint32_t *)&footer[12];
  global_header_offset = *(uint32_t *)&footer[16];
  comment_len = *(uint16_t *)&footer[20];
#endif
 
  UNUSED(disk_no); assert(disk_no == 0);
  UNUSED(disk_start); assert(disk_start == 0);
  UNUSED(nrecs_on_disk); assert(nrecs_on_disk == nrecs);
  UNUSED(comment_len); assert(comment_len == 0);
}
 
visp::cnpy::NpyArray load_the_npy_file(FILE *fp)
{
  std::vector<size_t> shape;
  size_t word_size;
  bool fortran_order, little_endian;
  char data_type = 'i'; // integer type
  visp::cnpy::parse_npy_header(fp, word_size, shape, fortran_order, little_endian, data_type);
 
  visp::cnpy::NpyArray arr(shape, word_size, fortran_order, data_type);
  size_t nread = fread(arr.data<char>(), 1, arr.num_bytes(), fp);
  if (nread != arr.num_bytes()) {
    std::ostringstream oss;
    oss << "load_the_npy_file: failed fread, nread=" << nread << " ; num_bytes=" << arr.num_bytes();
    throw std::runtime_error(oss.str());
  }
 
#ifdef VISP_LITTLE_ENDIAN
  if (!little_endian) {
    reverse_data(arr.data_holder, arr.shape, arr.word_size, data_type);
  }
#else
  if (little_endian) {
    reverse_data(arr.data_holder, arr.shape, arr.word_size, data_type);
  }
#endif
  return arr;
}
 
visp::cnpy::NpyArray load_the_npz_array(FILE *fp, uint32_t compr_bytes, uint32_t uncompr_bytes)
{
  std::vector<unsigned char> buffer_compr(compr_bytes);
  std::vector<unsigned char> buffer_uncompr(uncompr_bytes);
  size_t nread = fread(&buffer_compr[0], 1, compr_bytes, fp);
  if (nread != compr_bytes) {
    std::ostringstream oss;
    oss << "load_the_npz_array: failed fread, nread=" << nread << " ; compr_bytes=" << compr_bytes;
    throw std::runtime_error(oss.str());
  }
 
  z_stream d_stream;
 
  d_stream.zalloc = Z_NULL;
  d_stream.zfree = Z_NULL;
  d_stream.opaque = Z_NULL;
  d_stream.avail_in = 0;
  d_stream.next_in = Z_NULL;
  int err = inflateInit2(&d_stream, -MAX_WBITS);
  // https://github.com/rogersce/cnpy/commit/3ed2bc4063c455269b37af63442c595ee1bd60e1
  if (err != Z_OK) {
    std::ostringstream oss;
    oss << "load_the_npz_array: zlib inflateInit2 failed ; err=" << err;
    throw std::runtime_error(oss.str());
  }
 
  d_stream.avail_in = compr_bytes;
  d_stream.next_in = &buffer_compr[0];
  d_stream.avail_out = uncompr_bytes;
  d_stream.next_out = &buffer_uncompr[0];
 
  err = inflate(&d_stream, Z_FINISH);
  if (err != Z_OK) {
    std::ostringstream oss;
    oss << "load_the_npz_array: zlib inflate failed ; err=" << err;
    throw std::runtime_error(oss.str());
  }
  err = inflateEnd(&d_stream);
  if (err != Z_OK) {
    std::ostringstream oss;
    oss << "load_the_npz_array: zlib inflateEnd failed ; err=" << err;
    throw std::runtime_error(oss.str());
  }
 
  std::vector<size_t> shape;
  size_t word_size;
  bool fortran_order;
  bool little_endian = true;
  char data_type = 'i'; // integer type
  visp::cnpy::parse_npy_header(&buffer_uncompr[0], word_size, shape, fortran_order, little_endian, data_type);
 
  visp::cnpy::NpyArray array(shape, word_size, fortran_order, data_type);
 
  size_t offset = uncompr_bytes - array.num_bytes();
  memcpy(array.data<unsigned char>(), &buffer_uncompr[0]+offset, array.num_bytes());
 
#ifdef VISP_LITTLE_ENDIAN
  if (!little_endian) {
    reverse_data(array.data_holder, array.shape, array.word_size, data_type);
  }
#else
  if (little_endian) {
    reverse_data(array.data_holder, array.shape, array.word_size, data_type);
  }
#endif
 
  return array;
}

visp::cnpy::npz_t visp::cnpy::npz_load(const std::string &fname)
{
  if (!vpIoTools::checkFilename(fname)) {
    throw vpException(vpException::ioError, "This file does not exist: " + fname);
  }
 
  AutoCloser closer;
  closer.fp = fopen(fname.c_str(), "rb");
 
  if (!closer.fp) {
    throw std::runtime_error("npz_load: Error! Unable to open file " + fname + "!");
  }
 
  visp::cnpy::npz_t arrays;
  bool quit = false;
  const unsigned int index_2 = 2;
  const unsigned int index_3 = 3;
  const unsigned int index_26 = 26;
  const unsigned int index_28 = 28;
  const unsigned int val_8 = 8;
  const unsigned int val_18 = 18;
  const unsigned int val_22 = 22;
  const unsigned int val_30 = 30;
 
  bool host_is_LE = true;
#ifndef VISP_LITTLE_ENDIAN
  host_is_LE = false;
#endif
 
  const bool header_file_is_LE = true;
  bool same_endianness = (host_is_LE == header_file_is_LE);
  while (!quit) {
    std::vector<char> local_header(val_30);
    size_t headerres = fread(&local_header[0], sizeof(char), val_30, closer.fp);
    if (headerres != 30) {
      throw std::runtime_error("npz_load: failed fread 1");
    }
 
    //if we've reached the global header, stop reading
    if ((local_header[index_2] != 0x03) || (local_header[index_3] != 0x04)) {
      quit = true;
    }
    else {
      //read in the variable name
      uint16_t name_len = swap16bits_if(*(uint16_t *)&local_header[index_26], !same_endianness);
      std::string varname(name_len, ' ');
      size_t vname_res = fread(&varname[0], sizeof(char), name_len, closer.fp);
      if (vname_res != name_len) {
        throw std::runtime_error("npz_load: failed fread 2");
      }
 
      //erase the lagging .npy
      varname.erase(varname.end()-4, varname.end());
 
      //read in the extra field
      uint16_t extra_field_len = swap16bits_if(*(uint16_t *)&local_header[index_28], !same_endianness);
      if (extra_field_len > 0) {
        std::vector<char> buff(extra_field_len);
        size_t efield_res = fread(&buff[0], sizeof(char), extra_field_len, closer.fp);
        if (efield_res != extra_field_len) {
          throw std::runtime_error("npz_load: failed fread 3");
        }
      }
 
      uint16_t compr_method = swap16bits_if(*reinterpret_cast<uint16_t *>(&local_header[0] + val_8), !same_endianness);
      uint32_t compr_bytes = swap32bits_if(*reinterpret_cast<uint32_t *>(&local_header[0] + val_18), !same_endianness);
      uint32_t uncompr_bytes = swap32bits_if(*reinterpret_cast<uint32_t *>(&local_header[0] + val_22), !same_endianness);
 
      if (compr_method == 0) {
        arrays[varname] = load_the_npy_file(closer.fp);
      }
      else {
        arrays[varname] = load_the_npz_array(closer.fp, compr_bytes, uncompr_bytes);
      }
    }
  }
 
  return arrays;
}

visp::cnpy::NpyArray visp::cnpy::npz_load(const std::string &fname, const std::string &varname)
{
  if (!vpIoTools::checkFilename(fname)) {
    throw vpException(vpException::ioError, "This file does not exist: " + fname);
  }
 
  AutoCloser closer;
  closer.fp = fopen(fname.c_str(), "rb");
 
  if (!closer.fp) {
    throw std::runtime_error("npz_load: Unable to open file " + fname + "!");
  }
 
  bool quit = false;
  const unsigned int index_2 = 2;
  const unsigned int index_3 = 3;
  const unsigned int index_26 = 26;
  const unsigned int index_28 = 28;
  const unsigned int val_8 = 8;
  const unsigned int val_18 = 18;
  const unsigned int val_22 = 22;
  const unsigned int val_30 = 30;
 
  bool host_is_LE = true;
#ifndef VISP_LITTLE_ENDIAN
  host_is_LE = false;
#endif
 
  const bool header_file_is_LE = true;
  bool same_endianness = (host_is_LE == header_file_is_LE);
  while (!quit) {
    std::vector<char> local_header(val_30);
    size_t header_res = fread(&local_header[0], sizeof(char), val_30, closer.fp);
    if (header_res != 30) {
      throw std::runtime_error("npz_load 2: failed fread");
    }
 
    //if we've reached the global header, stop reading
    if ((local_header[index_2] != 0x03) || (local_header[index_3] != 0x04)) {
      quit = true;
    }
    else {
      //read in the variable name
      uint16_t name_len = swap16bits_if(*(uint16_t *)&local_header[index_26], !same_endianness);
      std::string vname(name_len, ' ');
      size_t vname_res = fread(&vname[0], sizeof(char), name_len, closer.fp);
      if (vname_res != name_len) {
        throw std::runtime_error("npz_load 2: failed fread");
      }
      vname.erase(vname.end()-4, vname.end()); //erase the lagging .npy
 
      //read in the extra field
      uint16_t extra_field_len = swap16bits_if(*(uint16_t *)&local_header[index_28], !same_endianness);
      fseek(closer.fp, extra_field_len, SEEK_CUR); //skip past the extra field
 
      uint16_t compr_method = swap16bits_if(*reinterpret_cast<uint16_t *>(&local_header[0] + val_8), !same_endianness);
      uint32_t compr_bytes = swap32bits_if(*reinterpret_cast<uint32_t *>(&local_header[0] + val_18), !same_endianness);
      uint32_t uncompr_bytes = swap32bits_if(*reinterpret_cast<uint32_t *>(&local_header[0] + val_22), !same_endianness);
 
      if (vname == varname) {
        NpyArray array = (compr_method == 0) ? load_the_npy_file(closer.fp) : load_the_npz_array(closer.fp, compr_bytes, uncompr_bytes);
        return array;
      }
      else {
        //skip past the data
        uint32_t size = swap32bits_if(*(uint32_t *)&local_header[22], !same_endianness);
        fseek(closer.fp, size, SEEK_CUR);
      }
    }
  }
 
  //if we get here, we haven't found the variable in the file
  throw std::runtime_error("npz_load 2: Variable name " + varname + " not found in " + fname);
}

visp::cnpy::NpyArray visp::cnpy::npy_load(const std::string &fname)
{
  if (!vpIoTools::checkFilename(fname)) {
    throw vpException(vpException::ioError, "This file does not exist: " + fname);
  }
 
  AutoCloser closer;
  closer.fp = fopen(fname.c_str(), "rb");
 
  if (!closer.fp) {
    throw std::runtime_error("npy_load: Unable to open file " + fname + "!");
  }
 
  NpyArray arr = load_the_npy_file(closer.fp);
 
  return arr;
}
 
namespace visp
{
namespace cnpy
{
// 000000e0  ff ff 12 00 14 00 4d 79  20 73 74 72 69 6e 67 20  |......My string |
// 000000f0  64 61 74 61 2e 6e 70 79  01 00 10 00 8c 00 00 00  |data.npy........|
// 00000100  00 00 00 00 8c 00 00 00  00 00 00 00 93 4e 55 4d  |.............NUM|
// 00000110  50 59 01 00 76 00 7b 27  64 65 73 63 72 27 3a 20  |PY..v.{'descr': |
// 00000120  27 3c 55 33 27 2c 20 27  66 6f 72 74 72 61 6e 5f  |'<U3', 'fortran_|
// 00000130  6f 72 64 65 72 27 3a 20  46 61 6c 73 65 2c 20 27  |order': False, '|
// 00000140  73 68 61 70 65 27 3a 20  28 29 2c 20 7d 20 20 20  |shape': (), }   |
// 00000150  20 20 20 20 20 20 20 20  20 20 20 20 20 20 20 20  |                |
// 00000160  20 20 20 20 20 20 20 20  20 20 20 20 20 20 20 20  |                |
// 00000170  20 20 20 20 20 20 20 20  20 20 20 20 20 20 20 20  |                |
// 00000180  20 20 20 20 20 20 20 20  20 20 20 0a 61 00 00 00  |           .a...|
// 00000190  62 00 00 00 63 00 00 00  50 4b 03 04 2d 00 00 00  |b...c...PK..-...|
// 000001a0  00 00 00 00 21 00 76 38  d2 58 ff ff ff ff ff ff  |....!.v8.X......|
// 000001b0  ff ff 16 00 14 00 4d 79  20 76 65 63 20 73 74 72  |......My vec str|
// 000001c0  69 6e 67 20 64 61 74 61  2e 6e 70 79 01 00 10 00  |ing data.npy....|
// 000001d0  40 01 00 00 00 00 00 00  40 01 00 00 00 00 00 00  |@.......@.......|
// 000001e0  93 4e 55 4d 50 59 01 00  76 00 7b 27 64 65 73 63  |.NUMPY..v.{'desc|
// 000001f0  72 27 3a 20 27 3c 55 38  27 2c 20 27 66 6f 72 74  |r': '<U8', 'fort|
// 00000200  72 61 6e 5f 6f 72 64 65  72 27 3a 20 46 61 6c 73  |ran_order': Fals|
// 00000210  65 2c 20 27 73 68 61 70  65 27 3a 20 28 32 2c 20  |e, 'shape': (2, |
// 00000220  33 29 2c 20 7d 20 20 20  20 20 20 20 20 20 20 20  |3), }           |
// 00000230  20 20 20 20 20 20 20 20  20 20 20 20 20 20 20 20  |                |
// 00000240  20 20 20 20 20 20 20 20  20 20 20 20 20 20 20 20  |                |
// 00000250  20 20 20 20 20 20 20 20  20 20 20 20 20 20 20 0a  |               .|
// 00000260  61 00 00 00 62 00 00 00  63 00 00 00 64 00 00 00  |a...b...c...d...|
// 00000270  20 00 00 00 00 00 00 00  00 00 00 00 00 00 00 00  | ...............|
// 00000280  65 00 00 00 66 00 00 00  67 00 00 00 68 00 00 00  |e...f...g...h...|
// 00000290  69 00 00 00 6a 00 00 00  6b 00 00 00 6c 00 00 00  |i...j...k...l...|
// 000002a0  6d 00 00 00 6e 00 00 00  00 00 00 00 00 00 00 00  |m...n...........|
// 000002b0  00 00 00 00 00 00 00 00  00 00 00 00 00 00 00 00  |................|
// 000002c0  6f 00 00 00 00 00 00 00  00 00 00 00 00 00 00 00  |o...............|
// 000002d0  00 00 00 00 00 00 00 00  00 00 00 00 00 00 00 00  |................|
// 000002e0  70 00 00 00 71 00 00 00  00 00 00 00 00 00 00 00  |p...q...........|
// 000002f0  00 00 00 00 00 00 00 00  00 00 00 00 00 00 00 00  |................|
// 00000300  72 00 00 00 73 00 00 00  74 00 00 00 00 00 00 00  |r...s...t.......|
// 00000310  00 00 00 00 00 00 00 00  00 00 00 00 00 00 00 00  |................|
// 00000320  50 4b 01 02 2d 03 2d 00  00 00 00 00 00 00 21 00  |PK..-.-.......!.|
// 00000330  4b f2 4d a7 84 00 00 00  84 00 00 00 12 00 00 00  |K.M.............|
//
// Create header for special string case
std::vector<char> create_npy_header_string(const std::vector<size_t> &shape, const std::vector<size_t> &lengths)
{
  size_t max_length = *std::max_element(lengths.begin(), lengths.end());
 
  std::vector<char> dict;
  dict += "{'descr': '";
  dict += BigEndianTest();
  dict += map_type(typeid(std::string));
  dict += std::to_string(max_length);
  if (shape.size() > 0) {
    dict += "', 'fortran_order': False, 'shape': (";
    dict += std::to_string(shape[0]);
    for (size_t i = 1; i < shape.size(); ++i) {
      dict += ", ";
      dict += std::to_string(shape[i]);
    }
    if (shape.size() == 1) dict += ",";
    dict += "), }";
  }
  else {
    dict += "', 'fortran_order': False, 'shape': (";
    dict += "), }";
  }
  //pad with spaces so that preamble+dict is modulo 16 bytes. preamble is 10 bytes. dict needs to end with \n
  int remainder = 16 - (10 + dict.size()) % 16;
  dict.insert(dict.end(), remainder, ' ');
  dict.back() = '\n';
 
  std::vector<char> header;
  header += static_cast<char>(0x93);
  header += "NUMPY";
  header += static_cast<char>(0x01); //major version of numpy format
  header += static_cast<char>(0x00); //minor version of numpy format
#ifdef VISP_BIG_ENDIAN
  header += vpEndian::swap16bits(static_cast<uint16_t>(dict.size()));
#else
  header += static_cast<uint16_t>(dict.size());
#endif
  header.insert(header.end(), dict.begin(), dict.end());
 
  return header;
}
 
// std::vector<char> utf8_to_utf32_vec(const std::string &utf8)
// {
//   // SO: https://stackoverflow.com/questions/52703630/convert-c-stdstring-to-utf-16-le-encoded-string/52703954#52703954
//   // SO and ChatGPT gives approximatively the same logic
 
//   // But since I don't know how C++ deals with utf-16/utf-32 and std::string and for simplicity we do something stupid
//   // instead (see utf8_to_utf32_vec() func below).
 
//   std::wstring_convert<std::codecvt_utf8<char32_t, 1114111UL, std::codecvt_mode::little_endian>, char32_t> cnv;
//   std::u32string s = cnv.from_bytes(utf8);
//   if (cnv.converted() < utf8.size()) {
//     throw std::runtime_error("incomplete conversion");
//   }
 
//   std::vector<char> utf32Vector;
//   for (char32_t uc : s) {
//     utf32Vector.push_back(static_cast<char>(uc & 0xFF));
//     utf32Vector.push_back(static_cast<char>((uc >> 8) & 0xFF));
//     utf32Vector.push_back(static_cast<char>((uc >> 16) & 0xFF));
//     utf32Vector.push_back(static_cast<char>((uc >> 24) & 0xFF));
//   }
 
//   return utf32Vector;
// }
 
// NumPy saves a string character using 32-bits:
// - https://github.com/numpy/numpy/issues/15347
// Here for simplicity we assume std::string is utf-8 and we simply pad the vector with 0-val (NULL character).
std::vector<char> utf8_to_utf32_vec_pad(const std::string &utf8, const std::size_t &max_size)
{
  std::vector<char> utf32Vector;
  utf32Vector.resize(4*max_size, 0);
 
  for (size_t i = 0, idx = 0; i < utf8.length(); i++, idx += 4) {
#ifdef VISP_BIG_ENDIAN
    utf32Vector[idx+3] = utf8[i];
#else
    utf32Vector[idx] = utf8[i];
#endif
  }
 
  return utf32Vector;
}
} // cnpy
} // visp

void visp::cnpy::npz_save(const std::string &zipname, std::string fname, const std::vector<std::string> &data_vec, const std::vector<size_t> &shape, const std::string &mode)
{
  if (data_vec.empty()) {
    vpException(vpException::badValue, "Input string data is empty.");
  }
  if (shape.empty()) {
    vpException(vpException::dimensionError, "Input string data shape vec is empty.");
  }
 
  //first, append a .npy to the fname
  fname += ".npy";
 
  //now, on with the show
  FILE *fp = NULL;
  uint16_t nrecs = 0;
  size_t global_header_offset = 0;
  std::vector<char> global_header;
 
  if (mode == "a") fp = fopen(zipname.c_str(), "r+b");
 
  if (fp) {
    //zip file exists. we need to add a new npy file to it.
    //first read the footer. this gives us the offset and size of the global header
    //then read and store the global header.
    //below, we will write the the new data at the start of the global header then append the global header and footer below it
    size_t global_header_size;
    parse_zip_footer(fp, nrecs, global_header_size, global_header_offset);
    fseek(fp, static_cast<long>(global_header_offset), SEEK_SET);
    global_header.resize(global_header_size);
    size_t res = fread(&global_header[0], sizeof(char), global_header_size, fp);
    if (res != global_header_size) {
      throw std::runtime_error("npz_save: header read error while adding to existing zip");
    }
    fseek(fp, static_cast<long>(global_header_offset), SEEK_SET);
  }
  else {
    fp = fopen(zipname.c_str(), "wb");
  }
 
  std::vector<size_t> lengths;
  lengths.reserve(data_vec.size());
  for (auto data_str : data_vec) {
    lengths.push_back(data_str.length());
  }
 
  std::vector<char> npy_header = create_npy_header_string(shape, lengths);
  size_t max_length = *std::max_element(lengths.begin(), lengths.end());
 
  std::vector<char> data_str_utf32_LE;
  data_str_utf32_LE.reserve(max_length*4); // NumPy stores string as UTF-32: https://github.com/numpy/numpy/issues/15347
  for (size_t i = 0; i < lengths.size(); i++) {
    std::vector<char> substr_utf32 = utf8_to_utf32_vec_pad(data_vec[i], max_length);
    data_str_utf32_LE.insert(data_str_utf32_LE.end(), substr_utf32.begin(), substr_utf32.end());
  }
 
  // https://github.com/rogersce/cnpy/pull/58/files
  size_t nels = data_str_utf32_LE.size();
  size_t nbytes = nels*sizeof(char) + npy_header.size();
 
  //get the CRC of the data to be added
  uint32_t crc = vp_mz_crc32(0L, (uint8_t *)&npy_header[0], npy_header.size());
  if (nels > 0) {
    crc = vp_mz_crc32(crc, (uint8_t *)&data_str_utf32_LE[0], nels*sizeof(uint8_t));
  }
 
  //build the local header
  std::vector<char> local_header;
  local_header += "PK"; //first part of sig
#ifdef VISP_BIG_ENDIAN
  local_header += vpEndian::swap16bits(static_cast<uint16_t>(0x0403)); //second part of sig
  local_header += vpEndian::swap16bits(static_cast<uint16_t>(20)); //min version to extract
  local_header += vpEndian::swap16bits(static_cast<uint16_t>(0)); //general purpose bit flag
  local_header += vpEndian::swap16bits(static_cast<uint16_t>(0)); //compression method
  local_header += vpEndian::swap16bits(static_cast<uint16_t>(0)); //file last mod time
  local_header += vpEndian::swap16bits(static_cast<uint16_t>(0));     //file last mod date
  local_header += vpEndian::swap32bits(static_cast<uint32_t>(crc)); //crc
  local_header += vpEndian::swap32bits(static_cast<uint32_t>(nbytes)); //compressed size
  local_header += vpEndian::swap32bits(static_cast<uint32_t>(nbytes)); //uncompressed size
  local_header += vpEndian::swap16bits(static_cast<uint16_t>(fname.size())); //fname length
  local_header += vpEndian::swap16bits(static_cast<uint16_t>(0)); //extra field length
#else
  local_header += static_cast<uint16_t>(0x0403); //second part of sig
  local_header += static_cast<uint16_t>(20); //min version to extract
  local_header += static_cast<uint16_t>(0); //general purpose bit flag
  local_header += static_cast<uint16_t>(0); //compression method
  local_header += static_cast<uint16_t>(0); //file last mod time
  local_header += static_cast<uint16_t>(0);     //file last mod date
  local_header += static_cast<uint32_t>(crc); //crc
  local_header += static_cast<uint32_t>(nbytes); //compressed size
  local_header += static_cast<uint32_t>(nbytes); //uncompressed size
  local_header += static_cast<uint16_t>(fname.size()); //fname length
  local_header += static_cast<uint16_t>(0); //extra field length
#endif
  local_header += fname;
 
  //build global header
  global_header += "PK"; //first part of sig
#ifdef VISP_BIG_ENDIAN
  global_header += vpEndian::swap16bits(static_cast<uint16_t>(0x0201)); //second part of sig
  global_header += vpEndian::swap16bits(static_cast<uint16_t>(20)); //version made by
  global_header.insert(global_header.end(), local_header.begin()+4, local_header.begin()+30);
  global_header += static_cast<uint16_t>(0); //file comment length
  global_header += static_cast<uint16_t>(0); //disk number where file starts
  global_header += static_cast<uint16_t>(0); //internal file attributes
  global_header += static_cast<uint32_t>(0); //external file attributes
  global_header += vpEndian::swap32bits(static_cast<uint32_t>(global_header_offset)); //relative offset of local file header, since it begins where the global header used to begin
#else
  global_header += static_cast<uint16_t>(0x0201); //second part of sig
  global_header += static_cast<uint16_t>(20); //version made by
  global_header.insert(global_header.end(), local_header.begin()+4, local_header.begin()+30);
  global_header += static_cast<uint16_t>(0); //file comment length
  global_header += static_cast<uint16_t>(0); //disk number where file starts
  global_header += static_cast<uint16_t>(0); //internal file attributes
  global_header += static_cast<uint32_t>(0); //external file attributes
  global_header += static_cast<uint32_t>(global_header_offset); //relative offset of local file header, since it begins where the global header used to begin
#endif
  global_header += fname;
 
  //build footer
  std::vector<char> footer;
  footer += "PK"; //first part of sig
#ifdef VISP_BIG_ENDIAN
  footer += vpEndian::swap16bits(static_cast<uint16_t>(0x0605)); //second part of sig
  footer += static_cast<uint16_t>(0); //number of this disk
  footer += static_cast<uint16_t>(0); //disk where footer starts
  footer += vpEndian::swap16bits(static_cast<uint16_t>(nrecs+1)); //number of records on this disk
  footer += vpEndian::swap16bits(static_cast<uint16_t>(nrecs+1)); //total number of records
  footer += vpEndian::swap32bits(static_cast<uint32_t>(global_header.size())); //nbytes of global headers
  footer += vpEndian::swap32bits(static_cast<uint32_t>(global_header_offset + nbytes + local_header.size())); //offset of start of global headers, since global header now starts after newly written array
#else
  footer += static_cast<uint16_t>(0x0605); //second part of sig
  footer += static_cast<uint16_t>(0); //number of this disk
  footer += static_cast<uint16_t>(0); //disk where footer starts
  footer += static_cast<uint16_t>(nrecs+1); //number of records on this disk
  footer += static_cast<uint16_t>(nrecs+1); //total number of records
  footer += static_cast<uint32_t>(global_header.size()); //nbytes of global headers
  footer += static_cast<uint32_t>(global_header_offset + nbytes + local_header.size()); //offset of start of global headers, since global header now starts after newly written array
#endif
  footer += static_cast<uint16_t>(0); //zip file comment length
 
  //write everything
  fwrite(&local_header[0], sizeof(char), local_header.size(), fp);
  fwrite(&npy_header[0], sizeof(char), npy_header.size(), fp);
  fwrite(&data_str_utf32_LE[0], sizeof(char), nels, fp);
  fwrite(&global_header[0], sizeof(char), global_header.size(), fp);
  fwrite(&footer[0], sizeof(char), footer.size(), fp);
  fclose(fp);
}

void visp::cnpy::npz_save(const std::string &zipname, const std::string &fname, const std::string &data_str, const std::string &mode)
{
  std::vector<std::string> data_vec;
  data_vec.push_back(data_str);
  std::vector<size_t> shape { 1 };
  npz_save(zipname, fname, data_vec, shape, mode);
}
 
#endif