// -*- C++ -*-
//
// This file is part of YODA -- Yet more Objects for Data Analysis
// Copyright (C) 2008-2023 The YODA collaboration (see AUTHORS for details)
//
#ifndef YODA_H5UTILS_H
#define YODA_H5UTILS_H

#include <YODA/Utils/BinningUtils.h>
#include <YODA/Config/BuildConfig.h>
#ifdef WITH_HIGHFIVE
#include <YODA/highfive/H5File.hpp>
#else
#include <highfive/H5File.hpp>
#define YODA_H5 HighFive
#endif

#include <map>
#include <memory>
#include <string>
#include <type_traits>
#include <vector>

using std::map;
using std::string;
using std::vector;

namespace YODA {

  /// @brief Helper method to construct and fill a YODA_H5::DataSet
  ///
  /// @note This version is suited best for 1D vectors
  /// that to be committed in one go
  template <typename T>
  YODA_H5::DataSet H5DataSet(YODA_H5::File& h5file, const string& label,
                              vector<T>&& data, bool compress) {

    YODA_H5::DataSetCreateProps props;
    if (compress && data.size()) {
      props.add(YODA_H5::Chunking(vector<hsize_t>{data.size()}));
      props.add(YODA_H5::Deflate(9));
    }
    return h5file.createDataSet(label, std::move(data), props);
  }


  /// @brief Helper method to construct an empty YODA_H5::DataSet
  ///
  /// @note This version is suited best for 2D vectors
  /// that to be committed row-by-row
  template <typename T>
  YODA_H5::DataSet H5DataSet(YODA_H5::File& h5file, const string& label,
                              size_t nrows, size_t chunksize, bool compress) {

    YODA_H5::DataSetCreateProps props;
    if (compress) { // enable compression
      props.add(YODA_H5::Chunking(vector<hsize_t>{1,chunksize}));
      props.add(YODA_H5::Deflate(9));
    }
    // create the dataset
    return h5file.createDataSet(label, YODA_H5::DataSpace({nrows,chunksize}),
                                YODA_H5::create_datatype<T>(), props);
  }

  /// @brief Helper class to extract information from YODA_H5::DataSets
  class H5DataSetReader {
    public:

    /// @brief Nullary constructor
    H5DataSetReader() { }

    // @brief Constructor from existing DataSet
    H5DataSetReader(const YODA_H5::File& h5file, const string& label)
          : _nrows(0), _thisrow(0), _ncols(0), _thiscol(0), _ds(nullptr) {
      _ds = std::make_unique<YODA_H5::DataSet>(h5file.getDataSet(label));
      const auto& dims = _ds->getDimensions();
      _nrows = dims.at(0);
      if (dims.size() > 1)  _ncols = dims.at(1);
    }

    /// @brief Move internal cursor by @a len elements
    void skip(size_t len) noexcept { _thisrow += len; }

    /// @brief Load next @a item and increment cursor
    template <typename T = size_t>
    T next() noexcept {
      T item;
      _ds->select({_thisrow++}, {1}).read(item);
      return item;
    }

    /// Method to read and return the entire 1D dataset
    template <typename T>
    vector<T> read() noexcept {
      if (_nrows == 0)  return vector<T>{};
      vector<T> data; data.reserve(_nrows);
      _ds->select({0}, {_nrows}).read(data);
      return data;
    }

    /// Method to read a subset of the 1D dataset
    template <typename T>
    vector<T> read(size_t len) noexcept {
      if (len == 0)  return {};
      vector<T> data; data.reserve(len);
      _ds->select({_thisrow}, {len}).read(data);
      _thisrow += len;
      return data;
    }

    /// Method to read a subset of the 1D dataset
    template <typename T>
    vector<T> readSlice(size_t len) noexcept {
      if (len == 0)  return {};
      vector<T> data; data.reserve(len);
      while (len) {
        vector<vector<T>> tmp;
        size_t ncols = std::min(len, _ncols - _thiscol);
        size_t nrows = 1+(ncols-1)/_ncols; // C-style ceil
        _ds->select({_thisrow,_thiscol}, {nrows,ncols}).read(tmp);
        for (size_t i=0; i < tmp.size(); ++i) {
          data.insert(data.end(), std::make_move_iterator(std::begin(tmp[i])),
                                  std::make_move_iterator(std::end(tmp[i])));
        }
        if ((_thiscol + ncols) == _ncols)  ++_thisrow;
        _thiscol = (_thiscol + ncols) % _ncols;
        len -= ncols;
      }
      return data;
    }

