__all__ = ['file_data_store', 'file_data_read', 'qsave', 'qload']

import pickle
import numpy as np
from .core import Qobj
from pathlib import Path


# -----------------------------------------------------------------------------
# Write matrix data to a file
#
def file_data_store(filename, data, numtype="complex", numformat="decimal",
                    sep=","):
    """Stores a matrix of data to a file to be read by an external program.

    Parameters
    ----------
    filename : str or pathlib.Path
        Name of data file to be stored, including extension.
    data: array_like
        Data to be written to file.
    numtype : str {'complex, 'real'}, default: 'complex'
        Type of numerical data.
    numformat : str {'decimal','exp'}, default: 'decimal'
        Format for written data.
    sep : str, default: ','
        Single-character field seperator.  Usually a tab, space, comma,
        or semicolon.

    """
    if filename is None or data is None:
        raise ValueError("filename or data is unspecified")

    M, N = np.shape(data)

    with open(filename, "w") as f:
        f.write(f"# Generated by QuTiP: {M}x{N} {numtype} matrix " +
                f"in {numformat} format ['{sep}' separated values].\n")

        if numtype == "complex":

            if numformat == "exp":

                for m in range(M):
                    for n in range(N):
                        f.write(f"{np.real(data[m, n]):.10e}{np.imag(data[m, n]):+.10e}j")
                        if n != N - 1:
                            f.write(sep)
                    f.write("\n")

            elif numformat == "decimal":

                for m in range(M):
                    for n in range(N):
                        f.write(f"{np.real(data[m, n]):.10f}{np.imag(data[m, n]):+.10f}j")
                        if n != N - 1:
                            f.write(sep)
                    f.write("\n")

            else:
                raise ValueError("Illegal numformat value (should be " +
                                "'exp' or 'decimal')")

        elif numtype == "real":

            if numformat == "exp":

                for m in range(M):
                    for n in range(N):
                        f.write(f"{np.real(data[m, n]):.10e}")
                        if n != N - 1:
                            f.write(sep)
                    f.write("\n")

            elif numformat == "decimal":

                for m in range(M):
                    for n in range(N):
                        f.write(f"{np.real(data[m, n]):.10f}")
                        if n != N - 1:
                            f.write(sep)
                    f.write("\n")

            else:
                raise ValueError("Illegal numformat value (should be " +
                                "'exp' or 'decimal')")

        else:
            raise ValueError("Illegal numtype value (should be " +
                            "'complex' or 'real')")


# -----------------------------------------------------------------------------
# Read matrix data from a file
#
def file_data_read(filename, sep=None):
    """Retrieves an array of data from the requested file.

    Parameters
    ----------
    filename : str or pathlib.Path
        Name of file containing reqested data.
    sep : str, optional
        Seperator used to store data.

    Returns
    -------
    data : array_like
        Data from selected file.

    """
    if filename is None:
        raise ValueError("filename is unspecified")

    with open(filename, "r") as f:

        # first count lines and numbers of
        M = N = 0
        for line in f:
            # skip comment lines
            if line[0] == '#' or line[0] == '%':
                continue
            # find delim
            if N == 0 and sep is None:
                if len(line.rstrip().split(",")) > 1:
                    sep = ","
                elif len(line.rstrip().split(";")) > 1:
                    sep = ";"
                elif len(line.rstrip().split(":")) > 1:
                    sep = ":"
                elif len(line.rstrip().split("|")) > 1:
                    sep = "|"
                elif len(line.rstrip().split()) > 1:
                    # sepical case for a mix of white space deliminators
                    sep = None
                else:
                    raise ValueError("Unrecognized column deliminator")
            # split the line
            line_vec = line.split(sep)
            n = len(line_vec)
            if N == 0 and n > 0:
                N = n
                # check type
                if ("j" in line_vec[0]) or ("i" in line_vec[0]):
                    numtype = "complex"
                else:
                    numtype = "np.real"

                # check format
                if ("e" in line_vec[0]) or ("E" in line_vec[0]):
                    numformat = "exp"
                else:
                    numformat = "decimal"

            elif N != n:
                raise ValueError("Badly formatted data file: " +
                                "unequal number of columns")
            M += 1

        #
        # read data and store in a matrix
        #
        f.seek(0)

        if numtype == "complex":
            data = np.zeros((M, N), dtype="complex")
            m = n = 0
            for line in f:
                # skip comment lines
                if line[0] == '#' or line[0] == '%':
                    continue
                n = 0
                for item in line.rstrip().split(sep):
                    data[m, n] = complex(item)
                    n += 1
                m += 1

        else:
            data = np.zeros((M, N), dtype="float")
            m = n = 0
            for line in f:
                # skip comment lines
                if line[0] == '#' or line[0] == '%':
                    continue
                n = 0
                for item in line.rstrip().split(sep):
                    data[m, n] = float(item)
                    n += 1
                m += 1

    return data


def qsave(data, filename='qutip_data'):
    """
    Saves given data to file named 'filename.qu' in current directory.

    Parameters
    ----------
    data : instance/array_like
        Input Python object to be stored.
    filename : str or pathlib.Path, default: "qutip_data"
        Name of output data file.

    """
    # open the file for writing
    path = Path(filename)
    path = path.with_suffix(path.suffix + ".qu")

    with open(path, "wb") as fileObject:
        # this writes the object a to the file named 'filename.qu'
        pickle.dump(data, fileObject)


def qload(filename):
    """
    Loads data file from file ``filename`` in current directory.

    Parameters
    ----------
    filename : str or pathlib.Path
        Name of data file to be loaded.

    Returns
    -------
    qobject : instance / array_like
        Object retrieved from requested file.

    """
    path = Path(filename)
    path = path.with_suffix(path.suffix + ".qu")

    with open(path, "rb") as fileObject:
        out = pickle.load(fileObject)
    return out
