"""Frequent itemsets -- paired counts fp-tree"""
# encoding: utf-8

__author__ = 'thor'

"""
A Python implementation of the FP-growth algorithm that, additionally to item counts also maintains
counts on some statistics.

Basic usage of the module is very simple:
    >>> from fp_growth import find_frequent_itemsets
    >>> find_frequent_itemsets(transactions, values, minimum_support)

Based on https://github.com/enaeseth/python-fp-growth (MIT License) from Eric Naeseth <eric@naeseth.com>
"""

from collections import defaultdict, namedtuple


# from matplotlib.cbook import flatten


def find_frequent_itemsets(transactions, transaction_values, minimum_support):
    """
    Find frequent itemsets in the given transactions using FP-growth. This
    function returns a generator instead of an eagerly-populated list of items.
    The `transactions` parameter can be any iterable of iterables of items.
    `minimum_support` should be an integer specifying the minimum number of
    occurrences of an itemset for it to be accepted.
    Each item must be hashable (i.e., it must be valid as a member of a
    dictionary or a set).
    """

    fptree = mk_fptree(transactions, transaction_values, minimum_support)

    return find_frequent_itemsets_from_fp_tree(fptree, minimum_support)


def mk_fptree(transactions, transaction_values, minimum_support, item_counts=None):
    # Load the passed-in transactions and count the support that individual
    if item_counts is None:
        item_counts = defaultdict(lambda: 0)  # mapping from items to their supports
        # processed_transactions = []
        for transaction in transactions:
            # processed = []
            for item in transaction:
                item_counts[item] += 1
                # processed.append(item)
            # processed_transactions.append(processed)

    # Remove infrequent items from the item support dictionary.
    item_counts = dict(
        (item, support)
        for item, support in item_counts.items()
        if support >= minimum_support
    )

    # Build our FP-tree. Before any transactions can be added to the tree, they
    # must be stripped of infrequent items and their surviving items must be
    # sorted in decreasing order of frequency.
    def clean_transaction(transaction):
        transaction = [v for v in transaction if v in item_counts]
        transaction.sort(key=lambda v: item_counts[v], reverse=True)
        return transaction

    fptree = FPTree()
    for transaction, transaction_value in zip(
        map(clean_transaction, transactions), transaction_values
    ):
        fptree.add(transaction, transaction_value)
        # print("transaction={}, transaction_value={}".format(transaction, transaction_value))
        # print(tree.inspect())
        # print("-----------------------")

    return fptree


def find_frequent_itemsets_from_fp_tree(fp_tree, minimum_support):
    """
    Find frequent itemsets from a given fp_tree
    """

    def find_with_suffix(tree, suffix):
        for item, nodes in list(tree.items()):
            # support = sum(n.count for n in nodes)
            support, value = list(
                map(sum, list(zip(*[(n.count, n.value) for n in nodes])))
            )
            if support >= minimum_support and item not in suffix:
                # New winner!
                found_set = [item] + suffix
                yield (found_set, support, value)
                # yield (found_set, support) if include_support else found_set

                # Build a conditional tree and recursively search for frequent
                # itemsets within it.
                cond_tree = conditional_tree_from_paths(
                    tree.prefix_paths(item), minimum_support
                )
                for s in find_with_suffix(cond_tree, found_set):
                    yield s  # pass along the good news to our caller

    # Search for frequent itemsets, and yield the results we find.
    for itemset in find_with_suffix(fptree, []):
        yield itemset


