//
// This file is part of 3DCellAtlasOvule.
// Copyright (C) 2015-2016 George W. Bassel and collaborators.
//
// 3DCellAtlasOvule is an AddOn for MorphoGraphX - http://www.MorphoGraphX.org
// Copyright (C) 2012-2012 Richard S. Smith and collaborators.
//
// MorphoGraphX is free software, and are licensed under under the terms of the
// GNU General (GPL) Public License version 2.0, http://www.gnu.org/licenses.
//
#include <HDF5Process.hpp>
#include <HDF5Tree.hpp>

#include <Dir.hpp>

#include "ui_HDF5Open.h"
#include "iterate.hpp"
#include "visit.hpp"

using namespace H5;

#ifdef WIN32
#  define FileDialogOptions QFileDialog::DontUseNativeDialog
#else
#  define FileDialogOptions 0
#endif

namespace mgx
{
  static const char *const ELEMENT_SIZE_UM = "element_size_um";

  bool StackSaveHDF5::run(Stack &stack, Store &store, QString fileName, QString dataSetName, int compressionLevel)
  {
    if(fileName.isEmpty())
      throw QString("%1 File name empty").arg(name());

    if(!(fileName.endsWith(".h5", Qt::CaseInsensitive) or fileName.endsWith(".hdf5", Qt::CaseInsensitive)
                           or fileName.endsWith(".hdf", Qt::CaseInsensitive))) 
      throw QString("%1 Invalid file type %2").arg(name()).arg(fileName);

    // Create the file
    H5File file(fileName.toStdString().c_str(), H5F_ACC_TRUNC);

    // Set the sizes for the dataset
    auto stackSize = stack.size();
    hsize_t dims[3] = {stackSize.z(), stackSize.y(), stackSize.x()};
    DataSpace dataspace(3, dims);

    DSetCreatPropList plist;
    hsize_t dimsC[3] = {20, 20, 20}; // Not sure what these are for
    plist.setChunk(3, dimsC);
    plist.setDeflate(compressionLevel);

    QStringList groups = dataSetName.split("/");
    QString group;
    for(int i = 0; i < groups.size() - 1; i++) {
      if(groups[i].isEmpty())
        continue;
      group += "/" + groups[i];
      file.createGroup(group.toStdString());
    }

    DataSet dataset = file.createDataSet(dataSetName.toStdString(), PredType::NATIVE_UINT16, dataspace, plist);

    // Create a dataset attribute to write the voxel sizes. 
    hsize_t rank = 3;
    DataSpace attr_dataspace(1, &rank);
    Attribute attribute = dataset.createAttribute("element_size_um", PredType::NATIVE_FLOAT, attr_dataspace);

    auto step = stack.step();
    float voxelSize[3] = {step.z(), step.y(), step.x()};
    attribute.write(PredType::NATIVE_FLOAT, voxelSize);

    // Write the data
    dataset.write(store.data().data(), PredType::NATIVE_UINT16, dataspace, dataspace);

    file.close();

    store.setFile(stripCurrentDir(fileName));
    actingFile(fileName);

    return true;
  }
  
  bool StackSaveHDF5::initialize(QWidget* parent)
  {
    if(!parent)
      return true;

    Stack *stk = stack(parm("Stack Number").toInt());
    Store *store = (stringToWorkStore(parm("Store")) ? stk->work() : stk->main());
    QString fileName = parm("File Name");
    QString dataSetName = parm("Data Set Name");
    int compressionLevel = parm("Compression Level").toInt();

    QString h5Filter("HDF5 (*.h5 *.hdf5 *.hdf)");

    // Use the active file if parameter empty
    if(fileName.isEmpty())
      fileName = store->file();

    // Launch a dialog to get the file name
    fileName = QFileDialog::getSaveFileName(parent, QString("Select HDF5 file"), fileName, h5Filter, 0, FileDialogOptions);
    if(fileName.isEmpty())
      return false;

    // Add extension if not there
    if(!fileName.contains("."))
      fileName += ".h5";

    setParm("File Name",fileName);

    return true;
  }
  REGISTER_PROCESS(StackSaveHDF5);

