#ifdef WITH_HDF5
#include "io_hdf5.h"

#include <cassert>
#include <iostream>

namespace tbplas::base {

HDFFile::HDFFile(const std::string& filePath, OptBinMode mode)
{
    bool status = open(filePath, mode);
    if (!status) {
        std::cerr << "ERROR: failed to open HDFFile: " << m_CurFilePath
                  << " in mode: " << static_cast<int>(mode) << "\n";
        std::exit(-1);
    }
}

HDFFile::~HDFFile() { close(); }

bool HDFFile::open(const std::string& filePath, OptBinMode mode)
{
    // 重复打开文件往往意味着bug，可能导致上一次缓存的数据未写入硬盘，因此应当避免
    assert(isClosed());
    std::string temp_file = filePath;
    temp_file.append(".h5");
    m_H5File.close();
    bool is_open = true;
    H5::Exception::dontPrint();

    if (mode == OptBinMode::RESET) {
        // RESET模式下不读取文件内容，直接新建文件
        m_H5File = H5::H5File(temp_file.c_str(), H5F_ACC_TRUNC);
        is_open = true;
    } else if (mode == OptBinMode::UPDATE) {
        // UPDATE模式下若文件打开失败则创建文件
        try {
            m_H5File.openFile(temp_file.c_str(), H5F_ACC_RDWR);
        } catch (H5::FileIException& error) {
            m_H5File = H5::H5File(temp_file.c_str(), H5F_ACC_TRUNC);
            is_open = true;
        } catch (H5::DataSetIException& error) {
            H5::DataSetIException::printErrorStack();
            is_open = false;
        } catch (H5::DataSpaceIException& error) {
            H5::DataSpaceIException::printErrorStack();
            is_open = false;
        }
    } else {
        // READONLY模式下打开失败则返回false
        try {
            m_H5File.openFile(temp_file.c_str(), H5F_ACC_RDONLY);
        } catch (H5::FileIException& error) {
            is_open = false;
        } catch (H5::DataSetIException& error) {
            H5::DataSetIException::printErrorStack();
            is_open = false;
        } catch (H5::DataSpaceIException& error) {
            H5::DataSpaceIException::printErrorStack();
            is_open = false;
        }
    }
    m_FileOptMode = mode;
    BaseFile::open(temp_file, is_open);
    return is_open;
}

void base::HDFFile::close()
{
    // 重复关闭文件通常没有问题，若用户手动关闭了文件，析构时还会关闭一次
    // 打开HDF5文件时已经指定了模式，所以无需关心
    if (m_IsOpen) {
        m_H5File.close();
        BaseFile::close();
    }
}

/**
 * @brief 将矩阵写入HDF5数据集中
 *
 * @param[in] key 组名+数据集名，格式必须为XX.XX.XX
 * @param[in] value 待写入的矩阵
 * @return 写入操作执行状态
 */
OptStatus HDFFile::setMatrixXi(
    const std::string& key,
    const Eigen::MatrixXi& value)
{
    assert(isOpen());
    assert(isWritable());
    assert(isValidKey(key));
    assert(value.rows() > 0);
    assert(value.cols() > 0);
    std::vector<std::string> key_parse = parseKey(key);
    std::string group_name = key_parse[0];
    std::string dataset_name = key_parse[1];
    H5::Group grp = buildH5Group(group_name);
    if (!H5::H5File::isValid(grp.getId())) {
        return OptStatus::ERROR;
    }
    H5::DataSet dbset;
    hsize_t dims[2];
    // HDF5使用行优先，但Eigen默认使用列优先，因此这里要手动转置，不然h5dump输出结果会乱
    dims[1] = value.rows();
    dims[0] = value.cols();
    H5::DataSpace dataspace(2, dims);
    H5::IntType data_type(H5::PredType::NATIVE_INT);

    // 直接覆盖原有dataset而不做检查
    if (grp.exists(dataset_name.c_str())) {
        dbset = grp.openDataSet(dataset_name.c_str());
    } else {
        dbset = grp.createDataSet(dataset_name.c_str(), data_type, dataspace);
    }
    dbset.write(value.data(), data_type);
    dbset.close();
    grp.close();
    return OptStatus::SUCCESS;
}

/**
 * @brief 将矩阵写入HDF5数据集中
 *
 * @param[in] key 组名+数据集名，格式必须为XX.XX.XX
 * @param[in] value 待写入的矩阵
 * @return 写入操作执行状态
 */
OptStatus HDFFile::setMatrixXd(
    const std::string& key,
    const Eigen::MatrixXd& value)
{
    assert(isOpen());
    assert(isWritable());
    assert(isValidKey(key));
    assert(value.rows() > 0);
    assert(value.cols() > 0);
    std::vector<std::string> key_parse = parseKey(key);
    std::string group_name = key_parse[0];
    std::string dataset_name = key_parse[1];
    H5::Group grp = buildH5Group(group_name);
    if (!H5::H5File::isValid(grp.getId())) {
        return OptStatus::ERROR;
    }
    H5::DataSet dbset;
    hsize_t dims[2];
    // HDF5使用行优先，但Eigen默认使用列优先，因此这里要手动转置，不然h5dump输出结果会乱
    dims[1] = value.rows();
    dims[0] = value.cols();
    H5::DataSpace dataspace(2, dims);
    H5::FloatType data_type(H5::PredType::NATIVE_DOUBLE);

    // 直接覆盖原有dataset而不做检查
    if (grp.exists(dataset_name.c_str())) {
        dbset = grp.openDataSet(dataset_name.c_str());
    } else {
        dbset = grp.createDataSet(dataset_name.c_str(), data_type, dataspace);
    }
    dbset.write(value.data(), data_type);
    dbset.close();
    grp.close();
    return OptStatus::SUCCESS;
}

/**
 * @brief 将矩阵写入HDF5数据集中
 *
 * @param[in] key 组名+数据集名，格式必须为XX.XX.XX
 * @param[in] value 待写入的矩阵
 * @return 写入操作执行状态
 */
OptStatus HDFFile::setMatrixXcd(
    const std::string& key,
    const Eigen::MatrixXcd& value)
{
    assert(isOpen());
    assert(isWritable());
    assert(isValidKey(key));
    assert(value.rows() > 0);
    assert(value.cols() > 0);
    std::vector<std::string> key_parse = parseKey(key);
    std::string group_name = key_parse[0];
    std::string dataset_name = key_parse[1];
    H5::Group grp = buildH5Group(group_name);
    if (!H5::H5File::isValid(grp.getId())) {
        return OptStatus::ERROR;
    }
    H5::DataSet dbset;
    hsize_t dims[2];
    // HDF5使用行优先，但Eigen默认使用列优先，因此这里要手动转置，不然h5dump输出结果会乱。
    // 复数矩阵当作实数矩阵处理，内存空间要翻倍。dim[1]*2或者dim[0]*2都可以，但只能选1个。
    // 如果是dim[1]*2，h5dump输出结果为矩阵的转置。dim[0]*2则线性输出矩阵元素。
    // getMatrixXcd中value.resize()需与此处一致。
    dims[1] = value.rows() * 2;
    dims[0] = value.cols();
    H5::DataSpace dataspace(2, dims);
    H5::FloatType data_type(H5::PredType::NATIVE_DOUBLE);

    // 直接覆盖原有dataset而不做检查
    if (grp.exists(dataset_name.c_str())) {
        dbset = grp.openDataSet(dataset_name.c_str());
    } else {
        dbset = grp.createDataSet(dataset_name.c_str(), data_type, dataspace);
    }
    dbset.write(value.data(), data_type);
    dbset.close();
    grp.close();
    return OptStatus::SUCCESS;
}

/**
 * @brief 从HDF5数据集中读取矩阵
 *
 * @param[in] key 组名+数据集名，格式必须为XX.XX.XX
 * @param[out] value 待读取的矩阵
 * @return 读取操作是否成功
 */
OptStatus HDFFile::getMatrixXi(
    const std::string& key,
    Eigen::MatrixXi& value) const
{
    assert(isOpen());
    assert(isValidKey(key));
    std::vector<std::string> key_parse = parseKey(key);
    std::string group_name = key_parse[0];
    std::string dataset_name = key_parse[1];
    H5::Group grp;
    if (!findH5Group(group_name, grp)) {
        return OptStatus::GROUP_NOT_FOUND;
    }
    H5::IntType data_type(H5::PredType::NATIVE_INT);
    if (grp.exists(dataset_name.c_str())) {
        H5::DataSet dbset = grp.openDataSet(dataset_name.c_str());
        H5T_class_t datatype = dbset.getTypeClass();
        assert(datatype == H5T_INTEGER);
        H5::DataSpace space = dbset.getSpace();
        int rank = space.getSimpleExtentNdims();
        assert(rank == 2);
        hsize_t dims_out[2];
        space.getSimpleExtentDims(dims_out, nullptr);
        // HDF5使用行优先，但Eigen默认使用列优先，由于在写入时做了转置，因此这里需要转回来，不然结果是错的
        value.resize(dims_out[1], dims_out[0]);
        dbset.read(value.data(), data_type);
        dbset.close();
        grp.close();
        return OptStatus::SUCCESS;
    }
    return OptStatus::DATASET_NOT_FOUND;
}

/**
 * @brief 从HDF5数据集中读取矩阵
 *
 * @param[in] key 组名+数据集名，格式必须为XX.XX.XX
 * @param[out] value 待读取的矩阵
 * @return 读取操作是否成功
 */
OptStatus HDFFile::getMatrixXd(
    const std::string& key,
    Eigen::MatrixXd& value) const
{
    assert(isOpen());
    assert(isValidKey(key));
    std::vector<std::string> key_parse = parseKey(key);
    std::string group_name = key_parse[0];
    std::string dataset_name = key_parse[1];
    H5::Group grp;
    if (!findH5Group(group_name, grp)) {
        return OptStatus::GROUP_NOT_FOUND;
    }
    H5::FloatType data_type(H5::PredType::NATIVE_DOUBLE);
    if (grp.exists(dataset_name.c_str())) {
        H5::DataSet dbset = grp.openDataSet(dataset_name.c_str());
        H5T_class_t datatype = dbset.getTypeClass();
        assert(datatype == H5T_FLOAT);
        H5::DataSpace space = dbset.getSpace();
        int rank = space.getSimpleExtentNdims();
        assert(rank == 2);
        hsize_t dims_out[2];
        space.getSimpleExtentDims(dims_out, nullptr);
        // HDF5使用行优先，但Eigen默认使用列优先，由于在写入时做了转置，因此这里需要转回来，不然结果是错的
        value.resize(dims_out[1], dims_out[0]);
        dbset.read(value.data(), data_type);
        dbset.close();
        grp.close();
        return OptStatus::SUCCESS;
    }
    return OptStatus::DATASET_NOT_FOUND;
}

/**
 * @brief 从HDF5数据集中读取矩阵
 *
 * @param[in] key 组名+数据集名，格式必须为XX.XX.XX
 * @param[out] value 待读取的矩阵
 * @return 读取操作是否成功
 */
OptStatus HDFFile::getMatrixXcd(
    const std::string& key,
    Eigen::MatrixXcd& value) const
{
    assert(isOpen());
    assert(isValidKey(key));
    std::vector<std::string> key_parse = parseKey(key);
    std::string group_name = key_parse[0];
    std::string dataset_name = key_parse[1];
    H5::Group grp;
    if (!findH5Group(group_name, grp)) {
        return OptStatus::GROUP_NOT_FOUND;
    }
    H5::FloatType data_type(H5::PredType::NATIVE_DOUBLE);
    if (grp.exists(dataset_name.c_str())) {
        H5::DataSet dbset = grp.openDataSet(dataset_name.c_str());
        H5T_class_t datatype = dbset.getTypeClass();
        assert(datatype == H5T_FLOAT);
        H5::DataSpace space = dbset.getSpace();
        int rank = space.getSimpleExtentNdims();
        assert(rank == 2);
        hsize_t dims_out[2];
        space.getSimpleExtentDims(dims_out, nullptr);
        // HDF5使用行优先，但Eigen默认使用列优先，由于在写入时做了转置，因此这里需要转回来，不然结果是错的。
        // 由于复数矩阵当作实数矩阵存储，尺寸被乘以2，所以要除回来。具体dims_out[1]还是dims_out[1]需与
        // setMatrixXcd中一致。
        value.resize(dims_out[1] / 2, dims_out[0]);
        dbset.read(value.data(), data_type);
        dbset.close();
        grp.close();
        return OptStatus::SUCCESS;
    }
    return OptStatus::DATASET_NOT_FOUND;
}

OptStatus HDFFile::getMatrixXdShape(
    const std::string& key,
    uint64_t& num_rows,
    uint64_t& num_cols) const
{
    assert(isOpen());
    assert(isValidKey(key));
    std::vector<std::string> key_parse = parseKey(key);
    std::string group_name = key_parse[0];
    std::string dataset_name = key_parse[1];
    H5::Group grp;
    if (!findH5Group(group_name, grp)) {
        return OptStatus::GROUP_NOT_FOUND;
    }
    H5::FloatType data_type(H5::PredType::NATIVE_DOUBLE);
    if (grp.exists(dataset_name.c_str())) {
        H5::DataSet dbset = grp.openDataSet(dataset_name.c_str());
        H5T_class_t datatype = dbset.getTypeClass();
        assert(datatype == H5T_FLOAT);
        H5::DataSpace space = dbset.getSpace();
        int rank = space.getSimpleExtentNdims();
        assert(rank == 2);
        hsize_t dims_out[2];
        space.getSimpleExtentDims(dims_out, nullptr);
        // HDF5使用行优先，但Eigen默认使用列优先，由于在写入时做了转置，因此这里需要转回来，不然结果是错的
        num_rows = dims_out[1];
        num_cols = dims_out[0];
        dbset.close();
        grp.close();
        return OptStatus::SUCCESS;
    }
    return OptStatus::DATASET_NOT_FOUND;
}

OptStatus HDFFile::getMatrixXcdShape(
    const std::string& key,
    uint64_t& num_rows,
    uint64_t& num_cols) const
{
    assert(isOpen());
    assert(isValidKey(key));
    std::vector<std::string> key_parse = parseKey(key);
    std::string group_name = key_parse[0];
    std::string dataset_name = key_parse[1];
    H5::Group grp;
    if (!findH5Group(group_name, grp)) {
        return OptStatus::GROUP_NOT_FOUND;
    }
    H5::FloatType data_type(H5::PredType::NATIVE_DOUBLE);
    if (grp.exists(dataset_name.c_str())) {
        H5::DataSet dbset = grp.openDataSet(dataset_name.c_str());
        H5T_class_t datatype = dbset.getTypeClass();
        assert(datatype == H5T_FLOAT);
        H5::DataSpace space = dbset.getSpace();
        int rank = space.getSimpleExtentNdims();
        assert(rank == 2);
        hsize_t dims_out[2];
        space.getSimpleExtentDims(dims_out, nullptr);
        // HDF5使用行优先，但Eigen默认使用列优先，由于在写入时做了转置，因此这里需要转回来，不然结果是错的。
        // 由于复数矩阵当作实数矩阵存储，尺寸被乘以2，所以要除回来。具体dims_out[1]还是dims_out[1]需与
        // setMatrixXcd中一致。
        num_rows = dims_out[1] / 2;
        num_cols = dims_out[0];
        dbset.close();
        grp.close();
        return OptStatus::SUCCESS;
    }
    return OptStatus::DATASET_NOT_FOUND;
}

OptStatus HDFFile::getMatrixXd(
    const std::string& key,
    double* data) const
{
    assert(isOpen());
    assert(isValidKey(key));
    std::vector<std::string> key_parse = parseKey(key);
    std::string group_name = key_parse[0];
    std::string dataset_name = key_parse[1];
    H5::Group grp;
    if (!findH5Group(group_name, grp)) {
        return OptStatus::GROUP_NOT_FOUND;
    }
    H5::FloatType data_type(H5::PredType::NATIVE_DOUBLE);
    if (grp.exists(dataset_name.c_str())) {
        H5::DataSet dbset = grp.openDataSet(dataset_name.c_str());
        H5T_class_t datatype = dbset.getTypeClass();
        assert(datatype == H5T_FLOAT);
        H5::DataSpace space = dbset.getSpace();
        int rank = space.getSimpleExtentNdims();
        assert(rank == 2);
        hsize_t dims_out[2];
        space.getSimpleExtentDims(dims_out, nullptr);
        // HDF5使用行优先，但Eigen默认使用列优先，由于在写入时做了转置，因此这里需要转回来，不然结果是错的
        // value.resize(dims_out[1], dims_out[0]);
        dbset.read(data, data_type);
        dbset.close();
        grp.close();
        return OptStatus::SUCCESS;
    }
    return OptStatus::DATASET_NOT_FOUND;
}

OptStatus HDFFile::getMatrixXcd(
    const std::string& key,
    std::complex<double>* data) const
{
    assert(isOpen());
    assert(isValidKey(key));
    std::vector<std::string> key_parse = parseKey(key);
    std::string group_name = key_parse[0];
    std::string dataset_name = key_parse[1];
    H5::Group grp;
    if (!findH5Group(group_name, grp)) {
        return OptStatus::GROUP_NOT_FOUND;
    }
    H5::FloatType data_type(H5::PredType::NATIVE_DOUBLE);
    if (grp.exists(dataset_name.c_str())) {
        H5::DataSet dbset = grp.openDataSet(dataset_name.c_str());
        H5T_class_t datatype = dbset.getTypeClass();
        assert(datatype == H5T_FLOAT);
        H5::DataSpace space = dbset.getSpace();
        int rank = space.getSimpleExtentNdims();
        assert(rank == 2);
        hsize_t dims_out[2];
        space.getSimpleExtentDims(dims_out, nullptr);
        // HDF5使用行优先，但Eigen默认使用列优先，由于在写入时做了转置，因此这里需要转回来，不然结果是错的。
        // 由于复数矩阵当作实数矩阵存储，尺寸被乘以2，所以要除回来。具体dims_out[1]还是dims_out[1]需与
        // setMatrixXcd中一致。
        // value.resize(dims_out[1] / 2, dims_out[0]);
        dbset.read(data, data_type);
        dbset.close();
        grp.close();
        return OptStatus::SUCCESS;
    }
    return OptStatus::DATASET_NOT_FOUND;
}

H5::Group HDFFile::buildH5Group(const std::string& groupName)
{
    // Create root group if not exists.
    H5::Group grp = m_H5File.openGroup("/");
    if (!H5::H5File::isValid(grp.getId())) {
        m_H5File.createGroup("/");
    }

    // Build groups iteratively.
    std::vector<std::string> str_vec = splitKey(groupName);
    H5::Group grp_temp;
    for (const auto& item : str_vec) {
        if (!grp.exists(item.c_str())) {
            grp_temp = grp.createGroup(item.c_str());
        } else {
            grp_temp = grp.openGroup(item.c_str());
        }
        grp.close();
        grp = grp_temp;
    }
    return grp;
}

bool HDFFile::findH5Group(const std::string& groupName, H5::Group& target) const
{
    // Exit if root group does not exist.
    H5::Group grp = m_H5File.openGroup("/");
    if (!H5::H5File::isValid(grp.getId())) {
        return false;
    }

    // Search for the group iteratively.
    std::vector<std::string> str_vec = splitKey(groupName);
    H5::Group grp_temp;
    bool hit = true;
    for (const auto& item : str_vec) {
        if (!grp.exists(item.c_str())) {
            hit = false;
            break;
        }
        grp_temp = grp.openGroup(item.c_str());
        grp.close();
        grp = grp_temp;
    }
    target = grp;
    return hit;
}

bool HDFFile::isWritable() const
{
    bool status = (m_FileOptMode != OptBinMode::READONLY);
    assert(status);
    return status;
}

void HDFFile::save_matrix_integer(const std::string& key, const Eigen::MatrixXi& value)
{
    OptStatus status = setMatrixXi(key, value);
    check_opt_status(__func__, key, status);
}

void HDFFile::save_matrix_real(const std::string& key, const Eigen::MatrixXd& value)
{
    OptStatus status = setMatrixXd(key, value);
    check_opt_status(__func__, key, status);
}

void HDFFile::save_matrix_complex(const std::string& key, const Eigen::MatrixXcd& value)
{
    OptStatus status = setMatrixXcd(key, value);
    check_opt_status(__func__, key, status);
}

void HDFFile::read_matrix_integer(const std::string& key, Eigen::MatrixXi& value) const
{
    OptStatus status = getMatrixXi(key, value);
    check_opt_status(__func__, key, status);
}

void HDFFile::read_matrix_real(const std::string& key, Eigen::MatrixXd& value) const
{
    OptStatus status = getMatrixXd(key, value);
    check_opt_status(__func__, key, status);
}

void HDFFile::read_matrix_complex(const std::string& key, Eigen::MatrixXcd& value) const
{
    OptStatus status = getMatrixXcd(key, value);
    check_opt_status(__func__, key, status);
}

void HDFFile::read_matrix_real(const std::string& key, double* data) const
{
    OptStatus status = getMatrixXd(key, data);
    check_opt_status(__func__, key, status);
}

void HDFFile::read_matrix_complex(const std::string& key, std::complex<double>* data) const
{
    OptStatus status = getMatrixXcd(key, data);
    check_opt_status(__func__, key, status);
}

void HDFFile::read_matrix_real_shape(const std::string& key, uint64_t& num_rows, uint64_t& num_cols) const
{
    OptStatus status = getMatrixXdShape(key, num_rows, num_cols);
    check_opt_status(__func__, key, status);
}

void HDFFile::read_matrix_complex_shape(const std::string& key, uint64_t& num_rows, uint64_t& num_cols) const
{
    OptStatus status = getMatrixXcdShape(key, num_rows, num_cols);
    check_opt_status(__func__, key, status);
}

void check_opt_status(
    const std::string& task,
    const std::string& key,
    OptStatus status)
{
    std::string msg;
    switch (status) {
    case OptStatus::SUCCESS:
        msg = "ok";
        break;
    case OptStatus::KEY_NOT_FOUND:
        msg = "KEY_NOT_FOUND";
        break;
    case OptStatus::FORMAT_NOT_MATCH:
        msg = "FORMAT_NOT_MATCH";
        break;
    case OptStatus::DATASET_NOT_FOUND:
        msg = "DATASET_NOT_FOUND";
        break;
    case OptStatus::GROUP_NOT_FOUND:
        msg = "GROUP_NOT_FOUND";
        break;
    case OptStatus::ERROR:
        msg = "ERROR";
        break;
    }
    if (status != OptStatus::SUCCESS) {
        std::cerr << "ERROR: operation " << task
                  << " for key " << key
                  << " failed with status " << msg
                  << std::endl;
        std::abort();
    }
}

} // namespace tbplas::base
#endif // WITH_HDF5