#cython: language_level=3
import inspect
import pickle
import scipy
from scipy.interpolate import make_interp_spline
import numpy as np
cimport numpy as cnp
cimport cython
import qutip

cdef extern from "<complex>" namespace "std" nogil:
    double complex conj(double complex x)
    double         norm(double complex x)


def coefficient_function_parameters(func, style=None):
    """
    Return the function style (either "pythonic" or not) and a list of
    additional parameters accepted.

    Used by :obj:`FunctionCoefficient` and :obj:`_FuncElement` to determine the
    call signature of the supplied function based on the given style (or
    ``qutip.settings.core["function_coefficient_style"]`` if no style is given)
    and the signature of the given function.

    Parameters
    ----------
    func : function
        The :obj:`FunctionCoefficient` to inspect. The first argument
        of the function is assumed to be ``t`` (the time at which to
        evaluate the coefficient). The remaining arguments depend on
        the signature style (see below).

    style : {None, "pythonic", "dict", "auto"}
        The style of the signature used. If style is ``None``,
        the value of ``qutip.settings.core["function_coefficient_style"]``
        is used. Otherwise the supplied value overrides the global setting.

    Returns
    -------
    (f_pythonic, f_parameters)

    f_pythonic : bool
        True if the function should be called as ``f(t, **kw)`` and False
        if the function should be called as ``f(t, kw_dict)``.

    f_parameters : set or None
        The set of parameters (other than ``t``) of the function or
        ``None`` if the function accepts arbitrary parameters.
    """
    sig = inspect.signature(func)
    f_has_kw = any(p.kind == p.VAR_KEYWORD for p in sig.parameters.values())
    if style is None:
        style = qutip.settings.core["function_coefficient_style"]
    if style == "auto":
        if tuple(sig.parameters.keys()) == ("t", "args") and not f_has_kw:
            # if the signature is exactly f(t, args), then assume parameters
            # are supplied in an argument dictionary
            style = "dict"
        else:
            style = "pythonic"
    if style == "dict" or f_has_kw:
        # f might accept any parameter
        f_parameters = None
    else:
        # f accepts only t and the named parameters
        f_parameters = set(list(sig.parameters.keys())[1:])
    return (style == "pythonic", f_parameters)


cdef class Coefficient:
    """
    `Coefficient` are the time-dependant scalar of a `[Qobj, coeff]` pair
    composing time-dependant operator in list format for :obj:`QobjEvo`.

    :obj:`Coefficient` are immutable.
    """
    def __init__(self):
        raise NotImplementedError("Only sub-classes should be initiated.")

    def replace_arguments(self, _args=None, **kwargs):
        """
        Replace the arguments (``args``) of a coefficient.

        Returns a new :obj:`Coefficient` if the coefficient has arguments, or
        the original coefficient if it does not. Arguments to replace may be
        supplied either in a dictionary as the first position argument, or
        passed as keywords, or as a combination of the two. Arguments not
        replaced retain their previous values.

        Parameters
        ----------
        _args : dict
            Dictionary of arguments to replace.

        **kwargs
            Arguments to replace.
        """
        return self

    def __call__(self, double t, dict _args=None, **kwargs):
        """
        Return the coefficient value at time `t`.
        Stored arguments can overwriten with `_args` or as keywords parameters.

        Parameters
        ----------
        t : float
            Time at which to evaluate the :obj:`Coefficient`.

        _args : dict
            Dictionary of arguments to use instead of the stored ones.

        **kwargs
            Arguments to overwrite for this call.
        """
        if _args is not None or kwargs:
            return (<Coefficient> self.replace_arguments(_args,
                                                         **kwargs))._call(t)
        return self._call(t)

    cdef double complex _call(self, double t) except *:
        """Core computation of the :obj:`Coefficient`."""
        raise NotImplementedError("All Coefficient sub-classes "
                                  "should overwrite this.")

    def __add__(left, right):
        if (
            isinstance(left, InterCoefficient)
            and isinstance(right, InterCoefficient)
        ):
            return add_inter(left, right)
        if isinstance(left, Coefficient) and isinstance(right, Coefficient):
            return SumCoefficient(left.copy(), right.copy())
        return NotImplemented

    def __mul__(left, right):
        if isinstance(left, Coefficient) and isinstance(right, Coefficient):
            return MulCoefficient(left.copy(), right.copy())
        if isinstance(left, qutip.Qobj):
            return qutip.QobjEvo([left.copy(), right.copy()])
        if isinstance(right, qutip.Qobj):
            return qutip.QobjEvo([right.copy(), left.copy()])
        return NotImplemented

    cpdef Coefficient copy(self):
        """Return a copy of the :obj:`Coefficient`."""
        return pickle.loads(pickle.dumps(self))

    def conj(self):
        """ Return a conjugate :obj:`Coefficient` of this"""
        return ConjCoefficient(self)

    def _cdc(self):
        """ Return a :obj:`Coefficient` being the norm of this"""
        return NormCoefficient(self)


