"This file implements a class to represent a sum."

__author__ = "Kristian B. Oelgaard (k.b.oelgaard@tudelft.nl)"
__date__ = "2009-07-12 -- 2009-08-08"
__copyright__ = "Copyright (C) 2009 Kristian B. Oelgaard"
__license__  = "GNU GPL version 3 or any later version"

# FFC common modules.
#from ffc.common.log import error

from new_symbol import create_float, create_product, create_sum, create_fraction
from expr import Expr

# FFC common modules
from ffc.common.log import error

#import psyco
#psyco.full()

# TODO: This function is needed to avoid passing around the 'format', but could
# it be done differently?
def set_format(_format):
    global format
    format = _format

class Sum(Expr):
    __slots__ = ("vrs", "_expanded", "_reduced")
    def __init__(self, variables):
        """Initialise a Sum object, it derives from Expr and contains the
        additional variables:

        vrs       - list, a list of variables.
        _expanded - object, an expanded object of self, e.g.,
                    self = 'x + x'-> self._expanded = 2*x (a product).
        _reduced  - object, a reduced object of self, e.g.,
                    self = '2*x + x*y'-> self._reduced = x*(2 + y) (a product).
        NOTE: self._prec = 3."""

        # Initialise value, list of variables, class, expanded and reduced.
        self.val = 1.0
        self.vrs = []
        self._prec = 3
        self._expanded = False
        self._reduced = False

        # Process variables if we have any.
        if variables:
            # Remove nested Sums.
            new_vars = []
            for v in variables:
                if v._prec == 3: # sum
                    new_vars += v.vrs
                    continue
                new_vars.append(v)

            floats = []
            # Loop variables and sort in positive and negative.
            # Also collect all floats in 1 variable.
            # We don't collect [x, x, x] into 3*x to avoid creating objects.
            # Instead we do this when expanding the object.
            for v in new_vars:
                # Skip zero terms.
                if v.val ==  0.0:
                    continue
                if v._prec == 0: # float
                    floats.append(v)
                else:
                    self.vrs.append(v)

            # Only create new float if we have more than one. Also ignore it
            # if 1 + 1 - 2 = 0.
            if len(floats) > 1:
                val = sum([v.val for v in floats])
                if val:
                    self.vrs.append(create_float(val))
            elif floats:
                self.vrs.append(floats[0])

        # If we don't have any variables the sum is zero.
        else:
            self.val = 0.0
            self.vrs = [create_float(0)]

        # Handle zero value.
        if not self.vrs:
            self.val = 0.0
            self.vrs = [create_float(0)]

        # Type is equal to the smallest type in both lists.
        self.t = min([v.t for v in self.vrs])


        # Sort variables, (for representation).
        self.vrs.sort()

        # Compute the representation now, such that we can use it directly
        # in the __eq__ and __ne__ methods (improves performance a bit, but
        # only when objects are cached).
        self._repr = "Sum([%s])" % ", ".join([v._repr for v in self.vrs])

        # Use repr as hash value.
        self._hash = hash(self._repr)

    # Print functions.
    def __str__(self):
        "Simple string representation which will appear in the generated code."
        # First add all the positive variables using plus, then add all
        # negative variables.
        s = format["add"]([str(v) for v in self.vrs if not v.val < 0]) + \
            "".join([str(v) for v in self.vrs if v.val < 0])
        # Group only if we have more that one variable.
        if len(self.vrs) > 1:
            return format["grouping"](s)
        return s

    # Binary operators.
    def __mul__(self, other):
        "Multiplication by other objects."
        # If product will be zero.
        if self.val == 0.0 or other.val == 0.0:
            return create_float(0)

        # NOTE: We expect expanded sub-expressions with no nested operators.
        # Create list of new products using the '*' operator
        # TODO: Is this efficient?
        new_prods = [v*other for v in self.vrs]

        # Remove zero valued terms.
        # TODO: Can this still happen?
        new_prods = [v for v in new_prods if v.val != 0.0]

        # Create new sum.
        if not new_prods:
            return create_float(0)
        elif len(new_prods) > 1:
            # Expand sum to collect terms.
            return create_sum(new_prods).expand()
        # TODO: Is it necessary to call expand?
        return new_prods[0].expand()

    def __div__(self, other):
        "Division by other objects."
        # If division is illegal (this should definitely not happen).
        if other.val == 0.0:
            error("Division by zero.")

        # If fraction will be zero.
        if self.val == 0.0:
            return create_float(0)

        # NOTE: assuming that we get expanded variables.
        # If other is a Sum we can only return a fraction.
        # TODO: We could check for equal sums if Sum.__eq__ could be trusted.
        # As it is now (2*x + y) == (3*x + y), which works for the other things I do.
        # NOTE: Expect that other is expanded i.e., x + x -> 2*x which can be handled.
        # TODO: Fix (1 + y) / (x + x*y) -> 1 / x
        # Will this be handled when reducing operations on a fraction?
        if other._prec == 3: # sum
            return create_fraction(self, other)

        # NOTE: We expect expanded sub-expressions with no nested operators.
        # Create list of new products using the '*' operator.
        # TODO: Is this efficient?
        new_fracs = [v/other for v in self.vrs]

        # Remove zero valued terms.
        # TODO: Can this still happen?
        new_fracs = [v for v in new_fracs if v.val != 0.0]

        # Create new sum.
        # TODO: No need to call expand here, using the '/' operator should have
        # taken care of this.
        if not new_fracs:
            return create_float(0)
        elif len(new_fracs) > 1:
            return create_sum(new_fracs)
        return new_fracs[0]

    # Public functions.
    def expand(self):
        "Expand all members of the sum."

        # If sum is already expanded, simply return the expansion.
        if self._expanded:
            return self._expanded

        # TODO: This function might need some optimisation.

        # Loop all variable and expand, then remove nested sums that got
        # created by expanding variables.
        variables = []
        for v in self.vrs:
            exp = v.expand()
            if exp._prec == 3: # sum
                variables += exp.vrs
                continue
            variables.append(exp)

        # Sort variables into symbols, products and fractions (add floats
        # directly to new list, will be handled later). Add fractions if
        # possible else add to list.
        new_variables = []
        syms = []
        prods = []
        frac_groups = {}
        # TODO: Rather than using '+', would it be more efficient to collect
        # the terms first?
        for v in variables:
            # TODO: Should we also group fractions, or put this in a separate function?
            if v._prec in (0, 4): # float or frac
                new_variables.append(v)
            elif v._prec == 1: # sym
                syms.append(v)
            elif v._prec == 2: # prod
                prods.append(v)
            # TODO: put this in another function, cannot group fractions
            # before we have reduced fractions with respect to Types etc.