class FPTree(object):
    """
    An FP tree.
    This object may only store transaction items that are hashable (i.e., all
    items must be valid as dictionary keys or set members).
    """

    Route = namedtuple('Route', 'head tail')

    def __init__(self):
        # The root node of the tree.
        self._root = FPNode(self, None, None, None)

        # A dictionary mapping items to the head and tail of a path of
        # "neighbors" that will hit every node containing that item.
        self._routes = {}

    @property
    def root(self):
        """The root node of the tree."""
        return self._root

    def add(self, transaction, transaction_value):
        """
        Adds a transaction to the tree.
        """

        point = self._root

        for item in transaction:
            next_point = point.search(item)
            if next_point:
                # There is already a node in this tree for the current transaction item; reuse it.
                # print("+++ before: ({}, {}, {})".format(next_point.item, next_point.count, next_point.value))
                next_point.increment_count()
                next_point.increment_value(transaction_value)
                # print("+++ after: ({}, {}, {})".format(next_point.item, next_point.count, next_point.value))
            else:
                # Create a new point and add it as a child of the point we're currently looking at.
                next_point = FPNode(self, item, count=1, value=transaction_value)
                # print("--- ({}, {}, {})".format(next_point.item, next_point.count, next_point.value))
                point.add(next_point)
                # Update the route of nodes that contain this item to include our new node.
                self._update_route(next_point)

            point = next_point

    def _update_route(self, point):
        """Add the given node to the route through all nodes for its item."""
        assert self is point.tree

        try:
            route = self._routes[point.item]
            route[1].neighbor = point  # route[1] is the tail
            self._routes[point.item] = self.Route(route[0], point)
        except KeyError:
            # First node for this item; start a new route.
            self._routes[point.item] = self.Route(point, point)

    def items(self):
        """
        Generate one 2-tuples for each item represented in the tree. The first
        element of the tuple is the item itself, and the second element is a
        generator that will yield the nodes in the tree that belong to the item.
        """
        for item in self._routes.keys():
            yield (item, self.nodes(item))

    def nodes(self, item):
        """
        Generates the sequence of nodes that contain the given item.
        """

        try:
            node = self._routes[item][0]
        except KeyError:
            return

        while node:
            yield node
            node = node.neighbor

    def prefix_paths(self, item):
        """Generates the prefix paths that end with the given item."""

        def collect_path(node):
            path = []
            # print("---{}".format(node))
            while node and not node.root:
                # print node
                path.append(node)
                node = node.parent
            path.reverse()
            return path

        return (collect_path(node) for node in self.nodes(item))

    def inspect(self):
        print('Tree:')
        self.root.inspect(1)

        print()
        print('Routes:')
        for item, nodes in list(self.items()):
            print('  %r' % item)
            for node in nodes:
                print('    %r' % node)

    def _removed(self, node):
        """Called when `node` is removed from the tree; performs cleanup."""

        head, tail = self._routes[node.item]
        if node is head:
            if node is tail or not node.neighbor:
                # It was the sole node.
                del self._routes[node.item]
            else:
                self._routes[node.item] = self.Route(node.neighbor, tail)
        else:
            for n in self.nodes(node.item):
                if n.neighbor is node:
                    n.neighbor = node.neighbor  # skip over
                    if node is tail:
                        self._routes[node.item] = self.Route(head, n)
                    break


def conditional_tree_from_paths(paths, minimum_support):
    """Builds a conditional FP-tree from the given prefix paths."""
    tree = FPTree()
    condition_item = None
    items = set()

    # Import the nodes in the paths into the new tree. Only the counts of the
    # leaf notes matter; the remaining counts will be reconstructed from the
    # leaf counts.
    for path in paths:
        if condition_item is None:
            condition_item = path[-1].item

        point = tree.root
        for node in path:
            next_point = point.search(node.item)
            if not next_point:
                # Add a new node to the tree.
                items.add(node.item)
                count = node.count if node.item == condition_item else 0
                value = node.value if node.item == condition_item else 0
                next_point = FPNode(tree, node.item, count, value)
                point.add(next_point)
                tree._update_route(next_point)
            point = next_point

    assert condition_item is not None

    # Calculate the counts of the non-leaf nodes.
    for path in tree.prefix_paths(condition_item):
        count = path[-1].count
        value = path[-1].value
        for node in reversed(path[:-1]):
            node._count += count
            node._value += value

    # Eliminate the nodes for any items that are no longer frequent.
    for item in items:
        support = sum(n.count for n in tree.nodes(item))
        if support < minimum_support:
            # Doesn't make the cut anymore
            for node in tree.nodes(item):
                if node.parent is not None:
                    node.parent.remove(node)

    # Finally, remove the nodes corresponding to the item for which this
    # conditional tree was generated.
    for node in tree.nodes(condition_item):
        if node.parent is not None:  # the node might already be an orphan
            node.parent.remove(node)

    return tree