@cython.auto_pickle(True)
cdef class FunctionCoefficient(Coefficient):
    """
    :obj:`Coefficient` wrapping a Python function.

    Parameters
    ----------
    func : callable(t : float, ...) -> complex
        Function returning the coefficient value at time ``t``.

    args : dict
        Values of the arguments to pass to ``func``.

    style : {None, "pythonic", "dict", "auto"}
        The style of function signature used. If style is ``None``,
        the value of ``qutip.settings.core["function_coefficient_style"]``
        is used. Otherwise the supplied value overrides the global setting.

    The parameters ``_f_pythonic`` and ``_f_parameters`` override function
    style and parameter detection and are not intended to be part of
    the public interface.
    """
    cdef object func
    cdef bint _f_pythonic
    cdef set _f_parameters

    _UNSET = object()

    def __init__(self, func, dict args, style=None, _f_pythonic=_UNSET,
                 _f_parameters=_UNSET):
        if _f_pythonic is self._UNSET or _f_parameters is self._UNSET:
            if not (_f_pythonic is self._UNSET
                    and _f_parameters is self._UNSET):
                raise TypeError(
                    "_f_pythonic and _f_parameters should "
                    "always be given together."
                )
            _f_pythonic, _f_parameters = coefficient_function_parameters(
                func, style=style)
            if _f_parameters is not None:
                args = {k: args[k] for k in _f_parameters & args.keys()}
            else:
                args = args.copy()
        self.func = func
        self.args = args
        self._f_pythonic = _f_pythonic
        self._f_parameters = _f_parameters

    cdef complex _call(self, double t) except *:
        if self._f_pythonic:
            return self.func(t, **self.args)
        return self.func(t, self.args)

    cpdef Coefficient copy(self):
        """Return a copy of the :obj:`Coefficient`."""
        return FunctionCoefficient(
            self.func,
            self.args.copy(),
            _f_pythonic=self._f_pythonic,
            _f_parameters=self._f_parameters,
        )

    def replace_arguments(self, _args=None, **kwargs):
        """
        Replace the arguments (``args``) of a coefficient.

        Returns a new :obj:`Coefficient` if the coefficient has arguments, or
        the original coefficient if it does not. Arguments to replace may be
        supplied either in a dictionary as the first position argument, or
        passed as keywords, or as a combination of the two. Arguments not
        replaced retain their previous values.

        Parameters
        ----------
        _args : dict
            Dictionary of arguments to replace.

        **kwargs
            Arguments to replace.
        """
        if _args:
            kwargs.update(_args)
        if self._f_parameters is not None:
            kwargs = {k: kwargs[k] for k in self._f_parameters & kwargs.keys()}
        if not kwargs:
            return self
        return FunctionCoefficient(
            self.func,
            {**self.args, **kwargs},
            _f_pythonic=self._f_pythonic,
            _f_parameters=self._f_parameters,
        )


def proj(x):
    if np.isfinite(x):
        return (x)
    else:
        return np.inf + 0j * np.imag(x)


