//
// This file is part of MorphoDynamX - http://www.MorphoDynamX.org
// Copyright (C) 2012-2021 Richard S. Smith and collaborators.
//
// If you use MorphoDynamX in your work, please cite:
//   http://dx.doi.org/10.7554/eLife.05864
//
// MorphoDynamX is free software, and is licensed under under the terms of the 
// GNU General (GPL) Public License version 2.0, http://www.gnu.org/licenses.
// 
// Derived from hdf_explorer, Copyright (C) 2016 Pedro Vicente
//
#include <HDF5Tree.hpp>
#include <iterate.hpp>

using namespace H5;

//
// hdf_dataset_t (common definition for HDF5 dataset and attribute)
// a HDF dataset is defined here simply has a:
// 1) location in memory to store dataset data
// 2) an array of dimensions
// 3) a full name path to open using a file id
// 4) size, sign and class of datatype
// the array of dimensions (of HDF defined type ' hsize_t') is defined in iteration
// the data buffer and datatype sizes are stored on per load variable 
// from tree using the HDF API from item input
//
class hdf_dataset_t
{
public:
  hdf_dataset_t(const char* path, const std::vector< hsize_t> &dim, size_t size, H5T_sign_t sign, H5T_class_t datatype_class) :
    m_path(path), m_dim(dim), m_datatype_size(size), m_datatype_sign(sign), m_datatype_class(datatype_class)
  {
    m_buf = NULL;
  }

  ~hdf_dataset_t() { free(m_buf); }

  void store(void *buf) { m_buf = buf; }
  std::string m_path;
  std::vector<hsize_t> m_dim;

  //needed to access HDF5 buffer data
  size_t m_datatype_size;
  H5T_sign_t m_datatype_sign;
  H5T_class_t  m_datatype_class;
  void *m_buf;
};

//get_item_data
ItemData* get_item_data(QTreeWidgetItem *item)
{
  QVariant data = item->data(0, Qt::UserRole);
  ItemData *item_data = data.value<ItemData*>();
  return item_data;
}

//struct to get name and type of an object
struct name_type_t
{
  char *name;
  int type;
};

//callback function function for H5Literate
static herr_t count_objects_cb(hid_t, const char *, const H5L_info_t *, void *_op_data)
{
  hsize_t *op_data = (hsize_t *)_op_data;

  (*op_data)++;

  return(H5_ITER_CONT);
}

//callback function function for H5Literate
static herr_t get_name_type_cb(hid_t loc_id, const char *name, const H5L_info_t *, void *op_data)
{
  H5G_stat_t stat_buf;

  if(H5Gget_objinfo(loc_id, name, 0, &stat_buf) < 0) {}

  ((name_type_t *)op_data)->type = stat_buf.type;
  ((name_type_t *)op_data)->name = (char *)strdup(name);

  // define H5_ITER_STOP for return. This will cause the iterator to stop 
  return H5_ITER_STOP;
}

namespace mgx
{
  HDF5Tree::HDF5Tree(const QString &fileName, QTreeWidget *tree)
  {
    h5File = new H5File(fileName.toStdString(), H5F_ACC_RDONLY);
    QTreeWidgetItem *item = new QTreeWidgetItem(QStringList() << "/");
    item->setFlags(Qt::ItemIsEnabled);
    item->setIcon(0, QIcon(":/images/HDF.png"));
    tree->addTopLevelItem(item);;
    
    size_t objCount = h5File->getObjCount();
    hid_t objList[objCount];
    h5File->getObjIDs(H5F_OBJ_ALL, objCount, objList);
    iterate(fileName.toStdString(), "/", objList[0], item);
    if(firstDataSet)
      tree->setCurrentItem(firstDataSet);
  }
  HDF5Tree::~HDF5Tree() 
  { 
    if(h5File) 
      delete h5File; 
  }

  H5O_info_added_t* HDF5Tree::find_object(haddr_t addr)
  {
    for(size_t idx = 0; idx < m_visit.visit_info.size(); idx++)
      if(addr == m_visit.visit_info[idx].oinfo.addr)
        return &(m_visit.visit_info[idx]);

    return NULL;
  }
  
