"""
Utilities to interact with, manipulate or call on the external solver, 
i.e., SPECFEM2D/3D/3D_GLOBE
"""
import os
import numpy as np
from glob import glob
from seisflows import logger
from seisflows.tools import msg


def check_source_names(path_specfem_data, source_prefix, ntask=None):
    """
    Determines names of sources by applying wildcard rule to user-supplied
    input files. Source names are only provided up to PAR.NTASK and are
    returned in alphabetical order.

    .. note::
        SeisFlows expects sources to be stored in the DATA/ directory with a
        prefix and a source name, e.g., {source_prefix}_{source_name} which
        would evaluate to something like CMTSOLUTION_001

    :type path_specfem_data: str
    :param path_specfem_data: path to a
    :type source_prefix: str
    :param source_prefix: type of SPECFEM input source, e.g., CMTSOLUTION
    :type ntask: int
    :parma ntask: if provided, curtails the list of sources up to `ntask`. If
        None, returns all files found matching the wildcard
    :rtype: list
    :return: alphabetically ordered list of source names up to PAR.NTASK
    """
    wildcard = f"{source_prefix}_*"
    fids = sorted(glob(os.path.join(path_specfem_data, wildcard)))
    if not fids:
        print(msg.cli("No matching source files when searching PATH for the "
                      "given WILDCARD", header="source error", border="=",
                      items=[f"PATH: {path_specfem_data}",
                             f"WILDCARD: {wildcard}"]))
        return
    if ntask is not None:
        assert(len(fids) >= ntask), (
            f"Number of requested tasks/events {ntask} exceeds number "
            f"of available sources {len(fids)}"
        )
        fids = fids[:ntask]

    # Create internal definition of sources names by stripping prefixes
    names = [os.path.basename(fid).split("_")[-1] for fid in fids]

    return names


def getpar(key, file, delim="=", match_partial=False):
    """
    Reads and returns parameters from a SPECFEM or SeisFlows parameter file
    Assumes the parameter file is formatted in the following way:

    # comment comment comment
    {key}      {delim} VAL

    :type key: str
    :param key: case-insensitive key to match in par_file. must be EXACT match
    :type file: str
    :param file: The SPECFEM Par_file to match against
    :type delim: str
    :param delim: delimiter between parameters and values within the file.
        default is '=', which matches for SPECFEM2D and SPECFEM3D_Cartesian
    :type match_partial: bool
    :param match_partial: allow partial key matches, e.g., allow key='tit' to
        return value for 'title'. Defaults to False as this can have
        unintended consequences
    :rtype: tuple (str, str, int)
    :return: a tuple of the key, value and line number (indexed from 0).
        The key will match exactly how it looks in the Par_file
        The value will be returned as a string, regardless of its expected type
        IF no matches found, returns (None, None, None)
    """
    lines = open(file, "r").readlines()

    for i, line in enumerate(lines):
        # Find the first occurence, CASE-INSENSITIVE search, strip whitespace
        # To allow for nested parameters
        if line.strip().upper().startswith(key.upper()):
            try:
                key_out, val = line.strip().split(delim)
            except ValueError as e:
                raise ValueError(
                    f"Could not split line with delimiter '{delim}'. Error "
                    f"message: {e}"
                )
            # Strip white space now so we can match keys
            key_out = key_out.strip()
            # Unless explcitely authorized, don't allow partial matches,
            # which str.find() will return on. Can have unintended consequences
            # e.g., 'tit' will return 'title'
            if (not match_partial) and (key_out.upper() != key.upper()):
                continue
            # Drop any trailing line comments
            if "#" in val:
                val = val.split("#")[0]
            # One last strip to remove any whitespace
            val = val.strip()
            # Address the fact that SPECFEM Par_file sometimes lists values as
            # formatted strings, e.g., 38.0d-2
            try:
                if len(val.split("d")) == 2:
                    num, exp = val.split("d")
                    val = str(float(num) * 10 ** int(exp))
            except ValueError as e:
                # This will break on anything other than than the above type str
                pass
            break
    else:
        raise KeyError(f"Could not find matching key '{key}' in file: {file}")
    return key_out, val, i


def setpar(key, val, file, delim="=", match_partial=False):
    """
    Overwrites parameter value to a SPECFEM Par_file.

    :type key: str
    :param key: case-insensitive key to match in par_file. must be EXACT match
    :type val: str
    :param val: value to OVERWRITE to the given key
    :type file: str
    :param file: The SPECFEM Par_file to match against
    :type delim: str
    :param delim: delimiter between parameters and values within the file.
        default is '=', which matches for SPECFEM2D and SPECFEM3D_Cartesian
    :type match_partial: bool
    :param match_partial: allow partial key matches, e.g., allow key='tit' to
        return value for 'title'. Defaults to False as this can have
        unintended consequences
    """
    key_out, val_out, i = getpar(key, file, delim, match_partial)
    if key_out is None:
        return

    lines = open(file, "r").readlines()

    # Replace value in place by splitting on the delimiter and replacing the 
    # first instance of the old value (avoids replacing comment values or 
    # matching values inside the key)
    if val_out != "":
        key, val_and_comment = lines[i].split(delim)
        val_and_comment = val_and_comment.replace(val_out, str(val), 1)
        lines[i] = delim.join([key, val_and_comment])
    else:
        # Special case where the initial parameter is empty so we just replace
        # the newline formatter at the end
        lines[i] = lines[i].replace("\n", f" {val}\n")
    with open(file, "w") as f:
        f.writelines(lines)