cdef class StrFunctionCoefficient(Coefficient):
    """
    A :obj:`Coefficient` defined by a string containing a simple Python
    expression.

    The string should contain a compilable Python expression that results in a
    complex number. The time ``t`` is available as a local variable, as are the
    individual arguments (i.e. the keys of ``args``). The ``args`` dictionary
    itself is not accessible.

    The following symbols are defined:
        ``sin``, ``cos``, ``tan``, ``asin``, ``acos``, ``atan``, ``pi``,
        ``sinh``, ``cosh``, ``tanh``, ``asinh``, ``acosh``, ``atanh``,
        ``exp``, ``log``, ``log10``, ``erf``, ``zerf``, ``sqrt``,
        ``real``, ``imag``, ``conj``, ``abs``, ``norm``, ``arg``, ``proj``,
        ``numpy`` as ``np`` and ``scipy.special`` as ``spe``.

    Examples
    --------
    >>> StrFunctionCoefficient("sin(w * pi * t)", {'w': 1j})

    Parameters
    ----------
    base : str
        A string representing a compilable Python expression that results in a
        complex number.

    args : dict
        A dictionary of variable used in the code string. It may include unused
        variables.
    """
    cdef object func
    cdef str base

    str_env = {
        "sin": np.sin,
        "cos": np.cos,
        "tan": np.tan,
        "asin": np.arcsin,
        "acos": np.arccos,
        "atan": np.arctan,
        "pi": np.pi,
        "sinh": np.sinh,
        "cosh": np.cosh,
        "tanh": np.tanh,
        "asinh": np.arcsinh,
        "acosh": np.arccosh,
        "atanh": np.arctanh,
        "exp": np.exp,
        "log": np.log,
        "log10": np.log10,
        "erf": scipy.special.erf,
        "zerf": scipy.special.erf,
        "sqrt": np.sqrt,
        "real": np.real,
        "imag": np.imag,
        "conj": np.conj,
        "abs": np.abs,
        "norm": lambda x: np.abs(x)**2,
        "arg": np.angle,
        "proj": proj,
        "np": np,
        "spe": scipy.special}

    def __init__(self, base, dict args):
        args2var = "\n".join(["    {} = args['{}']".format(key, key)
                              for key in args])
        code = f"""
def coeff(t, args):
{args2var}
    return {base}"""
        lc = {}
        exec(code, self.str_env, lc)
        self.base = base
        self.func = lc["coeff"]
        self.args = args

    cdef complex _call(self, double t) except *:
        return self.func(t, self.args)

    cpdef Coefficient copy(self):
        """Return a copy of the :obj:`Coefficient`."""
        return StrFunctionCoefficient(self.base, self.args.copy())

    def __reduce__(self):
        return (StrFunctionCoefficient, (self.base, self.args))

    def replace_arguments(self, _args=None, **kwargs):
        """
        Replace the arguments (``args``) of a coefficient.

        Returns a new :obj:`Coefficient` if the coefficient has arguments, or
        the original coefficient if it does not. Arguments to replace may be
        supplied either in a dictionary as the first position argument, or
        passed as keywords, or as a combination of the two. Arguments not
        replaced retain their previous values.

        Parameters
        ----------
        _args : dict
            Dictionary of arguments to replace.

        **kwargs
            Arguments to replace.
        """
        if _args:
            kwargs.update(_args)
        if kwargs:
            return StrFunctionCoefficient(self.base, {**self.args, **kwargs})
        return self


