"""
fileiobase provides general purpose NMR file IO functions and classes
used by multiple nmrglue.fileio modules.

Originally part of NMR glue -- an open source module for working with NMR data
in Python https://github.com/jjhelmus/nmrglue/

If you find nmrglue useful in your research please cite the package as:

J.J. Helmus, C.P. Jaroniec, Nmrglue: An open source Python package for the
analysis of multidimensional NMR data, J. Biomol. NMR 2013, 55, 355-367,
https://doi.org/10.1007/s10858-013-9718-x.
----------------

Copyright Notice and Statement for the nmrglue Project


Copyright (c) 2010-2015 Jonathan J. Helmus
All rights reserved.


Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:


a. Redistributions of source code must retain the above copyright
   notice, this list of conditions and the following disclaimer.


b. Redistributions in binary form must reproduce the above copyright
   notice, this list of conditions and the following disclaimer in the
   documentation and/or other materials provided with the
   distribution.


c. Neither the name of the author nor the names of contributors may
   be used to endorse or promote products derived from this software
   without specific prior written permission.


THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

---------------

"""


import os
import string
import itertools
from functools import reduce

import numpy as np


def create_blank_udic(ndim):
    """
    Create a blank universal dictionary for a spectrum of dimension ndim.
    """
    udic = {
        "ndim": ndim,
    }

    for i in range(ndim):
        udic[i] = {
            "sw": 999.99,       # spectral width in Hz
            "complex": True,    # Quadrature, True when dimension is complex
            "obs": 999.99,      # Observation frequency in MHz
            "car": 999.99,      # Carrier frequency in Hz
            # Number of points in dimension based on the shape of the data array.
            # As such the direct dimension (-1) size is R|I, all indirect
            # dimensions are R+I
            "size": 1,
            "label": ["X", "Y", "Z", "A"][i],   # name of dimension
            # encoding of dimension, ie states, tppi, etc.  The direct dimension
            # should be listed as direct.
            "encoding": "direct" if i == ndim - 1 else "states",
            # time and freq flags for domain of dimension
            "time": True,
            "freq": False,
        }

    return udic