#            else:
#                # Try to group fractions with similar denominators.
#                if v.denom in frac_groups:
#                    frac_groups[v.denom] += v
#                else:
#                    frac_groups[v.denom] = v

        # Sort all variables in groups: [2*x, -7*x], [(x + y), (2*x + 4*y)] etc.
        # First handle product in order to add symbols if possible.
        prod_groups = {}
        for v in prods:
            if v.get_vrs() in prod_groups:
                prod_groups[v.get_vrs()] += v
            else:
                prod_groups[v.get_vrs()] = v

        sym_groups = {}
        # Loop symbols and add to appropriate groups.
        for v in syms:
            # First try to add to a product group.
            if (v,) in prod_groups:
                prod_groups[(v,)] = v
            # Then to other symbols.
            elif v in sym_groups:
                sym_groups[v] += v
            # Create a new entry in the symbols group.
            else:
                sym_groups[v] = v

        # Loop groups and add to new variable list.
        for k,v in sym_groups.iteritems():
            new_variables.append(v)
        for k,v in prod_groups.iteritems():
            new_variables.append(v)
#        for k,v in frac_groups.iteritems():
#            new_variables.append(v)
#            append(v)

        if len(new_variables) > 1:
            # Return new sum (will remove multiple instances of floats during construction).
            self._expanded = create_sum(new_variables)
            return self._expanded
        elif new_variables:
            # If we just have one variable left, return it since it is already expanded.
            self._expanded = new_variables[0]
            return self._expanded
        error("Where did the variables go?")

    def get_unique_vars(self, var_type):
        "Get unique variables (Symbols) as a set."
        # Loop all variables of self update the set.
        var = set()
        for v in self.vrs:
            var.update(v.get_unique_vars(var_type))
        return var

    def get_var_occurrences(self):
        """Determine the number of minimum number of times all variables occurs
        in the expression. Returns a dictionary of variables and the number of
        times they occur. x*x + x returns {x:1}, x + y returns {}."""
        # NOTE: This function is only used if the numerator of a Fraction is a Sum.

        # Get occurrences in first expression.
        d0 = self.vrs[0].get_var_occurrences()
        for var in self.vrs[1:]:
            # Get the occurrences.
            d = var.get_var_occurrences()
            # Delete those variables in d0 that are not in d.
            for k, v in d0.items():
                if not k in d:
                    del d0[k]
            # Set the number of occurrences equal to the smallest number.
            for k, v in d.iteritems():
                if k in d0:
                    d0[k] = min(d0[k], v)
        return d0

    def ops(self):
        "Return number of operations to compute value of sum."
        # Subtract one operation as it only takes n-1 ops to sum n members.
        op = -1

        # Add the number of operations from sub-expressions.
        for v in self.vrs:
            #  +1 for the +/- symbol.
            op += v.ops() + 1
        return op

    def reduce_ops(self):
        "Reduce the number of operations needed to evaluate the sum."

        if self._reduced:
            return self._reduced
        # NOTE: Assuming that sum has already been expanded.
        # TODO: Add test for this and handle case if it is not.

        # TODO: The entire function looks expensive, can it be optimised?

        # TODO: It is not necessary to create a new Sum if we do not have more
        # than one Fraction.
        # First group all fractions in the sum.
        new_sum = group_fractions(self)
        if new_sum._prec != 3: # sum
            self._reduced = new_sum.reduce_ops()
            return self._reduced

        # Loop all variables of the sum and collect the number of common
        # variables that can be factored out.
        common_vars = {}
        for var in new_sum.vrs:

            # Get dictonary of occurrences and add the variable and the number
            # of occurrences to common dictionary.
            for k, v in var.get_var_occurrences().iteritems():
                if k in common_vars:
                    common_vars[k].append((v, var))
                else:
                    common_vars[k] = [(v, var)]

        # Determine the maximum reduction for each variable
        # sorted as: {(x*x*y, x*y*z, 2*y):[2, [y]]}.
        terms_reductions = {}
        for k, v in common_vars.iteritems():
            # If the number of expressions that can be reduced is only one
            # there is nothing to be done.
            if len(v) > 1:
                # TODO: Is there a better way to compute the reduction gain
                # and the number of occurrences we should remove?

                # Get the list of number of occurences of 'k' in expressions
                # in 'v'.
                occurrences = [t[0] for t in v]
                # Extract the terms of v.
                terms = [t[1] for t in v]

                # Determine the favorable number of occurences and an estimate
                # of the maximum reduction for current variable.
                fav_occur = 0
                reduc = 0
                for i in set(occurrences):
                    # Get number of terms that has a number of occcurences equal
                    # to or higher than the current number.
                    num_terms = len([o for o in occurrences if o >= i])

                    # An estimate of the reduction in operations is:
                    # (number_of_terms - 1) * number_occurrences.
                    new_reduc = (num_terms-1)*i
                    if new_reduc > reduc:
                        reduc = new_reduc
                        fav_occur = i

                # We need to reduce the expression with the favorable number of
                # occurrences of the current variable.
                red_vars = [k]*fav_occur

                # If the list of terms is already present in the dictionary,
                # add the reduction count and the variables.
                if tuple(terms) in terms_reductions:
                    terms_reductions[tuple(terms)][0] += reduc
                    terms_reductions[tuple(terms)][1] += red_vars
                else:
                    terms_reductions[tuple(terms)] = [reduc, red_vars]

        if terms_reductions:
            # Invert dictionary of terms.
            reductions_terms = dict([((v[0], tuple(v[1])), k) for k, v in terms_reductions.iteritems()])

            # Create a sorted list of those variables that give the highest
            # reduction.
            sorted_reduc_var = [k for k, v in reductions_terms.iteritems()]
            sorted_reduc_var.sort(lambda x, y: cmp(x[0], y[0]))
            sorted_reduc_var.reverse()

            # Create a new dictionary of terms that should be reduced, if some
            # terms overlap, only pick the one which give the highest reduction to
            # ensure that a*x*x + b*x*x + x*x*y + 2*y -> x*x*(a + b + y) + 2*y NOT 
            # x*x*(a + b) + y*(2 + x*x).
            reduction_vars = {}
            rejections = {}
            for var in sorted_reduc_var:
                terms = reductions_terms[var]
                if overlap(terms, reduction_vars) or overlap(terms, rejections):
                    rejections[var[1]] = terms
                else:
                    reduction_vars[var[1]] = terms

            # Reduce each set of terms with appropriate variables.
            all_reduced_terms = []
            reduced_expressions = []
            for reduc_var, terms in reduction_vars.iteritems():

                # Add current terms to list of all variables that have been reduced.
                all_reduced_terms += list(terms)

                # Create variable that we will use to reduce the terms.
                reduction_var = None
                if len(reduc_var) > 1:
                    reduction_var = create_product(list(reduc_var))
                else:
                    reduction_var = reduc_var[0]

                # Reduce all terms that need to be reduced.
                reduced_terms = [t.reduce_var(reduction_var) for t in terms]

                # Create reduced expression.
                reduced_expr = None
                if len(reduced_terms) > 1:
                    # Try to reduce the reduced terms further.
                    reduced_expr = create_product([reduction_var, create_sum(reduced_terms).reduce_ops()])
                else:
                    reduced_expr = create_product(reduction_var, reduced_terms[0])

                # Add reduced expression to list of reduced expressions.
                reduced_expressions.append(reduced_expr)

            # Create list of terms that should not be reduced.
            dont_reduce_terms = []
            for v in new_sum.vrs:
                if not v in all_reduced_terms:
                    dont_reduce_terms.append(v)

            # Create expression from terms that was not reduced.
            not_reduced_expr = None
            if dont_reduce_terms and len(dont_reduce_terms) > 1:
                # Try to reduce the remaining terms that were not reduced at first.
                not_reduced_expr = create_sum(dont_reduce_terms).reduce_ops()
            elif dont_reduce_terms:
                not_reduced_expr = dont_reduce_terms[0]

            # Create return expression.
            if not_reduced_expr:
                self._reduced = create_sum(reduced_expressions + [not_reduced_expr])
            elif len(reduced_expressions) > 1:
                self._reduced = create_sum(reduced_expressions)
            else:
                self._reduced = reduced_expressions[0]
            return self._reduced

        # Return self if we don't have any variables for which we can reduce
        # the sum.
        self._reduced = self
        return self._reduced

    def reduce_vartype(self, var_type):
        """Reduce expression with given var_type. It returns a list of tuples
        [(found, remain)], where 'found' is an expression that only has variables
        of type == var_type. If no variables are found, found=(). The 'remain'
        part contains the leftover after division by 'found' such that:
        self = Sum([f*r for f,r in self.reduce_vartype(Type)])."""
        found = {}
        # Loop members and reduce them by vartype.
        for v in self.vrs:
            f, r = v.reduce_vartype(var_type)
            if f in found:
                found[f].append(r)
            else:
                found[f] = [r]

        # Create the return value.
        returns = []
        for f, r in found.iteritems():
            if len(r) > 1:
                # Use expand to group expressions.
                r = create_sum(r).expand()
            elif r:
                r = r.pop()
            returns.append((f, r))
        return returns