cdef class InterCoefficient(Coefficient):
    """
    A :obj:`Coefficient` built from an interpolation of a numpy array.

    Parameters
    ----------
    coeff_arr : np.ndarray
        The array of coefficient values to interpolate.

    tlist : np.ndarray
        An array of times corresponding to each coefficient value. The times
        must be increasing, but do not need to be uniformly spaced.

    order : int
        Order of the interpolation. Order ``0`` uses the previous (i.e. left)
        value. The order will be reduced to ``len(tlist) - 1`` if it is larger.
    """
    cdef int order
    cdef double dt
    cdef double[::1] tlist
    cdef complex[:, :] poly
    cdef object np_arrays

    def __init__(self, coeff_arr, tlist, int order):
        tlist = np.array(tlist, dtype=np.float64)
        coeff_arr = np.array(coeff_arr, dtype=np.complex128)

        if coeff_arr.ndim != 1:
            raise ValueError("The array to interpolate must be a 1D array")
        if coeff_arr.shape != tlist.shape:
            raise ValueError("tlist must be the same len "
                             "as the array to interpolate")
        if order < 0:
            raise ValueError("order must be a positive integer")

        order = min(order, len(tlist) - 1)

        if order == 0:
            coeff_arr = coeff_arr.reshape((1, -1))
        elif order == 1:
            coeff_arr = np.vstack([
                    np.diff(coeff_arr, append=-1) / np.diff(tlist, append=-1),
                    coeff_arr
                ])
        elif order >= 2:
            # Use scipy to compute the spline and transform it to polynomes
            # as used in scipy's PPoly which is easier for us to use.
            spline = make_interp_spline(tlist, coeff_arr, k=order)
            # Scipy can move knots, we add them to tlist
            tlist = np.sort(np.unique(np.concatenate([spline.t, tlist])))
            a = np.arange(spline.k+1)
            a[0] = 1
            fact = np.cumprod(a)
            coeff_arr = np.concatenate([
                spline(tlist, i) / fact[i]
                for i in range(spline.k, -1, -1)
            ]).reshape((spline.k+1, -1))

        self._prepare(tlist, coeff_arr)

    def _prepare(self, np_tlist, np_poly, dt=None):
        self.np_arrays = (np_tlist, np_poly)
        self.tlist = np_tlist
        self.poly = np_poly
        self.order = self.poly.shape[0] - 1
        diff = np.diff(self.np_arrays[0])
        if dt is not None:
            self.dt = dt
        elif len(diff) >= 1 and np.allclose(diff[0], diff):
            self.dt = diff[0]
        else:
            self.dt = 0

    @cython.wraparound(False)
    @cython.boundscheck(False)
    @cython.cdivision(True)
    cdef size_t _binary_search(self, double x):
        # Binary search for the interval
        # return the indice of the of the biggest element where t <= x
        cdef size_t low = 0
        cdef size_t high = self.tlist.shape[0]
        cdef size_t middle
        cdef size_t count = 0
        while low+1 != high and count < 64:
            middle = (low + high)//2
            if x < self.tlist[middle]:
                high = middle
            else:
                low = middle
            # We keep a count to be sure that it never get into an infinit loop
            # even if tlist has an unexpected format.
            count += 1
        return low

    @cython.initializedcheck(False)
    @cython.cdivision(True)
    cdef double complex _call(self, double t) except *:
        cdef size_t idx, i
        cdef double factor
        cdef double complex out
        cdef double complex[:] slice
        if t <= self.tlist[0]:
            return self.poly[-1, 0]
        elif t >= self.tlist[-1]:
            return self.poly[-1, -1]
        if self.dt:
            idx = <size_t>((t - self.tlist[0]) / self.dt)
        else:
            idx = self._binary_search(t)
        if self.order == 0:
            out = self.poly[0, idx]
        else:
            factor = t - self.tlist[idx]
            slice = self.poly[:, idx]
            out = 0.
            for i in range(self.order+1):
                out *= factor
                out += slice[i]
        return out

    def __reduce__(self):
        return (InterCoefficient.restore, (*self.np_arrays, self.dt))

    @classmethod
    def restore(cls, np_tlist, np_poly, dt=None):
        cdef InterCoefficient out = cls.__new__(cls)
        out._prepare(np_tlist, np_poly, dt)
        return out

    @classmethod
    def from_PPoly(cls, ppoly):
        return cls.restore(ppoly.x, ppoly.c)

    @classmethod
    def from_Bspline(cls, spline):
        tlist = np.unique(spline.t)
        a = np.arange(spline.k+1)
        a[0] = 1
        fact = np.cumprod(a) + 0j
        poly = np.concatenate([
            spline(tlist, i) / fact[i]
            for i in range(spline.k, -1, -1)
        ]).reshape((spline.k+1, -1))
        return cls.restore(tlist, poly)

    cpdef Coefficient copy(self):
        """Return a copy of the :obj:`Coefficient`."""
        return InterCoefficient.restore(*self.np_arrays, self.dt)


cdef Coefficient add_inter(InterCoefficient left, InterCoefficient right):
    """ Add two array coefficient with matching tlist into one."""
    if (
        left.np_arrays[0].shape == right.np_arrays[0].shape
        and np.allclose(left.np_arrays[0], right.np_arrays[0],
                        rtol=1e-15, atol=1e-15)
        and (left.order == right.order)
    ):
        return InterCoefficient.restore(
            left.np_arrays[0], left.np_arrays[1] + right.np_arrays[1],
            left.dt
        )
    else:
        # It would be possible to add them by merging tlist.
        return SumCoefficient(left, right)


@cython.auto_pickle(True)
cdef class SumCoefficient(Coefficient):
    """
    A :obj:`Coefficient` built from the sum of two other coefficients.

    A :obj:`SumCoefficient` is returned as the result of the addition of two
    coefficients, e.g. ::

        coefficient("t * t") + coefficient("t")  # SumCoefficient
    """
    cdef Coefficient first
    cdef Coefficient second

    def __init__(self, Coefficient first, Coefficient second):
        self.first = first
        self.second = second

    cdef complex _call(self, double t) except *:
        return self.first._call(t) + self.second._call(t)

    cpdef Coefficient copy(self):
        """Return a copy of the :obj:`Coefficient`."""
        return SumCoefficient(self.first.copy(), self.second.copy())

    def replace_arguments(self, _args=None, **kwargs):
        """
        Replace the arguments (``args``) of a coefficient.

        Returns a new :obj:`Coefficient` if the coefficient has arguments, or
        the original coefficient if it does not. Arguments to replace may be
        supplied either in a dictionary as the first position argument, or
        passed as keywords, or as a combination of the two. Arguments not
        replaced retain their previous values.

        Parameters
        ----------
        _args : dict
            Dictionary of arguments to replace.

        **kwargs
            Arguments to replace.
        """
        return SumCoefficient(
            self.first.replace_arguments(_args, **kwargs),
            self.second.replace_arguments(_args, **kwargs)
        )