    /// Method to read a subset of the 1D dataset starting from row @a row
    template <typename T>
    vector<T> readAt(size_t row, size_t len) noexcept {
      _thisrow = row;
      return read<T>(len);
    }

    /// Method to read an attribute decorated onto this dataset
    template <typename T = size_t>
    vector<T> readAttribute(const string& label) const noexcept {
      vector<T> data;
      _ds->getAttribute(label).read(data);
      return data;
    }

    private:

    size_t _nrows, _thisrow, _ncols, _thiscol;

    std::unique_ptr<YODA_H5::DataSet> _ds;

  };


  /// @brief A helper class to deal with chunking of H5 datasets
  template <typename T>
  class H5DataSetWriter {
    public:

    // Constructor for empty DataSet
    H5DataSetWriter(YODA_H5::File& h5file, const string& label,
                    size_t datalen, size_t chunksize, bool compress)
          : _ncols(chunksize), _thiscol(0), _nrows(1+(datalen-1)/chunksize), _thisrow(0),
            _ds(H5DataSet<T>(h5file, label, _nrows, chunksize, compress)) {  }

    /// @brief Method to write a slice that may span multiple rows/columns
    void writeSlice(vector<T>&& data) noexcept {

      if (data.empty())  return;

      size_t offset = 0, len = data.size();
      const auto itr = data.cbegin();
      while (len) {
        vector<T> tmp;
        size_t ncols = std::min(len, _ncols - _thiscol);
        auto first = itr + offset;
        auto last = first + ncols;
        _ds.select({_thisrow,_thiscol},
                   {1,ncols}).write(vector<vector<T>>{{std::make_move_iterator(first),
                                                       std::make_move_iterator(last)}});
        offset += ncols;
        _thiscol += ncols;
        if (_thiscol == _ncols) {
          _thiscol = 0; ++_thisrow;
        }
        len -= ncols;
      }
    }

    /// @brief Method to decorate dataset with an attribute
    template <typename U>
    void createAttribute(const string& label, vector<U>&& data) noexcept {
      _ds.createAttribute(label, std::forward<vector<U>>(data));
    }

    size_t colPos() const noexcept { return _thiscol; }

    size_t rowPos() const noexcept { return _thisrow; }

    private:

    size_t _ncols, _thiscol, _nrows, _thisrow;

    YODA_H5::DataSet _ds;

  };




  /// @brief Base wrapper around a vector of edges
  class EdgeHandlerBase {
    public:

    /// Default constructor
    EdgeHandlerBase() { }

    /// Default destructor
    virtual ~EdgeHandlerBase() { }

    virtual void writeToFile(const string&, YODA_H5::File&, bool compress) = 0;

  };




  /// @brief Specialised wrapper for a vector of type T
  template<typename T>
  class EdgeHandler : public EdgeHandlerBase {
    public:

    using EdgeHandlerBase::EdgeHandlerBase;

    // Constructor from a vector of edges
    EdgeHandler(const vector<T>& edges)
        : _edges(edges) { }

    // Constructor from an rvalue vector of edges
    EdgeHandler(vector<T>&& edges)
        : _edges(std::move(edges)) { }

    // Method to extend the vector of edges
    void extend(vector<T>&& edges) {
      _edges.insert(_edges.end(), std::make_move_iterator(std::begin(edges)),
                                  std::make_move_iterator(std::end(edges)));
    }

    auto begin() const { return _edges.cbegin(); }

    // Method to commit vector of edges to H5 @a file using @a label
    void writeToFile(const string& label, YODA_H5::File& file, bool compress) {
      (void)H5DataSet(file, label, std::move(_edges), compress);
    }

    private:

    vector<T> _edges;
  };

  /// Convenience aliases
  template <typename T>
  using EdgeHandlerPtr = typename std::shared_ptr<EdgeHandler<T>>;

  using EdgeHandlerBasePtr = typename std::shared_ptr<EdgeHandlerBase>;




  /// @brief Helper class to extract AO information from a H5 file
  class H5FileManager {

    public:

    static constexpr size_t AO_META = 2;