  int HDF5Tree::get_attributes(const std::string& file_name, const std::string& path, const hid_t loc_id, QTreeWidgetItem *tree_item_parent)
  {
    H5O_info_t oinfo;
    hid_t aid;
    size_t datatype_size;
    H5T_sign_t datatype_sign;
    H5T_class_t datatype_class;
    hid_t sid;
    hid_t ftid;
    hid_t mtid;
    hsize_t dims[H5S_MAX_RANK];
    char name[1024];
    int rank;
    ItemData *item_data;
    QVariant data;
    QTreeWidgetItem *item = NULL;
  
    //get object info
    if(H5Oget_info(loc_id, &oinfo) < 0) {}
  
    for(hsize_t idx = 0; idx < oinfo.num_attrs; idx++) {
      if((aid = H5Aopen_by_idx(loc_id, ".", H5_INDEX_CRT_ORDER, H5_ITER_INC, idx, H5P_DEFAULT, H5P_DEFAULT)) < 0) {}
      if(H5Aget_name(aid, (size_t)1024, name) < 0) {}
      if((sid = H5Aget_space(aid)) < 0) {}
      if((rank = H5Sget_simple_extent_dims(sid, dims, NULL)) < 0) {}
      if((ftid = H5Aget_type(aid)) < 0) {}
      if((mtid = H5Tget_native_type(ftid, H5T_DIR_DEFAULT)) < 0) {}
  
      //store datatype sizes and metadata needed to display HDF5 buffer data
      if((datatype_size = H5Tget_size(mtid)) == 0) {}
      if((datatype_sign = H5Tget_sign(mtid)) < 0) {}
      if((datatype_class = H5Tget_class(mtid)) < 0) {}
      if(H5Sclose(sid) < 0) {}
      if(H5Tclose(ftid) < 0) {}
      if(H5Tclose(mtid) < 0) {}
      if(H5Aclose(aid) < 0) {}
  
      //store dimensions in ItemData
      std::vector< hsize_t> dim; //dataset dimensions
      for(int idx = 0; idx < rank; ++idx)
        dim.push_back(dims[idx]);
  
      //store a hdf_dataset_t with full path, dimensions and metadata
      hdf_dataset_t *dataset = new hdf_dataset_t(path.c_str(), dim, datatype_size, datatype_sign, datatype_class);
  
      //append item
      item = new QTreeWidgetItem(tree_item_parent);
      item->setText(0, name);
      item->setFlags(Qt::ItemIsEnabled);
      item->setIcon(0, QIcon(":/images/Attributes.png"));
      //item data
      std::string item_path = tree_item_parent->text(0).toStdString() + "/" + name;
      item_data = new ItemData(ItemData::Attribute, file_name, item_path, name, dataset);
      data.setValue(item_data);
      item->setData(0, Qt::UserRole, data);
    }
  
    return 0;
  }
  