@cython.auto_pickle(True)
cdef class MulCoefficient(Coefficient):
    """
    A :obj:`Coefficient` built from the product of two other coefficients.

    A :obj:`MulCoefficient` is returned as the result of the multiplication of
    two coefficients, e.g. ::

        coefficient("w * t", args={'w': 1}) * coefficient("t")
    """
    cdef Coefficient first
    cdef Coefficient second

    def __init__(self, Coefficient first, Coefficient second):
        self.first = first
        self.second = second

    cdef complex _call(self, double t) except *:
        return self.first._call(t) * self.second._call(t)

    cpdef Coefficient copy(self):
        """Return a copy of the :obj:`Coefficient`."""
        return MulCoefficient(self.first.copy(), self.second.copy())

    def replace_arguments(self, _args=None, **kwargs):
        """
        Replace the arguments (``args``) of a coefficient.

        Returns a new :obj:`Coefficient` if the coefficient has arguments, or
        the original coefficient if it does not. Arguments to replace may be
        supplied either in a dictionary as the first position argument, or
        passed as keywords, or as a combination of the two. Arguments not
        replaced retain their previous values.

        Parameters
        ----------
        _args : dict
            Dictionary of arguments to replace.

        **kwargs
            Arguments to replace.
        """
        return MulCoefficient(
            self.first.replace_arguments(_args, **kwargs),
            self.second.replace_arguments(_args, **kwargs)
        )


@cython.auto_pickle(True)
cdef class ConjCoefficient(Coefficient):
    """
    The conjugate of a :obj:`Coefficient`.

    A :obj:`ConjCoefficient` is returned by ``Coefficient.conj()`` and
    ``qutip.coefficent.conj(Coefficient)``.
    """
    cdef Coefficient base

    def __init__(self, Coefficient base):
        self.base = base

    cdef complex _call(self, double t) except *:
        return conj(self.base._call(t))

    cpdef Coefficient copy(self):
        """Return a copy of the :obj:`Coefficient`."""
        return ConjCoefficient(self.base.copy())

    def replace_arguments(self, _args=None, **kwargs):
        """
        Replace the arguments (``args``) of a coefficient.

        Returns a new :obj:`Coefficient` if the coefficient has arguments, or
        the original coefficient if it does not. Arguments to replace may be
        supplied either in a dictionary as the first position argument, or
        passed as keywords, or as a combination of the two. Arguments not
        replaced retain their previous values.

        Parameters
        ----------
        _args : dict
            Dictionary of arguments to replace.

        **kwargs
            Arguments to replace.
        """
        return ConjCoefficient(
            self.base.replace_arguments(_args, **kwargs)
        )


@cython.auto_pickle(True)
cdef class NormCoefficient(Coefficient):
    """
    The L2 :func:`norm` of a :obj:`Coefficient`. A shortcut for
    ``conj(coeff) * coeff``.

    :obj:`NormCoefficient` is returned by
    ``qutip.coefficent.norm(Coefficient)``.
    """
    cdef Coefficient base

    def __init__(self, Coefficient base):
        self.base = base

    def replace_arguments(self, _args=None, **kwargs):
        """
        Replace the arguments (``args``) of a coefficient.

        Returns a new :obj:`Coefficient` if the coefficient has arguments, or
        the original coefficient if it does not. Arguments to replace may be
        supplied either in a dictionary as the first position argument, or
        passed as keywords, or as a combination of the two. Arguments not
        replaced retain their previous values.

        Parameters
        ----------
        _args : dict
            Dictionary of arguments to replace.

        **kwargs
            Arguments to replace.
        """
        return NormCoefficient(
            self.base.replace_arguments(_args, **kwargs)
        )

    cdef complex _call(self, double t) except *:
        return norm(self.base._call(t))

    cpdef Coefficient copy(self):
        """Return a copy of the :obj:`Coefficient`."""
        return NormCoefficient(self.base.copy())