class unit_conversion():
    """
    Provides methods to convert between common NMR units

    Parameters
    ----------
    size : int
        Number of points in dimension (R|I).
    cplex : bool
        True if dimension is complex, False is real.
    sw : float
        Spectral width in Hz.
    obs : float
        Observation frequency in MHz.
    car : float
        Carrier frequency in Hz.

    """
    def __init__(self, size, cplx, sw, obs, car):
        """
        create and set up a unit_conversion object
        """
        # fundamental units
        self._size = int(size)
        self._cplx = bool(cplx)
        self._sw = float(sw)
        self._obs = float(obs)
        self._car = float(car)

        # derived units (these are in ppm)
        self._delta = -self._sw / (self._size * self._obs)
        self._first = self._car / self._obs - self._delta * self._size / 2.

    # individual unit conversion functions
    def __percent2pts(self, percent):
        return percent * (self._size - 1) / 100.0

    def __pts2percent(self, pts):
        return pts * 100 / (self._size - 1.0)

    def __hz2pts(self, hz):
        return ((hz / self._obs) - self._first) / self._delta

    def __pts2hz(self, pts):
        return (pts * self._delta + self._first) * self._obs

    def __ppm2pts(self, ppm):
        return (ppm - self._first) / self._delta

    def __pts2ppm(self, pts):
        return (pts * self._delta) + self._first

    # times based units: seconds, ms, and us
    def __sec2pts(self, sec):
        return sec * self._sw

    def __pts2sec(self, pts):
        return pts * 1. / self._sw

    def __ms2pts(self, ms):
        return ms * self._sw / 1.e3

    def __pts2ms(self, pts):
        return pts * 1.e3 / self._sw

    def __us2pts(self, us):
        return us * self._sw / 1.e6

    def __pts2us(self, pts):
        return pts * 1.e6 / self._sw

    # routers
    def __unit2pnt(self, val, units):
        """
        Convert units to points
        """
        units = units.upper()
        if units == "PPM":
            pts = self.__ppm2pts(val)
        elif units == "HZ":
            pts = self.__hz2pts(val)
        elif units in ("%", "PERCENT"):
            pts = self.__percent2pts(val)
        elif units in ("SEC", "SECOND", "S"):
            pts = self.__sec2pts(val)
        elif units == "MS":
            pts = self.__ms2pts(val)
        elif units == "US":
            pts = self.__us2pts(val)
        else:
            raise ValueError("invalid unit type")
        # if self._cplx:
        #     return pts+round(pts)
        # else:
        return pts

    def __pnt2unit(self, val, units):
        """
        Convert points to units
        """
        units = units.upper()
        # if self._cplx:
        #     val = val-round(val)
        if units == "PPM":
            k = self.__pts2ppm(val)
        elif units == "HZ":
            k = self.__pts2hz(val)
        elif units in ("%", "PERCENT"):
            k = self.__pts2percent(val)
        elif units in ("SEC", "SECOND", "S"):
            k = self.__pts2sec(val)
        elif units == "MS":
            k = self.__pts2ms(val)
        elif units == "US":
            k = self.__pts2us(val)
        else:
            raise ValueError("invalid units")
        return k

    def __str2pnt(self, s):
        """
        Convert string with units to points
        """
        units = s.strip(string.digits + string.whitespace + "." + "-").upper()
        val = float(s.strip(string.ascii_letters + string.whitespace + "%"))
        return self.__unit2pnt(val, units)

    def __convert(self, val, unit=None):
        """
        Convert string or value/unit pair
        """
        if isinstance(val, str):
            return self.__str2pnt(val)
        else:
            if unit is None:
                raise ValueError("invalid unit type")
            return self.__unit2pnt(val, unit)

    # User functions
    def f(self, val, unit=None):
        """
        Convert string or value/unit pair to float
        """
        return self.__convert(val, unit)

    def i(self, val, unit=None):
        """
        Convert string or value/unit pair to integer
        """
        return int(round(self.__convert(val, unit)))

    def ppm(self, val):
        """
        Convert to ppm
        """
        return self.__pnt2unit(val, "PPM")

    def hz(self, val):
        """
        Convert to Hz
        """
        return self.__pnt2unit(val, "HZ")

    def percent(self, val):
        """
        Convert to percent
        """
        return self.__pnt2unit(val, "PERCENT")

    def seconds(self, val):
        """
        Convert to seconds
        """
        return self.__pnt2unit(val, "SEC")

    def sec(self, val):
        """
        Convert to seconds
        """
        return self.__pnt2unit(val, "SEC")

    def ms(self, val):
        """
        Convert to milliseconds (ms)
        """
        return self.__pnt2unit(val, "MS")

    def us(self, val):
        """
        Convert to microseconds (us)
        """
        return self.__pnt2unit(val, "US")

    def unit(self, val, unit):
        """
        Convert val points to unit
        """
        return self.__pnt2unit(val, unit)

    # limits and scales
    @staticmethod
    def percent_limits():
        """
        Return tuple of left and right edges in percent
        """
        return 0.0, 100.0

    def percent_scale(self):
        """
        Return array of percent values
        """
        return np.linspace(0.0, 100.0, self._size)

    def ppm_limits(self):
        """
        Return tuple of left and right edges in ppm
        """
        return self.ppm(0), self.ppm(self._size - 1)

    def ppm_scale(self):
        """
        Return array of ppm values
        """
        x0, x1 = self.ppm_limits()
        return np.linspace(x0, x1, self._size)

    def hz_limits(self):
        """
        Return tuple of left and right edges in Hz
        """
        return self.hz(0), self.hz(self._size - 1)

    def hz_scale(self):
        """
        Return array of Hz values
        """
        x0, x1 = self.hz_limits()
        return np.linspace(x0, x1, self._size)

    def sec_limits(self):
        """
        Return tuple of left and right edges in seconds
        """
        return self.sec(0), self.sec(self._size - 1)

    def sec_scale(self):
        """
        Return array of seconds values
        """
        x0, x1 = self.sec_limits()
        return np.linspace(x0, x1, self._size)

    def ms_limits(self):
        """
        Return tuple of left and right edges in milliseconds
        """
        return self.ms(0), self.ms(self._size - 1)

    def ms_scale(self):
        """
        Return array of seconds values
        """
        x0, x1 = self.ms_limits()
        return np.linspace(x0, x1, self._size)

    def us_limits(self):
        """
        Return tuple of left and right edges in milliseconds
        """
        return self.us(0), self.us(self._size - 1)

    def us_scale(self):
        """
        Return array of seconds values
        """
        x0, x1 = self.us_limits()
        return np.linspace(x0, x1, self._size)

    __call__ = i    # calling the object x is the same as x.i