    /// @brief Constructor
    H5FileManager(const YODA_H5::File& file)
        : _index(-1), _cachepos(-1),
          _labelindex(0), _h5file(file),
          _aoinfo(file, "aoinfo"),
          _layout(file, "sizeinfo"),
          _content(file, "content"),
          _annos(file, "annotations"),
          _labels(H5DataSetReader(file, "labels").read<string>()) {

      _meta = _aoinfo.readAttribute("meta");
      if (_meta.size() < 2)
        throw ReadError("No file metadata found!");

      _annosizes = _layout.read<size_t>(_meta.at(1));
      _datasizes = _layout.read<size_t>(_meta.at(1));
      _labelsizes = _layout.read<size_t>(_layout.next());
    }

    /// @brief H5 YODA format version
    size_t version() const {
      return _meta.at(0);
    }

    /// @brief Number of AOs in this H5 file
    size_t size() const {
      return _meta.at(1);
    }

    /// @brief Loads next AO from file
    bool next() {
      ++_index;
      if ((size_t)_index == size())  return false;

      _aodims = _aoinfo.readAt<string>(AO_META*_index, AO_META);
      if (_aodims.empty()) {
        throw ReadError("No AO information found!");
      }

      _cachepos = -1;

      return true;
    }

    /// @brief Path of current AO
    const string& path() const {
      return _aodims.at(0);
    }

    /// @brief Type of current AO
    const string& type() const {
      return _aodims.at(1);
    }

    /// @brief Serialized annotations of current AO
    vector<string> loadAnnotations() noexcept {
      return _annos.readSlice<string>(_annosizes[_index]);
    }

    /// @brief Serialized content of current AO
    vector<double> loadContent() noexcept {
      return _content.readSlice<double>(_datasizes[_index]);
    }

    /// @brief Skips next set of annotations and content of current AO
    void skipCommon() noexcept {
      _annos.skip(_annosizes[_index]);
      _content.skip(_datasizes[_index]);
    }

    /// @brief Indices of masked bins in current AO
    vector<size_t> loadMasks() noexcept {
      return _layout.read<size_t>(_layout.next());
    }

    /// @brief Skips next set of masked indices of current AO
    void skipMasks() noexcept {
      _layout.skip(_layout.next());
    }

    /// @brief Labels of error sources of current AO
    vector<string> loadSources() noexcept {

      if (_cachepos < 0) { // fill label cache
        _labelcache = _layout.read<size_t>(_labelsizes[_labelindex++]);
        _cachepos = 0;
      }

      size_t len = _labelcache[_cachepos++];
      if (len == 0) return {};

      vector<size_t> indices(_labelcache.begin()+_cachepos, _labelcache.begin()+_cachepos+len);
      vector<string> sources; sources.reserve(indices.size());
      for (size_t idx : indices) {
        sources.emplace_back(_labels[idx]);
      }
      _cachepos += len;
      return sources;
    }

    /// @brief Skips next set of error sources of current AO
    void skipSources() noexcept {
      _layout.skip(_labelsizes[_labelindex++]);
    }

    /// @brief Returns next set of edges of type EdgeT
    template <typename EdgeT>
    vector<EdgeT> loadEdges() noexcept {

      const string label = string("edges_") + TypeID<EdgeT>::name();
      auto itr = datasets.find(label);
      if (itr == datasets.end()) {
        // put H5::DataSet into DataSet cache
        datasets[label] = H5DataSetReader(_h5file, label);
        itr = datasets.find(label);
      }
      return itr->second.read<EdgeT>(_layout.next());
    }

    /// @brief Skips next set of edges of type EdgeT
    template <typename EdgeT>
    void skipEdges() noexcept {
      const string label = string("edges_") + TypeID<EdgeT>::name();
      auto itr = datasets.find(label);
      if (itr == datasets.end()) {
        // put H5::DataSet into DataSet cache
        datasets[label] = H5DataSetReader(_h5file, label);
        itr = datasets.find(label);
      }
      itr->second.skip(_layout.next());
    }

    private:

    ssize_t _index, _cachepos;

    size_t _labelindex;

    const YODA_H5::File _h5file;

    H5DataSetReader _aoinfo, _layout, _content, _annos;

    vector<string> _labels, _aodims;

    vector<size_t> _meta, _annosizes, _datasizes, _labelsizes, _labelcache;

    map<string,H5DataSetReader> datasets;
  };


}

#endif
