# Copyright 2011-2017 Kwant authors.
#
# This file is a modified version of a file that is part of Kwant.
# It is subject to the license terms below. A list of Kwant authors
# can be found at http://kwant-project.org/authors.
#
# =============
# Kwant license
# =============
#
# Copyright 2011-2015 C. W. Groth (CEA), M. Wimmer, A. R. Akhmerov,
# X. Waintal (CEA), and others.  All rights reserved.
#
# (CEA = Commissariat à l'énergie atomique et aux énergies alternatives)
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * 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.
#
# 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 keyword
from collections import defaultdict

import numpy as np

import sympy
import sympy.abc
import sympy.physics.quantum
from sympy.core.function import AppliedUndef
from sympy.core.sympify import converter
from sympy.core.core import all_classes as sympy_classes
from sympy.physics.matrices import msigma as _msigma

import warnings

# TODO: remove when sympy correctly includes MutableDenseMatrix (lol).
sympy_classes = set(sympy_classes) | {sympy.MutableDenseMatrix}

momentum_operators = sympy.symbols('k_x k_y k_z', commutative=False)
position_operators = sympy.symbols('x y z', commutative=False)

pauli = [sympy.eye(2), _msigma(1), _msigma(2), _msigma(3)]

extra_ns = sympy.abc._clash.copy()
extra_ns.update({s.name: s for s in momentum_operators})
extra_ns.update({s.name: s for s in position_operators})
extra_ns.update({'kron': sympy.physics.quantum.TensorProduct,
                 'eye': sympy.eye, 'identity': sympy.eye})
extra_ns.update({'sigma_{}'.format(c): p for c, p in zip('0xyz', pauli)})


# workaroud for https://github.com/sympy/sympy/issues/12060
del extra_ns['I']
del extra_ns['pi']


################  Helpers to handle sympy

def lambdify(expr, locals=None):
    """Return a callable object for computing a continuum Hamiltonian.

    .. warning::
        This function uses ``eval`` (because it calls ``sympy.sympify``), and
        thus should not be used on unsanitized input.

    If necessary, the given expression is sympified using
    `kwant.continuum.sympify`.  It is then converted into a callable object.

    Parameters
    ----------
    expr : str or SymPy expression
        Expression to be converted into a callable object
    locals : dict or ``None`` (default)
        Additional definitions for `~kwant.continuum.sympify`.

    Examples
    --------
    >>> f = lambdify('a + b', locals={'b': 'b + c'})
    >>> f(1, 3, 5)
    9

    >>> ns = {'sigma_plus': [[0, 2], [0, 0]]}
    >>> f = lambdify('k_x**2 * sigma_plus', ns)
    >>> f(0.25)
    array([[ 0.   ,  0.125],
           [ 0.   ,  0.   ]])
    """
    expr = sympify(expr, locals)

    args = [s.name for s in expr.atoms(sympy.Symbol)]
    args += [str(f.func) for f in expr.atoms(AppliedUndef, sympy.Function)]

    return sympy.lambdify(sorted(args), expr)