  int HDF5Tree::iterate(const std::string& file_name, const std::string& grp_path, const hid_t loc_id, QTreeWidgetItem *tree_item_parent)
  {
    hsize_t nbr_objects = 0;
    QTreeWidgetItem *item_grp = NULL;
    QTreeWidgetItem *item_var = NULL;
    hsize_t index;
    name_type_t info;
    ItemData *item_data;
    QVariant data;
    std::string path;
    bool do_iterate;
    size_t datatype_size;
    H5T_sign_t datatype_sign;
    H5T_class_t datatype_class;
    hid_t gid;
    hid_t did;
    hid_t sid;
    hid_t ftid;
    hid_t mtid;
    hsize_t dims[H5S_MAX_RANK];
    int rank;

    firstDataSet = 0;
  
    if(H5Literate(loc_id, H5_INDEX_NAME, H5_ITER_INC, NULL, count_objects_cb, &nbr_objects) < 0)
      ;
  
    for(hsize_t idx_obj = 0; idx_obj < nbr_objects; idx_obj++) { 
      index = idx_obj;
  
      //'index' allows an interrupted iteration to be resumed; it is passed in by the application with a starting point 
      //and returned by the library with the point at which the iteration stopped
  
      if(H5Literate(loc_id, H5_INDEX_NAME, H5_ITER_INC, &index, get_name_type_cb, &info) < 0) {}
  
      // initialize path 
      path = grp_path;
      if(grp_path != "/")
        path += "/";
      path += info.name;
  
      switch(info.type)
      {
      case H5G_GROUP:
        if((gid = H5Gopen2(loc_id, info.name, H5P_DEFAULT)) < 0) {}
  
        H5O_info_t oinfo_buf;
        do_iterate = true;
  
        //get object info
        if(H5Oget_info(gid, &oinfo_buf) < 0) {}
  
        //group item
        item_grp = new QTreeWidgetItem(tree_item_parent);
        item_grp->setText(0, info.name);
        item_grp->setFlags(Qt::ItemIsEnabled);
        item_grp->setIcon(0, QIcon(":/images/Group.png"));
        //item data
        item_data = new ItemData(ItemData::Group, file_name, path, info.name, (hdf_dataset_t*)NULL);
        data.setValue(item_data);
        item_grp->setData(0, Qt::UserRole, data);
  
        if(oinfo_buf.rc > 1) {
          H5O_info_added_t *oinfo_added = find_object(oinfo_buf.addr);
          if(oinfo_added->added > 0) {
            //avoid infinite recursion due to a circular path in the file.
            do_iterate = false;
          } else {
            oinfo_added->added++;
          }
        }
  
        //iterate in sub group passing QTreeWidgetItem for group as parent
        if(do_iterate && iterate(file_name, path, gid, item_grp) < 0) {}
        if(get_attributes(file_name, path, gid, item_grp) < 0) {}
        if(H5Gclose(gid) < 0) {}
  
        free(info.name);
        break;
  
      case H5G_DATASET:
        if((did = H5Dopen2(loc_id, info.name, H5P_DEFAULT)) < 0) {}
        if((sid = H5Dget_space(did)) < 0) {}
        if((rank = H5Sget_simple_extent_dims(sid, dims, NULL)) < 0) {}
        if((ftid = H5Dget_type(did)) < 0) {}
        if((mtid = H5Tget_native_type(ftid, H5T_DIR_DEFAULT)) < 0) {}
  
        //store datatype sizes and metadata needed to display HDF5 buffer data
        if((datatype_size = H5Tget_size(mtid)) == 0) {}
        if((datatype_sign = H5Tget_sign(mtid)) < 0) {}
        if((datatype_class = H5Tget_class(mtid)) < 0) {}
        if(H5Sclose(sid) < 0) {}
        if(H5Tclose(ftid) < 0) {}
        if(H5Tclose(mtid) < 0) {}
  
        //store dimensions in ItemData
        std::vector< hsize_t> dim; //dataset dimensions
        for(int idx = 0; idx < rank; ++idx)
          dim.push_back(dims[idx]);
  
        //store a hdf_dataset_t with full path, dimensions and metadata
        hdf_dataset_t *dataset = new hdf_dataset_t(path.c_str(), dim, datatype_size, datatype_sign, datatype_class);
  
        //append item
        item_var = new QTreeWidgetItem(tree_item_parent);
        item_var->setFlags(Qt::ItemIsEnabled | Qt::ItemIsSelectable);
        item_var->setText(0, info.name);
        if(!firstDataSet)
          firstDataSet = item_var;
        item_var->setIcon(0, QIcon(":/images/Image.png"));
        //item data
        item_data = new ItemData(ItemData::Variable, file_name, path, info.name, dataset);
        data.setValue(item_data);
        item_var->setData(0, Qt::UserRole, data);
  
        if(get_attributes(file_name, path, did, item_var) < 0) {}
        if(H5Dclose(did) < 0) {}
  
        break;
      }
    }
    return 0;
  }
}