def _getidx_vel_model(lines):
    """
    Get the line indices of a velocity model, which can be used to retrieve 
    or replace the model values in a SPECFEM2D paramter file. Used by 
    `getpar_vel_model` and `setpar_vel_model`

    :type lines: list
    :param lines: list of strings read from the par_file
    :rtype idxs: list
    :param idxs: list of integer indices of the velocity model lines
    """
    idxs = []
    for l, line in enumerate(lines):
        # Skip over all other parameters, comments, newlines etc.
        if "=" in line:
            continue
        elif line.startswith(" "):
            continue
        elif line.startswith("#"):
            continue
        elif line.startswith("\n"):
            continue

        # Matches formatting expected by velocity model and starts with integer
        # Should be enough to avoid matching other strings
        lines = line.strip().split()
        if len(lines) == 15 and lines[0].isdigit():
            idxs.append(l)

    return idxs


def getpar_vel_model(file, strip=False):
    """
    SPECFEM2D doesn't follow standard key = val formatting when defining its 
    internal velocity models so we need a special function to address this 
    specifically.
    
    Velocity models are ASSUMED to be formatted in the following way

    1 1 2700.d0 3000.d0 1732.051d0 0 0 9999 9999 0 0 0 0 0 0

    That is, 15 entries separated by spaces. We use that to find all relevant 
    lines of the model.
    
    :type file: str
    :param file: The SPECFEM Par_file to match against
    :type strip: bool
    :param strip: strip newline '\n' from each of the model lines
    :rtype: list of str
    :return: list of all the layers of the velocity model as strings
    """
    # i_start + 1 to skip over the 'nbmodels' parameter
    lines = open(file, "r").readlines()
    idxs = _getidx_vel_model(lines)
    vel_model = []
    for idx in idxs:
        if strip:
            line = lines[idx].strip()
        else:
            line = lines[idx]
        vel_model.append(line)

    return vel_model


def setpar_vel_model(file, model):
    """
    Set velocity model values in a SPECFEM2D Par_file, see getpar_vel_model
    for more information.

    Deletes the old model from the Par_file, writes the new model in the same
    place, and then changes the value of 'nbmodels'

    :type file: str
    :param file: The SPECFEM Par_file to match against
    :type model: list of str
    :param model: input model
    :rtype: list of str
    :return: list of all the layers of the velocity model as strings, e.g.:
        model = ["1 1 2700.d0 3000.d0 1732.051d0 0 0 9999 9999 0 0 0 0 0 0",
                 "2 1 2500.d0 2700.d0 0 0 0 9999 9999 0 0 0 0 0 0"]
    """
    lines = open(file, "r").readlines()
    model_lines = _getidx_vel_model(open(file, "r").readlines())

    # one-liner to get rid of the original model from the file
    lines = [i for j, i in enumerate(lines) if j not in model_lines]
    model_idx_start = model_lines[0]

    # Throw a new line onto the last line of the model to get proper formatting
    model[-1] = model[-1] + "\n"

    # Drop in the new model one line at a time
    for i, val in enumerate(model):
        lines.insert(model_idx_start + i, f"{val.strip()}\n")

    # Overwrite with new lines containing updated velocity model
    with open(file, "w") as f:
        f.writelines(lines)

    # Set nbmodels to the correct value
    setpar(key="nbmodels", val=len(model), file=file)


def read_fortran_binary(filename):
    """
    Reads Fortran-style unformatted binary data into numpy array.

    .. note::
        The FORTRAN runtime system embeds the record boundaries in the data by
        inserting an INTEGER*4 byte count at the beginning and end of each
        unformatted sequential record during an unformatted sequential WRITE.
        see: https://docs.oracle.com/cd/E19957-01/805-4939/6j4m0vnc4/index.html

    :type filename: str
    :param filename: full path to the Fortran unformatted binary file to read
    :rtype: np.array
    :return: numpy array with data with data read in as type Float32
    """
    nbytes = os.path.getsize(filename)
    with open(filename, "rb") as file:
        # read size of record
        file.seek(0)
        n = np.fromfile(file, dtype="int32", count=1)[0]

        if n == nbytes - 8:
            file.seek(4)
            data = np.fromfile(file, dtype="float32")
            return data[:-1]
        else:
            file.seek(0)
            data = np.fromfile(file, dtype="float32")
            return data


def write_fortran_binary(arr, filename):
    """
    Writes Fortran style binary files. Data are written as single precision
    floating point numbers.

    .. note::
        FORTRAN unformatted binaries are bounded by an INT*4 byte count. This
        function mimics that behavior by tacking on the boundary data.
        https://docs.oracle.com/cd/E19957-01/805-4939/6j4m0vnc4/index.html

    :type arr: np.array
    :param arr: data array to write as Fortran binary
    :type filename: str
    :param filename: full path to file that should be written in format
        unformatted Fortran binary
    """
    buffer = np.array([4 * len(arr)], dtype="int32")
    data = np.array(arr, dtype="float32")

    with open(filename, "wb") as file:
        buffer.tofile(file)
        data.tofile(file)
        buffer.tofile(file)