def sympify(expr, locals=None):
    """Sympify object using special rules for Hamiltonians.

    If `'expr`` is already a type that SymPy understands, it will do nothing
    but return that value. Note that ``locals`` will not be used in this
    situation.

    Otherwise, it is sympified by ``sympy.sympify`` with a modified namespace
    such that

    * the position operators "x", "y" or "z" and momentum operators "k_x",
      "k_y", and "k_z" do not commute,
    * all single-letter identifiers and names of Greek letters (e.g. "pi" or
      "gamma") are treated as symbols,
    * "kron" corresponds to ``sympy.physics.quantum.TensorProduct``, and
      "identity" to ``sympy.eye``,
    * "sigma_0", "sigma_x", "sigma_y", "sigma_z" are the Pauli matrices.

    In addition, Python list literals are interpreted as SymPy matrices.

    .. warning::
        This function uses ``eval`` (because it calls ``sympy.sympify``), and
        thus should not be used on unsanitized input.

    Parameters
    ----------
    expr : str or SymPy expression
        Expression to be converted to a SymPy object.
    locals : dict or ``None`` (default)
        Additional entries for the namespace under which `expr` is sympified.
        The keys must be valid Python variable names.  The values may be
        strings, since they are all are sent through `continuum.sympify`
        themselves before use.  (Note that this is a difference to how
        ``sympy.sympify`` behaves.)

        .. note::
            When a value of `locals` is already a SymPy object, it is used
            as-is, and the caller is responsible to set the commutativity of
            its symbols appropriately.  This possible source of errors is
            demonstrated in the last example below.

    Returns
    -------
    result : SymPy object

    Examples
    --------
        >>> sympify('k_x * A(x) * k_x + V(x)')
        k_x*A(x)*k_x + V(x)     # as valid sympy object

        >>> sympify('k_x**2 + V', locals={'V': 'V_0 + V(x)'})
        k_x**2 + V(x) + V_0

        >>> ns = {'sigma_plus': [[0, 2], [0, 0]]}
        >>> sympify('k_x**2 * sigma_plus', ns)
        Matrix([
        [0, 2*k_x**2],
        [0,        0]])

        >>> sympify('k_x * A(c) * k_x', locals={'c': 'x'})
        k_x*A(x)*k_x
        >>> sympify('k_x * A(c) * k_x', locals={'c': sympy.Symbol('x')})
        A(x)*k_x**2

    """
    stored_value = None

    # if ``expr`` is already a ``sympy`` object we may terminate a code path
    if isinstance(expr, tuple(sympy_classes)):
        if locals:
            warnings.warn('Input expression is already SymPy object: '
                          '"locals" will not be used.',
                          RuntimeWarning,
                          stacklevel=2)

        # We assume that all present functions, like "sin", "cos", will be
        # provided by user during the final evaluation through "params".
        # Therefore we make sure they are defined as AppliedUndef, not built-in
        # sympy types.
        subs = {r: sympy.Function(str(r.func))(*r.args)
                for r in expr.atoms(sympy.Function)}

        return expr.subs(subs)

    # if ``expr`` is not a "sympy" then we proceed with sympifying process
    if locals is None:
        locals = {}

    for k in locals:
        if (not isinstance(k, str)
            or not k.isidentifier() or keyword.iskeyword(k)):
            raise ValueError(
                "Invalid key in 'locals': {}\nKeys must be "
                "identifiers and may not be keywords".format(repr(k)))

    # sympify values of locals before updating it with extra_ns
    locals = {k: sympify(v) for k, v in locals.items()}
    for k, v in extra_ns.items():
        locals.setdefault(k, v)
    try:
        stored_value = converter.pop(list, None)
        converter[list] = lambda x: sympy.Matrix(x)
        hamiltonian = sympy.sympify(expr, locals=locals)
        # if input is for example ``[[k_x * A(x) * k_x]]`` after the first
        # sympify we are getting list of sympy objects, so we call sympify
        # second time to obtain ``sympy`` matrices.
        hamiltonian = sympy.sympify(hamiltonian)
    finally:
        if stored_value is not None:
            converter[list] = stored_value
        else:
            del converter[list]

    return hamiltonian


def make_commutative(expr, *symbols):
    """Make sure that specified symbols are defined as commutative.

    Parameters
    ----------
    expr: sympy.Expr or sympy.Matrix
    symbols: sequace of symbols
        Set of symbols that are requiered to be commutative. It doesn't matter
        of symbol is provided as commutative or not.

    Returns
    -------
    input expression with all specified symbols changed to commutative.
    """
    symbols = [sympy.Symbol(s.name, commutative=False) for s in symbols]
    expr = expr.subs({s: sympy.Symbol(s.name) for s in symbols})
    return expr


def monomials(expr, gens=None):
    """Parse ``expr`` into monomials in the symbols in ``gens``.

    Parameters
    ----------
    expr: sympy.Expr or sympy.Matrix
        Sympy expression to be parsed into monomials.
    gens: sequence of sympy.Symbol objects or strings (optional)
        Generators of monomials. If unset it will default to all
        symbols used in ``expr``.

    Returns
    -------
    dictionary (generator: monomial)

    Examples
    --------
        >>> expr = kwant.continuum.sympify("A * (x**2 + y) + B * x + C")
        >>> monomials(expr, gens=('x', 'y'))
        {1: C, x: B, x**2: A, y: A}
    """
    if gens is None:
        gens = expr.atoms(sympy.Symbol)
    else:
        gens = [sympify(g) for g in gens]

    if not isinstance(expr, sympy.MatrixBase):
        return _expression_monomials(expr, gens)
    else:
        output = defaultdict(lambda: sympy.zeros(*expr.shape))
        for (i, j), e in np.ndenumerate(expr):
            mons = _expression_monomials(e, gens)
            for key, val in mons.items():
                output[key][i, j] += val
        return dict(output)


def _expression_monomials(expr, gens):
    """Parse ``expr`` into monomials in the symbols in ``gens``.

    Parameters
    ----------
    expr: sympy.Expr
        Sympy expr to be parsed.
    gens: sequence of sympy.Symbol
        Generators of monomials.

    Returns
    -------
    dictionary (generator: monomial)
    """
    expr = sympy.expand(expr)
    output = defaultdict(lambda: sympy.Integer(0))
    for summand in expr.as_ordered_terms():
        key = []
        val = []
        for factor in summand.as_ordered_factors():
            symbol, exponent = factor.as_base_exp()
            if symbol in gens:
                key.append(factor)
            else:
                val.append(factor)
        output[sympy.Mul(*key)] += sympy.Mul(*val)

    return dict(output)


################ general help functions

def gcd(*args):
    if len(args) == 1:
        return args[0]

    L = list(args)

    while len(L) > 1:
        a = L[len(L) - 2]
        b = L[len(L) - 1]
        L = L[:len(L) - 2]

        while a:
            a, b = b%a, a

        L.append(b)

    return abs(b)
