from collections import defaultdict
from functools import reduce

from .compatibility import is_sequence, default_sort_key
from .logic import _fuzzy_group
from .singleton import S
from .operations import AssocOp
from .cache import cacheit
from .numbers import ilcm, igcd
from .expr import Expr


def _unevaluated_Add(*args):
    """Return a well-formed unevaluated Add: Numbers are collected and
    put in slot 0 and args are sorted. Use this when args have changed
    but you still want to return an unevaluated Add.

    Examples
    ========

    >>> from diofant.core.add import _unevaluated_Add as uAdd
    >>> from diofant import S, Add, Integer, Float
    >>> from diofant.abc import x, y
    >>> a = uAdd(*[Float(1.0), x, Integer(2)])
    >>> a.args[0]
    3.00000000000000
    >>> a.args[1]
    x

    Beyond the Number being in slot 0, there is no other assurance of
    order for the arguments since they are hash sorted. So, for testing
    purposes, output produced by this in some other function can only
    be tested against the output of this function or as one of several
    options:

    >>> opts = (Add(x, y, evaluated=False), Add(y, x, evaluated=False))
    >>> a = uAdd(x, y)
    >>> assert a in opts and a == uAdd(x, y)
    >>> uAdd(x + 1, x + 2)
    x + x + 3
    """
    args = list(args)
    newargs = []
    co = S.Zero
    while args:
        a = args.pop()
        if a.is_Add:
            # this will keep nesting from building up
            # so that x + (x + 1) -> x + x + 1 (3 args)
            args.extend(a.args)
        elif a.is_Number:
            co += a
        else:
            newargs.append(a)
    newargs.sort(key=default_sort_key)
    if co:
        newargs.insert(0, co)
    return Add._from_args(newargs)


