# cython: language_level=3
# cython: warn.undeclared=True
# cython: warn.unreachable=True
# cython: warn.maybe_uninitialized=True
# cython: warn.unused=True
# cython: warn.unused_arg=True
# cython: warn.unused_result=True
# cython: warn.multiple_declarators=True
# distutils: language = c++

import cython
from libc.stdint cimport int32_t, int64_t, uint64_t
from libcpp.string cimport string
import numpy as np


# C++ functions
cdef extern from "interface_hdf5.h" namespace "tbplas::python":
    void hdf5_open(const string& file_name)
    void hdf5_close()
    void hdf5_get_matrix_double_shape(
        const string& key,
        uint64_t& num_rows, 
        uint64_t& num_cols)
    void hdf5_get_matrix_complex_shape(
        const string& key,
        uint64_t& num_rows,
        uint64_t& num_cols)
    void hdf5_get_matrix_double(
        const string& key,
        double* data)
    void hdf5_get_matrix_complex(
        const string& key,
        double complex* data)


# Python functions
def get_matrix_double(const string& filePath, const string& key):
    cdef uint64_t num_rows, num_cols
    cdef double [::1, :] data

    hdf5_open(filePath)
    hdf5_get_matrix_double_shape(key, num_rows, num_cols)
    # Allocation should be done at Cython level to avoid memory leakage.
    data = np.zeros((num_rows, num_cols), dtype=np.float64, order="F")
    hdf5_get_matrix_double(key, &data[0][0])
    hdf5_close()
    return np.asarray(data)


def get_matrix_complex(const string& filePath, const string& key):
    cdef uint64_t num_rows, num_cols
    cdef double complex [::1, :] data

    hdf5_open(filePath)
    hdf5_get_matrix_complex_shape(key, num_rows, num_cols)
    # Allocation should be done at Cython level to avoid memory leakage.
    data = np.zeros((num_rows, num_cols), dtype=np.complex128, order="F")
    hdf5_get_matrix_complex(key, &data[0][0])
    hdf5_close()
    return np.asarray(data)


def test_matrix_shape():
    """
    Test function for determining shape of matrices with only 1 ROW or COLUMN.
    """
    cdef double [:,::1] data1, data2
    data1 = np.zeros((10, 1), dtype=np.float64)
    data2 = np.zeros((1,10), dtype=np.float64)
    return np.asarray(data1), np.asarray(data2)