def uc_from_udic(udic, dim=-1):
    """
    Create a unit conversion object from a Universal dictionary.

    Parameters
    ----------
    udic : dic
        Universal dictionary of spectral parameters.
    dim : int. optional
        Dimension number to create unit conversion object for.  Default is for
        last dimension.

    Returns
    -------
    uc : unit conversion object.
        Unit conversion object for given dimension.
    """
    if dim == -1:
        dim = udic['ndim'] - 1  # last dimension
    adic = udic[dim]
    return unit_conversion(adic['size'], adic['complex'], adic['sw'],
                           adic['obs'], adic['car'])


def uc_from_freqscale(scale, obs, unit='ppm'):
    """
    Create a unit conversion object from a spectrum frequency scale axis.

    Parameters
    ----------
    scale : array like
        array of spectrum axis
    obs : float
        Observation frequency in MHz.
    unit: {'ppm', 'hz', 'khz'}
        The unit of the scale axis.

    Returns
    -------
    uc : unit conversion object.
        Unit conversion object for given axis.
    """
    scale = np.array(scale)
    size = len(scale)

    if unit in ('ppm', 'hz', 'khz'):
        cplex = False

        min = scale.min()
        max = scale.max()

        # The scale needs be corrected by extending each extremum by half the
        # bin width (to convert from centers to edges).
        dx = abs(scale[1] - scale[0])

        if unit == 'ppm':
            sw = ((max + dx / 2.0) - (min - dx / 2.0)) * obs
            car = (min - dx / 2.0 + (max - min) / 2.0) * obs
        elif unit == 'hz':
            sw = ((max + dx / 2.0) - (min - dx / 2.0))
            car = (min - dx / 2.0 + (max - min) / 2.0)
        else:
            # unit is 'kHz':
            sw = ((max + dx / 2.0) - (min - dx / 2.0)) / 1.e3
            car = (min - dx / 2.0 + (max - min) / 2.0) / 1.e3

    else:
        mesg = f'{unit} is not a supported unit.'
        raise ValueError(mesg)

    return unit_conversion(size, cplex, sw, obs, car)


def open_towrite(filename, overwrite=False, mode='wb'):
    """
    Open filename for writing and return file object

    Function checks if file exists (and raises IOError if overwrite=False) and
    creates necessary directories as needed.
    """
    # check if file exists and overwrite if False
    if os.path.exists(filename) and (overwrite is False):
        raise OSError("File exists, recall with overwrite=True")

    p, fn = os.path.split(filename)  # split into filename and path
    # create directories if needed
    if p != '' and os.path.exists(p) is False:
        os.makedirs(p)

    return open(filename, mode)

################################################
# numpy ndarray emulation and helper functions #
################################################

# iterators for ND array


def ndfrom_iter(shape, slices):
    ch = [range(lenx)[sX] for lenx, sX in zip(shape, slices)]
    return itertools.product(*ch)


def ndto_iter(shape, slices):
    ich = [range(len(range(lenx)[sX])) for lenx, sX in zip(shape, slices)]
    return itertools.product(*ich)


def ndtofrom_iter(shape, slices):
    ch = [range(lenx)[sX] for lenx, sX in zip(shape, slices)]
    ich = [range(len(i)) for i in ch]
    return zip(itertools.product(*ich), itertools.product(*ch))


def size_and_ndtofrom_iter(shape, slices):
    ch = [range(lenx)[sX] for lenx, sX in zip(shape, slices)]
    s = [len(i) for i in ch]
    ich = [range(i) for i in s]
    return s, zip(itertools.product(*ich), itertools.product(*ch))


# index2trace and trace2index functions