class Add(Expr, AssocOp):

    is_Add = True

    @classmethod
    def flatten(cls, seq):
        """
        Takes the sequence "seq" of nested Adds and returns a flatten list.

        Returns: (commutative_part, noncommutative_part, order_symbols)

        Applies associativity, all terms are commutable with respect to
        addition.

        See also
        ========

        diofant.core.mul.Mul.flatten

        """
        from diofant.series.order import Order

        rv = None
        if len(seq) == 2:
            a, b = seq
            if b.is_Rational:
                a, b = b, a
            if a.is_Rational:
                if b.is_Mul:
                    rv = [a, b], [], None
            if rv:
                if all(s.is_commutative for s in rv[0]):
                    return rv
                return [], rv[0], None

        # term -> coeff
        # e.g. x**2 -> 5   for ... + 5*x**2 + ...
        terms = {}

        # coefficient (Number or zoo) to always be in slot 0, e.g. 3 + ...
        coeff = S.Zero

        order_factors = []

        for o in seq:

            # O(x)
            if o.is_Order:
                for o1 in order_factors:
                    if o1.contains(o):
                        o = None
                        break
                if o is None:
                    continue
                order_factors = [o] + [
                    o1 for o1 in order_factors if not o.contains(o1)]
                continue

            # 3 or NaN
            elif o.is_Number:
                if (o is S.NaN or coeff is S.ComplexInfinity and
                        o.is_finite is False):
                    # we know for sure the result will be nan
                    return [S.NaN], [], None
                if coeff.is_Number:
                    coeff += o
                    if coeff is S.NaN:
                        # we know for sure the result will be nan
                        return [S.NaN], [], None
                o  # XXX "peephole" optimization, http://bugs.python.org/issue2506
                continue

            elif o is S.ComplexInfinity:
                if coeff.is_finite is False:
                    # we know for sure the result will be nan
                    return [S.NaN], [], None
                coeff = S.ComplexInfinity
                continue

            # Add([...])
            elif o.is_Add:
                # NB: here we assume Add is always commutative
                seq.extend(o.args)  # TODO zerocopy?
                continue

            elif o.has(Order):
                c, s = S.One, o

            # Mul([...])
            elif o.is_Mul:
                c, s = o.as_coeff_Mul()

            # check for unevaluated Pow, e.g. 2**3 or 2**(-1/2)
            elif o.is_Pow:
                b, e = o.as_base_exp()
                if b.is_Number and (e.is_Integer or
                                   (e.is_Rational and e.is_negative)):
                    seq.append(b**e)
                    continue
                c, s = S.One, o

            else:
                # everything else
                c = S.One
                s = o

            # now we have:
            # o = c*s, where
            #
            # c is a Number
            # s is an expression with number factor extracted
            # let's collect terms with the same s, so e.g.
            # 2*x**2 + 3*x**2  ->  5*x**2
            if s in terms:
                terms[s] += c
                if terms[s] is S.NaN:
                    # we know for sure the result will be nan
                    return [S.NaN], [], None
            else:
                terms[s] = c

        # now let's construct new args:
        # [2*x**2, x**3, 7*x**4, pi, ...]
        newseq = []
        noncommutative = False
        for s, c in terms.items():
            # 0*s
            if c is S.Zero:
                continue
            # 1*s
            elif c is S.One:
                newseq.append(s)
            # c*s
            else:
                if s.is_Mul:
                    # Mul, already keeps its arguments in perfect order.
                    # so we can simply put c in slot0 and go the fast way.
                    cs = s._new_rawargs(*((c,) + s.args))
                    newseq.append(cs)
                elif s.is_Add:
                    # we just re-create the unevaluated Mul
                    newseq.append(Mul(c, s, evaluate=False))
                else:
                    # alternatively we have to call all Mul's machinery (slow)
                    newseq.append(Mul(c, s))

            noncommutative = noncommutative or not s.is_commutative

        # oo, -oo
        if coeff is S.Infinity:
            newseq = [f for f in newseq if not
                      (f.is_nonnegative or f.is_extended_real and f.is_finite)]

        elif coeff is S.NegativeInfinity:
            newseq = [f for f in newseq if not
                      (f.is_nonpositive or f.is_extended_real and f.is_finite)]

        if coeff is S.ComplexInfinity:
            # zoo might be
            #   infinite_real + finite_im
            #   finite_real + infinite_im
            #   infinite_real + infinite_im
            # addition of a finite real or imaginary number won't be able to
            # change the zoo nature; adding an infinite qualtity would result
            # in a NaN condition if it had sign opposite of the infinite
            # portion of zoo, e.g., infinite_real - infinite_real.
            newseq = [c for c in newseq if not (c.is_finite and
                                                (c.is_extended_real is not None
                                                 or c.is_imaginary is not None))]

        # process O(x)
        if order_factors:
            newseq2 = []
            for t in newseq:
                for o in order_factors:
                    # x + O(x) -> O(x)
                    if o.contains(t):
                        t = None
                        break
                # x + O(x**2) -> x + O(x**2)
                if t is not None:
                    newseq2.append(t)
            newseq = newseq2 + order_factors
            # 1 + O(1) -> O(1)
            for o in order_factors:
                if o.contains(coeff):
                    coeff = S.Zero
                    break

        # order args canonically
        newseq.sort(key=default_sort_key)

        # current code expects coeff to be first
        if coeff is not S.Zero:
            newseq.insert(0, coeff)

        # we are done
        if noncommutative:
            return [], newseq, None
        else:
            return newseq, [], None

    @classmethod
    def class_key(cls):
        """Nice order of classes"""
        return 4, 1, cls.__name__

    def as_coefficients_dict(self):
        """Return a dictionary mapping terms to their Rational coefficient.

        Since the dictionary is a defaultdict, inquiries about terms which
        were not present will return a coefficient of 0. If an expression is
        not an Add it is considered to have a single term.

        Examples
        ========

        >>> from diofant.abc import a, x
        >>> (3*x + a*x + 4).as_coefficients_dict()
        {1: 4, x: 3, a*x: 1}
        >>> _[a]
        0
        >>> (3*a*x).as_coefficients_dict()
        {a*x: 3}
        """

        d = defaultdict(list)
        for ai in self.args:
            c, m = ai.as_coeff_Mul()
            d[m].append(c)
        for k, v in d.items():
            if len(v) == 1:
                d[k] = v[0]
            else:
                d[k] = Add(*v)
        di = defaultdict(int)
        di.update(d)
        return di

    @cacheit
    def as_coeff_add(self, *deps):
        """
        Returns a tuple (coeff, args) where self is treated as an Add and coeff
        is the Number term and args is a tuple of all other terms.

        Examples
        ========

        >>> from diofant.abc import x
        >>> (7 + 3*x).as_coeff_add()
        (7, (3*x,))
        >>> (7*x).as_coeff_add()
        (0, (7*x,))
        """
        if deps:
            l1 = []
            l2 = []
            for f in self.args:
                if f.has(*deps):
                    l2.append(f)
                else:
                    l1.append(f)
            return self._new_rawargs(*l1), tuple(l2)
        coeff, notrat = self.args[0].as_coeff_add()
        if coeff is not S.Zero:
            return coeff, notrat + self.args[1:]
        return S.Zero, self.args

    def as_coeff_Add(self):
        """Efficiently extract the coefficient of a summation. """
        coeff, args = self.args[0], self.args[1:]

        if coeff.is_Number:
            if len(args) == 1:
                return coeff, args[0]
            else:
                return coeff, self._new_rawargs(*args)
        else:
            return S.Zero, self

    # Note, we intentionally do not implement Add.as_coeff_mul().  Rather, we
    # let Expr.as_coeff_mul() just always return (S.One, self) for an Add.  See
    # issue 5524.

    @cacheit
    def _eval_derivative(self, s):
        return self.func(*[a.diff(s) for a in self.args])

    def _eval_nseries(self, x, n, logx):
        terms = [t.nseries(x, n=n, logx=logx) for t in self.args]
        return self.func(*terms)

    def _matches_simple(self, expr, repl_dict):
        # handle (w+3).matches('x+5') -> {w: x+2}
        coeff, terms = self.as_coeff_add()
        if len(terms) == 1:
            return terms[0].matches(expr - coeff, repl_dict)
        return

    def matches(self, expr, repl_dict={}):
        return AssocOp._matches_commutative(self, expr, repl_dict)

    @staticmethod
    def _combine_inverse(lhs, rhs):
        """
        Returns lhs - rhs, but treats arguments like symbols, so things like
        oo - oo return 0, instead of a nan.
        """
        from diofant import oo, I, expand_mul
        if lhs == oo and rhs == oo or lhs == oo*I and rhs == oo*I:
            return S.Zero
        return expand_mul(lhs - rhs)

    @cacheit
    def as_two_terms(self):
        """Return head and tail of self.

        This is the most efficient way to get the head and tail of an
        expression.

        - if you want only the head, use self.args[0];
        - if you want to process the arguments of the tail then use
          self.as_coef_add() which gives the head and a tuple containing
          the arguments of the tail when treated as an Add.
        - if you want the coefficient when self is treated as a Mul
          then use self.as_coeff_mul()[0]

        >>> from diofant.abc import x, y
        >>> (3*x*y).as_two_terms()
        (3, x*y)
        """
        return self.args[0], self._new_rawargs(*self.args[1:])

    def as_numer_denom(self):

        # clear rational denominator
        content, expr = self.primitive()
        ncon, dcon = content.as_numer_denom()

        # collect numerators and denominators of the terms
        nd = defaultdict(list)
        for f in expr.args:
            ni, di = f.as_numer_denom()
            nd[di].append(ni)

        # check for quick exit
        if len(nd) == 1:
            d, n = nd.popitem()
            return self.func(
                *[_keep_coeff(ncon, ni) for ni in n]), _keep_coeff(dcon, d)

        # sum up the terms having a common denominator
        for d, n in nd.items():
            if len(n) == 1:
                nd[d] = n[0]
            else:
                nd[d] = self.func(*n)

        # assemble single numerator and denominator
        denoms, numers = [list(i) for i in zip(*iter(nd.items()))]
        n, d = self.func(*[Mul(*(denoms[:i] + [numers[i]] + denoms[i + 1:]))
                   for i in range(len(numers))]), Mul(*denoms)

        return _keep_coeff(ncon, n), _keep_coeff(dcon, d)

    def _eval_is_polynomial(self, syms):
        return all(term._eval_is_polynomial(syms) for term in self.args)

    def _eval_is_rational_function(self, syms):
        return all(term._eval_is_rational_function(syms) for term in self.args)

    def _eval_is_algebraic_expr(self, syms):
        return all(term._eval_is_algebraic_expr(syms) for term in self.args)

    # assumption methods

    def _eval_is_extended_real(self):
        return _fuzzy_group((a.is_extended_real for a in self.args),
                            quick_exit=True)

    def _eval_is_complex(self):
        return _fuzzy_group((a.is_complex for a in self.args), quick_exit=True)

    def _eval_is_antihermitian(self):
        return _fuzzy_group((a.is_antihermitian for a in self.args),
                            quick_exit=True)

    def _eval_is_finite(self):
        return _fuzzy_group((a.is_finite for a in self.args), quick_exit=True)

    def _eval_is_hermitian(self):
        return _fuzzy_group((a.is_hermitian for a in self.args), quick_exit=True)

    def _eval_is_integer(self):
        return _fuzzy_group((a.is_integer for a in self.args), quick_exit=True)

    def _eval_is_rational(self):
        return _fuzzy_group((a.is_rational for a in self.args), quick_exit=True)

    def _eval_is_algebraic(self):
        return _fuzzy_group((a.is_algebraic for a in self.args), quick_exit=True)

    def _eval_is_imaginary(self):
        rv = _fuzzy_group((a.is_imaginary for a in self.args), quick_exit=True)
        if rv is False:
            return rv
        iargs = [a*S.ImaginaryUnit for a in self.args]
        return _fuzzy_group((a.is_real for a in iargs), quick_exit=True)

    def _eval_is_odd(self):
        l = [f for f in self.args if not f.is_even]
        if not l:
            return False
        if l[0].is_odd:
            return self._new_rawargs(*l[1:]).is_even

    def _eval_is_irrational(self):
        for t in self.args:
            a = t.is_irrational
            if a:
                if all(x.is_rational for x in self.args if x != t):
                    return True
                return
            if a is None:
                return
        return False

    def _eval_is_positive(self):
        if self.is_number:
            return super(Add, self)._eval_is_positive()

        if any(a.is_infinite for a in self.args):
            args = [a for a in self.args if not a.is_finite]
        else:
            args = self.args

        nonpos = nonneg = 0
        for a in args:
            if a.is_positive:
                continue
            elif a.is_nonnegative:
                nonneg += 1
                if a.is_zero:
                    nonpos += 1
            elif a.is_nonpositive:
                nonpos += 1
            else:
                break
        else:
            if not nonpos and nonneg < len(args):
                return True
            elif nonpos == len(args):
                return False

    def _eval_is_negative(self):
        if self.is_number:
            return super(Add, self)._eval_is_negative()

        if any(a.is_infinite for a in self.args):
            args = [a for a in self.args if not a.is_finite]
        else:
            args = self.args

        nonneg = nonpos = 0
        for a in args:
            if a.is_negative:
                continue
            elif a.is_nonpositive:
                nonpos += 1
                if a.is_zero:
                    nonneg += 1
            elif a.is_nonnegative:
                nonneg += 1
            else:
                break
        else:
            if not nonneg and nonpos < len(args):
                return True
            elif nonneg == len(args):
                return False

    def _eval_subs(self, old, new):
        if not old.is_Add:
            return

        coeff_self, terms_self = self.as_coeff_Add()
        coeff_old, terms_old = old.as_coeff_Add()

        if coeff_self.is_Rational and coeff_old.is_Rational:
            if terms_self == terms_old:   # (2 + a).subs( 3 + a, y) -> -1 + y
                return self.func(new, coeff_self, -coeff_old)
            if terms_self == -terms_old:  # (2 + a).subs(-3 - a, y) -> -1 - y
                return self.func(-new, coeff_self, coeff_old)

        if coeff_self.is_Rational and coeff_old.is_Rational \
                or coeff_self == coeff_old:
            args_old, args_self = self.func.make_args(
                terms_old), self.func.make_args(terms_self)
            if len(args_old) < len(args_self):  # (a+b+c).subs(b+c,x) -> a+x
                self_set = set(args_self)
                old_set = set(args_old)

                if old_set < self_set:
                    ret_set = self_set - old_set
                    return self.func(new, coeff_self, -coeff_old,
                               *[s._subs(old, new) for s in ret_set])

                args_old = self.func.make_args(
                    -terms_old)     # (a+b+c+d).subs(-b-c,x) -> a-x+d
                old_set = set(args_old)
                if old_set < self_set:
                    ret_set = self_set - old_set
                    return self.func(-new, coeff_self, coeff_old,
                               *[s._subs(old, new) for s in ret_set])

    def removeO(self):
        args = [a for a in self.args if not a.is_Order]
        return self._new_rawargs(*args)

    def getO(self):
        args = [a for a in self.args if a.is_Order]
        if args:
            return self._new_rawargs(*args)

    @cacheit
    def extract_leading_order(self, symbols):
        """Returns the leading term and its order.

        Examples
        ========

        >>> from diofant.abc import x
        >>> (x + 1 + 1/x**5).extract_leading_order(x)
        ((x**(-5), O(x**(-5))),)
        >>> (1 + x).extract_leading_order(x)
        ((1, O(1)),)
        >>> (x + x**2).extract_leading_order(x)
        ((x, O(x)),)
        """
        from diofant import Order
        lst = []
        symbols = list(symbols if is_sequence(symbols) else [symbols])
        point = [0]*len(symbols)
        seq = [(f, Order(f, *zip(symbols, point))) for f in self.args]
        for ef, of in seq:
            for e, o in lst:
                if o.contains(of) and o != of:
                    of = None
                    break
            if of is None:
                continue
            new_lst = [(ef, of)]
            for e, o in lst:
                if of.contains(o) and o != of:
                    continue
                new_lst.append((e, o))
            lst = new_lst
        return tuple(lst)

    def as_real_imag(self, deep=True, **hints):
        """
        returns a tuple representing a complex number

        Examples
        ========

        >>> from diofant import I
        >>> (7 + 9*I).as_real_imag()
        (7, 9)
        >>> ((1 + I)/(1 - I)).as_real_imag()
        (0, 1)
        >>> ((1 + 2*I)*(1 + 3*I)).as_real_imag()
        (-5, 5)
        """
        sargs, terms = self.args, []
        re_part, im_part = [], []
        for term in sargs:
            re, im = term.as_real_imag(deep=deep)
            re_part.append(re)
            im_part.append(im)
        return self.func(*re_part), self.func(*im_part)

    def _eval_as_leading_term(self, x):
        from diofant import factor_terms

        expr = self.func(*[t.as_leading_term(x) for t in self.args]).removeO()
        if not expr:
            # simple leading term analysis gave us 0 but we have to send
            # back a term, so compute the leading term (via series)
            return self.compute_leading_term(x)
        elif not expr.is_Add:
            return expr
        else:
            plain = expr.func(*[s for s, _ in expr.extract_leading_order(x)])
            rv = factor_terms(plain, fraction=False)
            rv_simplify = rv.simplify()
            # if it simplifies to an x-free expression, return that;
            # tests don't fail if we don't but it seems nicer to do this
            if x not in rv_simplify.free_symbols:
                if rv_simplify.is_zero and plain.is_zero is not True:
                    return (expr - plain)._eval_as_leading_term(x)
                return rv_simplify
            return rv

    def _eval_adjoint(self):
        return self.func(*[t.adjoint() for t in self.args])

    def _eval_conjugate(self):
        return self.func(*[t.conjugate() for t in self.args])

    def _eval_transpose(self):
        return self.func(*[t.transpose() for t in self.args])

    def __neg__(self):
        return self.func(*[-t for t in self.args])

    def primitive(self):
        """
        Return ``(R, self/R)`` where ``R``` is the Rational GCD of ``self```.

        ``R`` is collected only from the leading coefficient of each term.

        Examples
        ========

        >>> from diofant.abc import x, y

        >>> (2*x + 4*y).primitive()
        (2, x + 2*y)

        >>> (2*x/3 + 4*y/9).primitive()
        (2/9, 3*x + 2*y)

        >>> (2*x/3 + 4.2*y).primitive()
        (1/3, 2*x + 12.6*y)

        No subprocessing of term factors is performed:

        >>> ((2 + 2*x)*x + 2).primitive()
        (1, x*(2*x + 2) + 2)

        Recursive subprocessing can be done with the as_content_primitive()
        method:

        >>> ((2 + 2*x)*x + 2).as_content_primitive()
        (2, x*(x + 1) + 1)

        See Also
        ========

        diofant.polys.polytools.primitive
        """

        terms = []
        inf = False
        for a in self.args:
            c, m = a.as_coeff_Mul()
            if not c.is_Rational:
                c = S.One
                m = a
            inf = inf or m is S.ComplexInfinity
            terms.append((c.p, c.q, m))

        if not inf:
            ngcd = reduce(igcd, [t[0] for t in terms], 0)
            dlcm = reduce(ilcm, [t[1] for t in terms], 1)
        else:
            ngcd = reduce(igcd, [t[0] for t in terms if t[1]], 0)
            dlcm = reduce(ilcm, [t[1] for t in terms if t[1]], 1)

        if ngcd == dlcm == 1:
            return S.One, self
        for i, (p, q, term) in enumerate(terms):
            terms[i] = _keep_coeff(Rational((p//ngcd)*(dlcm//q)), term)

        # we don't need a complete re-flattening since no new terms will join
        # so we just use the same sort as is used in Add.flatten. When the
        # coefficient changes, the ordering of terms may change, e.g.
        #     (3*x, 6*y) -> (2*y, x)
        #
        # We do need to make sure that term[0] stays in position 0, however.
        #
        if terms[0].is_Number or terms[0] is S.ComplexInfinity:
            c = terms.pop(0)
        else:
            c = None
        terms.sort(key=default_sort_key)
        if c:
            terms.insert(0, c)
        return Rational(ngcd, dlcm), self._new_rawargs(*terms)

    def as_content_primitive(self, radical=False):
        """Return the tuple (R, self/R) where R is the positive Rational
        extracted from self. If radical is True (default is False) then
        common radicals will be removed and included as a factor of the
        primitive expression.

        Examples
        ========

        >>> from diofant import sqrt
        >>> (3 + 3*sqrt(2)).as_content_primitive()
        (3, 1 + sqrt(2))

        Radical content can also be factored out of the primitive:

        >>> (2*sqrt(2) + 4*sqrt(10)).as_content_primitive(radical=True)
        (2, sqrt(2)*(1 + 2*sqrt(5)))

        See Also
        ========

        diofant.core.expr.Expr.as_content_primitive
        """
        con, prim = self.func(*[_keep_coeff(*a.as_content_primitive(
            radical=radical)) for a in self.args]).primitive()
        if radical and prim.is_Add:
            # look for common radicals that can be removed
            args = prim.args
            rads = []
            common_q = None
            for m in args:
                term_rads = defaultdict(list)
                for ai in Mul.make_args(m):
                    if ai.is_Pow:
                        b, e = ai.as_base_exp()
                        if e.is_Rational and b.is_Integer:
                            term_rads[e.q].append(abs(int(b))**e.p)
                if not term_rads:
                    break
                if common_q is None:
                    common_q = set(term_rads.keys())
                else:
                    common_q = common_q & set(term_rads.keys())
                    if not common_q:
                        break
                rads.append(term_rads)
            else:
                # process rads
                # keep only those in common_q
                for r in rads:
                    for q in list(r.keys()):
                        if q not in common_q:
                            r.pop(q)
                    for q in r:
                        r[q] = prod(r[q])
                # find the gcd of bases for each q
                G = []
                for q in common_q:
                    g = reduce(igcd, [r[q] for r in rads], 0)
                    if g != 1:
                        G.append(g**Rational(1, q))
                if G:
                    G = Mul(*G)
                    args = [ai/G for ai in args]
                    prim = G*prim.func(*args)

        return con, prim

    @property
    def _sorted_args(self):
        return tuple(sorted(self.args, key=default_sort_key))

from .mul import Mul, _keep_coeff, prod
from .numbers import Rational
