__copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner"

__license__ = """
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

from pymbolic.mapper import RecursiveMapper
from pymbolic.primitives import Variable


def unify_map(map1, map2):
    result = map1.copy()
    for name, value in map2.items():
        if name in map1:
            if map1[name] != value:
                return None
        else:
            result[name] = value

    return result


class UnificationRecord:

    def __init__(self, equations, lmap=None, rmap=None):
        self.equations = equations

        # lmap and rmap just serve as a tool to reject
        # some unifications early.

        if lmap is None or rmap is None:
            lmap = {}
            rmap = {}

            for lhs, rhs in equations:
                if isinstance(lhs, Variable):
                    lmap[lhs.name] = rhs
                if isinstance(rhs, Variable):
                    rmap[rhs.name] = lhs

        self.lmap = lmap
        self.rmap = rmap

    def unify(self, other):
        new_lmap = unify_map(self.lmap, other.lmap)
        if new_lmap is None:
            return None

        new_rmap = unify_map(self.rmap, other.rmap)
        if new_rmap is None:
            return None

        # Merge redundant equations.
        new_equations = set(self.equations)
        new_equations.update(other.equations)

        return UnificationRecord(
            list(new_equations), new_lmap, new_rmap)

    def __repr__(self):
        return "UnificationRecord({})".format(
                ", ".join(f"{lhs} = {rhs}" for lhs, rhs in self.equations)
                )


def unify_many(unis1, uni2):
    result = []
    for uni1 in unis1:
        unif_result = uni1.unify(uni2)
        if unif_result is not None:
            result.append(unif_result)

    return result


class UnifierBase(RecursiveMapper):
    # The idea of the algorithm here is that the unifier accumulates a set of
    # unification possibilities (:class:`UnificationRecord`) as it descends the
    # expression tree. :func:`unify_many` above then checks if these possibilities
    # are consistent with new incoming information (also encoded as a
    # :class:`UnificationRecord`) and either augments or abandons them.

    def __init__(self, lhs_mapping_candidates=None,
            rhs_mapping_candidates=None,
            force_var_match=True):
        """
        :arg lhs_mapping_candidates: list or set of  variable names that may be
          assigned in the left-hand ("first") expression
        :arg rhs_mapping_candidates: list or set of  variable names that may be
          assigned in the right-hand ("second") expression
        :arg force_var_match: In the (unimplemented) case of bidirectional
          unification, only assign to variable names, don't make matches
          between higher-level expressions.
        """

        self.lhs_mapping_candidates = lhs_mapping_candidates
        self.rhs_mapping_candidates = rhs_mapping_candidates
        self.force_var_match = force_var_match

    def treat_mismatch(self, expr, other, urecs):
        raise NotImplementedError

    def unification_record_from_equation(self, lhs, rhs):
        if isinstance(lhs, (tuple, list)) or isinstance(rhs, (tuple, list)):
            # Always force lists/tuples to agree elementwise, never
            # generate a unification record between them directly.
            # This pushes the matching process down to the elementwise
            # level.

            return None

        lhs_is_var = isinstance(lhs, Variable)
        rhs_is_var = isinstance(rhs, Variable)

        if self.force_var_match and not (lhs_is_var or rhs_is_var):
            return None

        if (self.lhs_mapping_candidates is not None
                and lhs_is_var
                and lhs.name not in self.lhs_mapping_candidates):
            return None
        if (self.rhs_mapping_candidates is not None
                and rhs_is_var
                and rhs.name not in self.rhs_mapping_candidates):
            return None

        return UnificationRecord([(lhs, rhs)])

    def map_constant(self, expr, other, urecs):
        if expr == other:
            return urecs
        else:
            return []

    def map_variable(self, expr, other, urecs):
        new_uni_record = self.unification_record_from_equation(
                expr, other)

        if new_uni_record is None:
            # Check if the variables match literally--that's ok, too.
            if (isinstance(other, Variable)
                    and other.name == expr.name
                    and expr.name not in self.lhs_mapping_candidates):
                return urecs
            else:
                return []
        else:
            return unify_many(urecs, new_uni_record)

    def map_call(self, expr, other, urecs):
        if not isinstance(other, type(expr)):
            return self.treat_mismatch(expr, other, urecs)

        return self.rec(expr.function, other.function,
                self.rec(expr.parameters, other.parameters, urecs))

    def map_subscript(self, expr, other, urecs):
        if not isinstance(other, type(expr)):
            return self.treat_mismatch(expr, other, urecs)

        # {{{ unpack length-1 index tuples to avoid ambiguity

        expr_index = expr.index
        if isinstance(expr_index, tuple) and len(expr_index) == 1:
            expr_index, = expr_index

        other_index = other.index
        if isinstance(other_index, tuple) and len(other_index) == 1:
            other_index, = other_index

        # }}}

        return self.rec(expr.aggregate, other.aggregate,
                self.rec(expr_index, other_index, urecs))

    def map_lookup(self, expr, other, urecs):
        if not isinstance(other, type(expr)):
            return self.treat_mismatch(expr, other, urecs)
        if expr.name != other.name:
            return []

        return self.rec(expr.aggregate, other.aggregate, urecs)

    def map_sum(self, expr, other, urecs):
        if (not isinstance(other, type(expr))
                or len(expr.children) != len(other.children)):
            return []

        result = []

        from pytools import generate_permutations
        had_structural_match = False
        for perm in generate_permutations(range(len(expr.children))):
            it_assignments = urecs

            for my_child, other_child in zip(
                    expr.children,
                    (other.children[i] for i in perm)):
                it_assignments = self.rec(my_child, other_child, it_assignments)
                if not it_assignments:
                    break

            if it_assignments:
                had_structural_match = True
                result.extend(it_assignments)

        if not had_structural_match:
            return self.treat_mismatch(expr, other, urecs)

        return result

    map_product = map_sum

    def map_quotient(self, expr, other, urecs):
        if not isinstance(other, type(expr)):
            return self.treat_mismatch(expr, other, urecs)

        return self.rec(expr.numerator, other.numerator,
                self.rec(expr.denominator, other.denominator, urecs))

    map_floor_div = map_quotient
    map_remainder = map_quotient

    def map_power(self, expr, other, urecs):
        if not isinstance(other, type(expr)):
            return self.treat_mismatch(expr, other, urecs)

        return self.rec(expr.base, other.base,
                self.rec(expr.exponent, other.exponent, urecs))

    def map_left_shift(self, expr, other, urecs):
        if not isinstance(other, type(expr)):
            return self.treat_mismatch(expr, other, urecs)

        return self.rec(expr.shiftee, other.shiftee,
                self.rec(expr.shift, other.shift, urecs))

    map_right_shift = map_left_shift

    def map_bitwise_not(self, expr, other, urecs):
        if not isinstance(other, type(expr)):
            return self.treat_mismatch(expr, other, urecs)

        return self.rec(expr.child, other.child, urecs)

    map_bitwise_or = map_sum
    map_bitwise_xor = map_sum
    map_bitwise_and = map_sum

    map_logical_not = map_bitwise_not
    map_logical_or = map_sum
    map_logical_and = map_sum

    def map_comparison(self, expr, other, urecs):
        if (not isinstance(other, type(expr))
                or expr.operator != other.operator):
            return self.treat_mismatch(expr, other, urecs)

        return self.rec(expr.left, other.left,
                self.rec(expr.right, other.right, urecs))

    def map_if_positive(self, expr, other, urecs):
        if not isinstance(other, type(expr)):
            return self.treat_mismatch(expr, other, urecs)

        return self.rec(expr.criterion, other.criterion,
                self.rec(expr.then, other.then,
                    self.rec(expr.else_, other.else_, urecs)))

    def map_if(self, expr, other, urecs):
        if not isinstance(other, type(expr)):
            return self.treat_mismatch(expr, other, urecs)

        return self.rec(expr.condition, other.condition,
                self.rec(expr.then, other.then,
                    self.rec(expr.else_, other.else_, urecs)))

    map_min = map_sum
    map_max = map_sum

    def map_list(self, expr, other, urecs):
        if (not isinstance(other, type(expr))
                or len(expr) != len(other)):
            return []

        for my_child, other_child in zip(expr, other):
            urecs = self.rec(my_child, other_child, urecs)
            if not urecs:
                break

        return urecs

    map_tuple = map_list

    def __call__(self, expr, other, urecs=None):
        if urecs is None:
            urecs = [UnificationRecord([])]
        return self.rec(expr, other, urecs)


class UnidirectionalUnifier(UnifierBase):
    """Finds assignments of variables encountered in the
    first ("left") expression to subexpression of the second
    ("right") expression.
    """

    def treat_mismatch(self, expr, other, urecs):
        return []

    def map_commut_assoc(self, expr, other, urecs, factory):
        if not isinstance(other, type(expr)):
            return

        # Partition expr into terms that are plain (free) variables and those
        # that are not.
        plain_var_candidates = []
        non_var_children = []
        for child in expr.children:
            if (isinstance(child, Variable)
                    and child.name in self.lhs_mapping_candidates):
                plain_var_candidates.append(child)
            else:
                non_var_children.append(child)

        # list (with indices matching non_var_children) of
        #   list of tuples (other_index, unifiers)
        unification_candidates = []

        # Unify non-free-variable children of expr with children of the other
        # expr.
        for my_child in non_var_children:
            i_matches = []
            for j, other_child in enumerate(other.children):
                result = self.rec(my_child, other_child, urecs)
                if result:
                    i_matches.append((j, result))
            unification_candidates.append(i_matches)

        # Combine the unification candidates of children in all possible ways.
        def match_children(urec, next_cand_idx, other_leftovers):
            if next_cand_idx >= len(non_var_children):
                yield from match_plain_var_candidates(urec, other_leftovers)
                return

            for other_idx, pair_urecs in unification_candidates[next_cand_idx]:
                if other_idx not in other_leftovers:
                    # Don't re-match any elements.
                    continue

                new_urecs = unify_many(pair_urecs, urec)
                new_rhs_leftovers = other_leftovers - {other_idx}

                for cand_urec in new_urecs:
                    yield from match_children(
                            cand_urec, next_cand_idx + 1, new_rhs_leftovers)

        def match_plain_var_candidates(urec, other_leftovers):
            if len(plain_var_candidates) == len(other_leftovers) == 0:
                yield urec
                return

            # At this point, the values in plain_var_candidates have not
            # been matched in the lhs, and the values in other_leftovers
            # have not been matched in the rhs. Try all possible
            # combinations of matches (this part may become a performance
            # bottleneck and if needed could be optimized further).
            def subsets(s, max_size):
                from itertools import combinations
                for size in range(1, max_size + 1):
                    yield from combinations(s, size)

            def partitions(s, k):
                if k == 1:
                    yield [s]
                    return
                for subset in map(set, subsets(s, len(s) - k + 1)):
                    for partition in partitions(s - subset, k - 1):
                        yield [subset] + partition

            for partition in partitions(
                    other_leftovers, len(plain_var_candidates)):
                result = urec
                for subset, var in zip(partition, plain_var_candidates):
                    rec = self.unification_record_from_equation(
                        var, factory(other.children[i] for i in subset))
                    result = result.unify(rec)
                    if not result:
                        break
                else:
                    if len(non_var_children) != 0:
                        # urecs was merged in.
                        yield result
                        return
                    # urecs was not merged in, do it here.
                    yield from unify_many(urecs, result)

        for urec in match_children(
                UnificationRecord([]), 0, set(range(len(other.children)))):
            yield urec

    def map_sum(self, expr, other, unis):
        from pymbolic.primitives import flattened_sum
        return list(self.map_commut_assoc(expr, other, unis, flattened_sum))

    def map_product(self, expr, other, unis):
        from pymbolic.primitives import flattened_product
        return list(self.map_commut_assoc(expr, other, unis, flattened_product))