  bool StackOpenHDF5::initialize(QWidget* parent)
  {
    if(!parent)
      return true;
    Stack *stk = stack(parm("Stack Number").toInt());
    bool isWorkStore = stringToWorkStore(parm("Store"));
    Store* store = (isWorkStore ? stk->work() : stk->main());
    QString dataSetName = parm("Data Set Name");

    QString h5Filter("HDF5 Stack files (*.h5 *.hdf5 *.hdf)"), filter;
    QString fileName = parm("File Name");

    if(fileName.isEmpty()) {
      QString currentFile;
      if(isWorkStore)
        currentFile = stk->work()->file();
      else
        currentFile = stk->main()->file();
      if(!currentFile.isEmpty())
        fileName = currentFile;
    }

    // Get the file name
    if(fileName.isEmpty())
      fileName = QFileDialog::getOpenFileName(parent, QString("Select HDF5 file"), fileName, QString(h5Filter), 0, FileDialogOptions);

    if(fileName.isEmpty())
      return false;

    setParm("File Name",fileName);

    if(dataSetName.isEmpty()) {
      Ui_HDF5OpenDialog ui;
      QDialog dlg(parent);
      ui.setupUi(&dlg);

      HDF5Tree h5Tree(fileName, ui.dataSets);
      ui.dataSets->expandAll();

      if(dlg.exec() == QDialog::Accepted) {
        auto currentItem = ui.dataSets->currentItem();
        if(currentItem > 0) {
          ItemData *itemData = get_item_data(currentItem);
          if(itemData)
            setParm("Data Set Name", itemData->m_item_path.c_str());
          else
            throw QString("%1 Bad item: %2").arg(name()).arg(currentItem->text(0));
        } else 
          return false;
      } else
        return false;
    }

    return true;
  }

  bool StackOpenHDF5::run(Stack* stack, Store* store, QString fileName, QString dataSetName)
  {
    if(fileName.isEmpty())
      throw QString("%1 File Name empty").arg(name());

    actingFile(fileName);
    store->setFile(fileName);
    QFile file(fileName);
    if(!file.open(QIODevice::ReadOnly))
      throw QString("%1 Cannot open input file: %2").arg(name()).arg(fileName);

    if(dataSetName.isEmpty())
      throw QString("%1 H5 dataset name cannot be empty").arg(name());

    // open H5 file
    H5File fileH5(fileName.toStdString(), H5F_ACC_RDONLY);
    // open dataset
    DataSet dataset = fileH5.openDataSet(dataSetName.toStdString());
    // get data type
    H5T_class_t typeClass = dataset.getTypeClass();

    if(typeClass == H5T_INTEGER)
      IntType intype = dataset.getIntType();
    else if(typeClass == H5T_FLOAT)
      FloatType fltype = dataset.getFloatType();
    else
      throw QString("%1 Unsupported dataset type:").arg(name()).arg(typeClass);

    DataSpace dataspace = dataset.getSpace();
    const int rank = dataspace.getSimpleExtentNdims();
    hsize_t dims[rank];
    dataspace.getSimpleExtentDims(dims, NULL);

    // get output vector size
    hsize_t dims_1d = 1;
    for(int i=0; i < rank; i++)
      dims_1d *= dims[i];

    std::vector<short> buf(dims_1d);

    // read depending on type
    if(typeClass == H5T_INTEGER)
      dataset.read(buf.data(), PredType::NATIVE_UINT16, dataspace, dataspace);
    else if(typeClass == H5T_FLOAT){
      std::vector<float> bufF(dims_1d);
      dataset.read(bufF.data(), PredType::NATIVE_FLOAT, dataspace, dataspace);

      // convert float (0...1) to short (0...65535)
      for(uint i = 0; i < bufF.size(); i++)
        buf[i] = (ushort)(bufF[i]*USHRT_MAX);
    }

    // read element_size_um if present
    float element_size_um[3];
    std::fill_n(element_size_um, rank, 1.);
    if (dataset.attrExists(ELEMENT_SIZE_UM)) {
      const Attribute attribute = dataset.openAttribute(ELEMENT_SIZE_UM);
      const DataType dataType = attribute.getDataType();
      attribute.read(dataType, element_size_um);
    } else
      Information::out << "No Element Size attribute found! Set voxel size to 1um x 1um x 1um" << endl;

    Point3u size(dims[2], dims[1], dims[0]);
    if(rank == 4) 
      size = Point3u(dims[3], dims[2], dims[1]);

    // set voxel size and step
    Point3f step(element_size_um[2], element_size_um[1], element_size_um[0]);
    stack->setSize(size);
    stack->setStep(step);

    // Copy the data
    memcpy(store->data().data(), buf.data(), store->data().size()*2); 

    fileH5.close();

    if(store->labels()) {
      ushort maxLabel = 0;
      const HVecUS& data = store->data();
      for(ushort u : data) {
        ushort label = u;
        if(label > maxLabel)
          maxLabel = label;
      }
      stack->setLabel(maxLabel);
    }

    store->setFile(fileName);
    store->changed();

    return true;
  }
  REGISTER_PROCESS(StackOpenHDF5);
}