def overlap(l, d):
    "Check if a member in list l is in the value (list) of dictionary d."
    for m in l:
        for k, v in d.iteritems():
            if m in v:
                return True
    return False

def group_fractions(expr):
    "Group Fractions in a Sum: 2/x + y/x -> (2 + y)/x."
    if expr._prec != 3: # sum
        return expr

    # Loop variables and group those with common denominator.
    not_frac = []
    fracs = {}
    for v in expr.vrs:
        if v._prec == 4: # frac
            if v.denom in fracs:
                fracs[v.denom][1].append(v.num)
                fracs[v.denom][0] += 1
            else:
                fracs[v.denom] = [1, [v.num], v]
            continue
        not_frac.append(v)
    if not fracs:
        return expr

    # Loop all fractions and create new ones using an appropriate numerator.
    for k, v in fracs.iteritems():
        if v[0] > 1:
            # TODO: Is it possible to avoid expanding the Sum?
            # I think we have to because x/a + 2*x/a -> 3*x/a.
            not_frac.append(create_fraction(create_sum(v[1]).expand(), k))
        else:
            not_frac.append(v[2])

    # Create return value.
    if len(not_frac) > 1:
        return create_sum(not_frac)
    return not_frac[0]

from floatvalue_obj import FloatValue
from symbol_obj     import Symbol
from product_obj    import Product
from fraction_obj   import Fraction