class FPNode(object):
    """A node in an FP tree."""

    def __init__(self, tree, item, count, value):
        self._tree = tree
        self._item = item
        self._count = count
        self._value = value
        self._parent = None
        self._children = {}
        self._neighbor = None

    def add(self, child):
        """Adds the given FPNode `child` as a child of this node."""

        if not isinstance(child, FPNode):
            raise TypeError('Can only add other FPNodes as children')

        if not child.item in self._children:
            self._children[child.item] = child
            child.parent = self

    def search(self, item):
        """
        Checks to see if this node contains a child node for the given item.
        If so, that node is returned; otherwise, `None` is returned.
        """

        try:
            return self._children[item]
        except KeyError:
            return None

    def remove(self, child):
        try:
            if self._children[child.item] is child:
                del self._children[child.item]
                child.parent = None
                self._tree._removed(child)
                for sub_child in child.children:
                    try:
                        # Merger case: we already have a child for that item, so
                        # add the sub-child's count to our child's count.
                        self._children[sub_child.item]._count += sub_child.count
                        self._children[sub_child.item]._value += sub_child.value
                        sub_child.parent = None  # it's an orphan now
                    except KeyError:
                        # Turns out we don't actually have a child, so just add
                        # the sub-child as our own child.
                        self.add(sub_child)
                child._children = {}
            else:
                raise ValueError('that node is not a child of this node')
        except KeyError:
            raise ValueError('that node is not a child of this node')

    def __contains__(self, item):
        return item in self._children

    @property
    def tree(self):
        """The tree in which this node appears."""
        return self._tree

    @property
    def item(self):
        """The item contained in this node."""
        return self._item

    @property
    def count(self):
        """The count associated with this node's item."""
        return self._count

    @property
    def value(self):
        """The count associated with this node's item."""
        return self._value

    def increment_count(self):
        """Increments the count associated with this node's item."""
        if self._count is None:
            raise ValueError('Root nodes have no associated count.')
        self._count += 1

    def increment_value(self, value_to_add=1.0):
        """Increments the count associated with this node's item."""
        if self._value is None:
            raise ValueError('Root nodes have no associated value.')
        self._value += value_to_add

    @property
    def root(self):
        """True if this node is the root of a tree; false if otherwise."""
        return self._item is None and self._count is None

    @property
    def leaf(self):
        """True if this node is a leaf in the tree; false if otherwise."""
        return len(self._children) == 0

    def parent():
        doc = "The node's parent."

        def fget(self):
            return self._parent

        def fset(self, value):
            if value is not None and not isinstance(value, FPNode):
                raise TypeError('A node must have an FPNode as a parent.')
            if value and value.tree is not self.tree:
                raise ValueError('Cannot have a parent from another tree.')
            self._parent = value

        return locals()

    parent = property(**parent())

    def neighbor():
        doc = '''
        The node's neighbor; the one with the same value that is "to the right"
        of it in the tree.
        '''

        def fget(self):
            return self._neighbor

        def fset(self, value):
            if value is not None and not isinstance(value, FPNode):
                raise TypeError('A node must have an FPNode as a neighbor.')
            if value and value.tree is not self.tree:
                raise ValueError('Cannot have a neighbor from another tree.')
            self._neighbor = value

        return locals()

    neighbor = property(**neighbor())

    @property
    def children(self):
        """The nodes that are children of this node."""
        return tuple(self._children.values())

    def inspect(self, depth=0):
        print(('  ' * depth) + repr(self))
        for child in self.children:
            child.inspect(depth + 1)

    def __repr__(self):
        if self.root:
            return '<%s (root)>' % type(self).__name__
        return '<%s %r count=%r, value=%r>' % (
            type(self).__name__,
            self.item,
            self.count,
            self.value,
        )


# if __name__ == '__main__':
#     from optparse import OptionParser
#     import csv

#     p = OptionParser(usage='%prog data_file')
#     p.add_option('-s', '--minimum-support', dest='minsup', type='int',
#         help='Minimum itemset support (default: 2)')
#     p.set_defaults(minsup=2)

#     options, args = p.parse_args()
#     if len(args) < 1:
#         p.error('must provide the path to a CSV file to read')

#     f = open(args[0])
#     try:
#         for itemset, support in find_frequent_itemsets(csv.reader(f), options.minsup, True):
#             print '{' + ', '.join(itemset) + '} ' + str(support)
#     finally:
#         f.close()
