"""
Common io functions for parsing data.

NOTES:
    1. np.loadtxt() yields a ROW vector if the data has only one ROW or COLUMN.
    2. Both a[:, i] or a[i, :] yield a ROW vector.
    3. a[0:1, :] yields a matrix with only one ROW vector.
    4. a[:, 0:1] yields a matrix with only one COLUMN vector.
    5. v[np.newaxis, :] is similar to a[0:1,:].
    6. v[:, np.newaxis] is similar to a[:,0:1].
"""

import glob
import os
import shutil

import numpy as np

from tbplas.cython import interface_hdf5 as interface


__all__ = ["h5_repack_file", "h5_repack_directory", "flatten_matrix",
           "np_load_double", "np_load_complex",
           "h5_load_double", "h5_load_complex"]


# Check if hdf5 is installed during import
H5_INSTALLED = (shutil.which("h5repack") is not None)


def h5_repack_file(file_name: str) -> None:
    """
    Repack h5 file after shrinking dataset to free disk usage.
    :param file_name: name of h5 file
    :return: None
    :raises RuntimeError: if h5repack command is not found
    :raises FileNotFoundError: if h5 file is not found
    """
    if not H5_INSTALLED:
        raise RuntimeError("h5repack command not found")
    if os.path.exists(file_name):
        tmp_name = f"{file_name}.tmp"
        os.system(f"h5repack {file_name} {tmp_name}")
        os.rename(tmp_name, file_name)
    else:
        raise FileNotFoundError(f"{file_name} not found")


def h5_repack_directory(dir_name: str = os.getcwd()) -> None:
    """
    Repack all h5 files under given directory.
    :param dir_name: directory name
    :return: None
    """
    h5_files = glob.glob(f"{dir_name}/*.h5")
    for f in h5_files:
        h5_repack_file(f)


def flatten_matrix(data: np.ndarray) -> np.ndarray:
    """
    Convert matrix with only one ROW or COLUMN to a ROW vector.
    :param data: (1, num_col) or (num_row, 1) array
        matrix to convert
    :return: (num_col, ) or (num_row, ) vector
    """
    shape = data.shape
    if len(shape) == 2 and (shape[0] == 1 or shape[1] == 1):
        data = data.flatten()
    return data


def np_load_double(file_name: str,
                   transpose: bool = False,
                   flatten: bool = False) -> np.ndarray:
    """
    Load a double matrix from text file.

    WARNING: DO NOT change the default value of arguments. Otherwise, A LOT OF
    functions will not work properly.s

    :param file_name: name of data file
    :param transpose: whether to transpose the matrix
    :param flatten: whether to flatten matrix with only 1 ROW or COLUMN to a
        ROW vector
    :return: (num_rows, num_cols) float64 matrix or ROW vector, depending on
        raw data shape and the 'flatten' parameter
    """
    # Enforce data to be a matrix even if it has only 1 row or column.
    data = np.loadtxt(file_name, ndmin=2)
    if transpose:
        data = data.T
    if flatten:
        data = flatten_matrix(data)
    return data


def np_load_complex(file_name: str,
                    transpose: bool = False,
                    flatten: bool = False) -> np.ndarray:
    """
    Load a double complex matrix from text file.

    This functions assuming the real and complex parts are in contiguous
    columns, i.e., real(a(:,0)), imag(a(:,0)), real(a(:,1)), imag(a(:,1)), ...

    WARNING: DO NOT change the default value of arguments. Otherwise, A LOT OF
    functions will not work properly.

    :param file_name: name of data file
    :param transpose: whether to transpose the matrix
    :param flatten: whether to flatten matrix with only 1 ROW or COLUMN to a
        ROW vector
    :return: (num_rows, num_cols) complex128 matrix or ROW vector, depending on
        raw data shape and the 'flatten' parameter
    """
    # Enforce data to be a matrix even if it has only 1 row or column.
    raw_data = np.loadtxt(file_name, ndmin=2)
    data = raw_data[:, 0::2] + 1j * raw_data[:, 1::2]
    if transpose:
        data = data.T
    if flatten:
        data = flatten_matrix(data)
    return data


def h5_load_double(file_name: str,
                   key: str,
                   transpose: bool = False,
                   flatten: bool = False) -> np.ndarray:
    """
    Load a double matrix from hdf5 file.

    WARNING: DO NOT change the default value of arguments. Otherwise, A LOT OF
    functions will not work properly.

    :param file_name: name of hdf5 file without the h5 suffix
    :param key: name of matrix
    :param transpose: whether to transpose the matrix
    :param flatten: whether to flatten matrix with only 1 ROW or COLUMN to a
        ROW vector
    :return: (num_rows, num_cols) float64 matrix or ROW vector, depending on
        raw data shape and the 'flatten' parameter
    """
    # Data is always a matrix even if it has only 1 row or column.
    data = interface.get_matrix_double(file_name.encode(), key.encode())
    if transpose:
        data = data.T
    if flatten:
        data = flatten_matrix(data)
    return data


def h5_load_complex(file_name: str,
                    key: str,
                    transpose: bool = False,
                    flatten: bool = False) -> np.ndarray:
    """
    Load a double complex matrix from hdf5 file.

    WARNING: DO NOT change the default value of arguments. Otherwise, A LOT OF
    functions will not work properly.

    :param file_name: name of hdf5 file without the h5 suffix
    :param key: name of matrix
    :param transpose: whether to transpose the matrix
    :param flatten: whether to flatten matrix with only 1 ROW or COLUMN to a
        ROW vector
    :return: (num_rows, num_cols) complex128 matrix or ROW vector, depending on
        raw data shape and the 'flatten' parameter
    """
    # Data is always a matrix even if it has only 1 row or column.
    data = interface.get_matrix_complex(file_name.encode(), key.encode())
    if transpose:
        data = data.T
    if flatten:
        data = flatten_matrix(data)
    return data