def index2trace_flat(shape, index):
    """
    Calculate trace number from shape and index of all indirect dimensions
    assuming a flat structure
    """
    # We need to perform:
    # index[0]*shape[1]*...shape[-1] + index[1]*shape[2]*...shape[-1] + ...
    # + index[-1]*shape[-1] + index[-1]
    # To do this we calculate the product of shape[X] elements and multiple
    # by the corresponding index element, index[-1] as added at the beginning
    a = index[-1]
    for i, v in enumerate(index[:-1]):
        mult = reduce(lambda x, y: x * y, shape[i + 1:])
        a = a + mult * v
    return a


def trace2index_flat(shape, ntrace):
    """
    Calculate the index of a trace assuming a flat structure
    """
    # algorithm is to take quotient/remainders of sizes in reverse
    q = ntrace  # seed quotient with remained
    index = []
    for s in shape[:0:-1]:  # loop from last size to 2nd size
        q, r = divmod(q, s)
        index.insert(0, r)
    index.insert(0, q)
    return tuple(index)


def index2trace_opp(shape, index):
    """
    Calculate trace number from shape and index of all indirect dimensions
    assuming a phase ordering opposite the time increments.
    """
    n = len(shape)
    # deal with the phase component
    phases = [v % 2 for v in index]
    nphase = index2trace_flat([2] * n, phases[::-1])
    # deal with the remainder
    pindex = [v // 2 for v in index]
    pshape = [i // 2 for i in shape]
    nbase = index2trace_flat(pshape, pindex)
    return nbase * 2 ** n + nphase


def trace2index_opp(shape, ntrace):
    """
    Calculate the index of a trace assuming opposite phase/time increment
    ordering
    """
    n = len(shape)
    q, r = divmod(ntrace, 2 ** n)
    to_add = list(trace2index_flat([2] * n, r))[::-1]
    pshape = [i // 2 for i in shape]
    base = list(trace2index_flat(pshape, q))
    total = [b * 2 + a for b, a in zip(base, to_add)]
    return tuple(total)


def index2trace_reg(shape, index):
    """
    Calculate trace number from shape and index of all indirect dimensions
    assuming the same  phase and time ordering.
    """
    n = len(shape)
    # deal with the phase component
    phases = [v % 2 for v in index]
    nphase = index2trace_flat([2] * n, phases)
    # deal with the remainder
    pindex = [v // 2 for v in index]
    pshape = [i // 2 for i in shape]
    nbase = index2trace_flat(pshape, pindex)
    return nbase * 2 ** n + nphase


def trace2index_reg(shape, ntrace):
    """
    Calculate the index of a trace assuming the same phase/time increment
    ordering
    """
    n = len(shape)
    q, r = divmod(ntrace, 2 ** n)
    to_add = list(trace2index_flat([2] * n, r))
    pshape = [i // 2 for i in shape]
    base = list(trace2index_flat(pshape, q))
    total = [b * 2 + a for b, a in zip(base, to_add)]
    return tuple(total)

#
# data_nd class
#
# inherited classes should define:
#
#    __init__ which sets up the object and defines at minimum
#
#       self.fshape shape of data on disk (shape when order = (0,1,2...)
#       self.order order of axes, default is (0,1,2,...)
#       self.dtype
#
#   self.__setdimandshape__ can be called to set self.dim and self.shape
#    if they are not set by __init__
#
#    __fgetitem__ which takes well formatted tuples of slices
#    and returns ndarray objects
#
#    __fcopy__ which creates a copy provided only self and order parameters
#


class data_nd:
    """
    Base class for building objects which emulate ndarray objects without
    loading data into memory.  These object have the following properties:

    * slicing operations return ndarray objects
    * can iterate over with expected results
    * transpose and swapaxes functions create a new data_nd object with the
      new axes ordering
    * has ndim, shape, and dtype attributes.

    Notes
    -----

    Classes which are use this class as a base should define the following
    methods:

    __init__ which must set up the object and defines at minimum:

        self.fshape : tuple
            Shape of the data on disk, the shape when order = (0, 1, 2, ..)
        self.order : tuple
            Ordering of the axes
        self.dtype : dtype
            Dtype of the emulated ndarray

        self.__setdimandshape__ should be called if self.dim and self.shape
            are not set up __init__.

    __fgetitem__ which takes a well formatted tuple of slices and returns a
        ndarray object with the selected data

    __fcopy__(self, order) which created a copy of the object with the new
        order

    """

    def __init__(self, order):
        pass

    def __setdimandshape__(self):
        """ Set the ndim and shape attributes from fshape """
        # set ndim and shape
        self.ndim = len(self.fshape)
        self.shape = tuple(self.fshape[i] for i in self.order)

    def __copy__(self):
        """
        create a copy
        """
        return self.__fcopy__(self, self.order)

    def __getitem__(self, key):
        """
        x.__getitem__(y) <==> x[y]
        """
        # convert the key into a list
        if not isinstance(key, tuple):
            rlist = [key]
        else:
            rlist = list(key)

        # remove Ellipsis
        while Ellipsis in rlist:
            i = rlist.index(Ellipsis)
            rlist.pop(i)
            for j in range(self.ndim - len(rlist)):
                rlist.insert(i, slice(None))

        if len(rlist) > self.ndim:
            raise IndexError("invalid index")

        # replace none-slices with slice objects
        for i, v in enumerate(rlist):
            if not isinstance(v, slice):

                # check for out of range indexes
                if v >= self.shape[i]:
                    raise IndexError("index({}) out of range(0 <= index < {}) \
                    in dimension {}".format(v, self.shape[i] - 1, i))

                if v <= (-1 * self.shape[i] - 1):
                    raise IndexError("index({}) out of range(0 <= index < {}) \
                    in dimension {}".format(v, self.shape[i] - 1, i))

                if v < 0:
                    w = self.shape[i] + v
                    rlist[i] = slice(w, w + 1, 1)
                else:
                    rlist[i] = slice(v, v + 1, 1)

        # pad the list with additional dimentions
        for i in range(len(rlist), self.ndim):
            rlist.append(slice(None))

        # reorder the slices into file order
        frlist = [rlist[self.order.index(i)] for i in range(self.ndim)]

        # get the data from disk
        data = self.__fgetitem__(tuple(frlist))

        # re-order the data
        if data.shape != (0,):
            return np.squeeze(data.transpose(self.order))
        else:
            return data

    def __len__(self):
        """
        x._len__ <==> len(x)
        """
        return self.shape[0]

    def __iter__(self):
        """ x.__iter__() <==> iter(x) """
        for index in range(0, self.shape[0]):
            yield self[index]

    def swapaxes(self, axis1, axis2):
        """
        Return object with `axis1` and `axis2` interchanged.
        """
        axis1, axis2 = int(axis1), int(axis2)
        if axis1 < 0:
            axis1 = self.ndim - axis1
        if axis2 < 0:
            axis2 = self.ndim - axis2
        if axis1 >= self.ndim:
            raise ValueError("bad axis1 argument to swapaxes")
        if axis2 >= self.ndim:
            raise ValueError("bad axis2 argument to swapaxes")

        order = list(self.order)
        order[axis1], order[axis2] = order[axis2], order[axis1]
        n = self.__fcopy__(order=order)
        return n

    def transpose(self, *axes):
        """
        Return object with axes transposed.

        Parameters
        ----------
        axes : None, tuple or ints, or `n` ints
            * None or no arguments: reverse order of the axes

            * tuple of ints: `i` in the `j`-th place in the tuples means the
            'i'-th axis becomes the new objects `j`-th axis.

            * `n` ints: same as an n-tuple.

        Returns
        -------
        out : data_nd object
            Object whose axes are permuted.

        """
        if axes == ():    # default is to switch order of axes
            axes = range(self.ndim)[::-1]

        if len(axes) == 1:    # if a single tuple is given unpack
            axes = axes[0]

        try:    # convert to integers
            axes = [int(i) for i in axes]
        except Exception:
            raise TypeError("an integer is required")

        if len(axes) != self.ndim:   # check for to few/many axes
            raise ValueError("axes don't match array")

        # replace negatives axes values with positives
        for i, v in enumerate(axes):
            if v < 0:
                axes[i] = self.ndim + v

        # check for invalid axes
        for v in axes:
            if v >= self.ndim:
                raise ValueError("invalid axis for this array")

        # check for repeated axes
        if len(set(axes)) != self.ndim:
            raise ValueError("repeated axis in tranpose")

        # create a new data_nd object with transposed order
        return self.__fcopy__(order=tuple(self.order[i] for i in axes))
