import sys
import time

from theano.gof.python25 import partial

import graph


class AlreadyThere(Exception):
    pass


class ReplacementDidntRemovedError(Exception):
    """This exception should be thrown by replace_all_validate_remove
    when an optimization wanted to remove a Variable or a Node from
    the graph, but the replacement it gived didn't do that.

    """
    pass


class Bookkeeper:

    def on_attach(self, fgraph):
        for node in graph.io_toposort(fgraph.inputs, fgraph.outputs):
            self.on_import(fgraph, node)

    def on_detach(self, fgraph):
        for node in graph.io_toposort(fgraph.inputs, fgraph.outputs):
            self.on_prune(fgraph, node)


class History:

    def __init__(self):
        self.history = {}

    def on_attach(self, fgraph):
        if hasattr(fgraph, 'checkpoint') or hasattr(fgraph, 'revert'):
            raise AlreadyThere("History feature is already present or in"
                               " conflict with another plugin.")
        self.history[fgraph] = []
        fgraph.checkpoint = lambda: len(self.history[fgraph])
        fgraph.revert = partial(self.revert, fgraph)

    def on_detach(self, fgraph):
        del fgraph.checkpoint
        del fgraph.revert
        del self.history[fgraph]

    def on_change_input(self, fgraph, node, i, r, new_r, reason=None):
        if self.history[fgraph] is None:
            return
        h = self.history[fgraph]
        h.append(lambda: fgraph.change_input(node, i, r,
                                          reason=("Revert", reason)))

    def revert(self, fgraph, checkpoint):
        """
        Reverts the graph to whatever it was at the provided
        checkpoint (undoes all replacements).  A checkpoint at any
        given time can be obtained using self.checkpoint().
        """
        h = self.history[fgraph]
        self.history[fgraph] = None
        while len(h) > checkpoint:
            f = h.pop()
            f()
        self.history[fgraph] = h


class Validator:

    def on_attach(self, fgraph):
        for attr in ('validate', 'validate_time'):
            if hasattr(fgraph, attr):
                raise AlreadyThere("Validator feature is already present or in"
                                   " conflict with another plugin.")

        def validate():
            t0 = time.time()
            ret = fgraph.execute_callbacks('validate')
            t1 = time.time()
            if fgraph.profile:
                fgraph.profile.validate_time += t1 - t0
            return ret

        fgraph.validate = validate

        def consistent():
            try:
                fgraph.validate()
                return True
            except Exception:
                return False
        fgraph.consistent = consistent

    def on_detach(self, fgraph):
        del fgraph.validate
        del fgraph.consistent


class ReplaceValidate(History, Validator):

    def on_attach(self, fgraph):
        History.on_attach(self, fgraph)
        Validator.on_attach(self, fgraph)
        for attr in ('replace_validate', 'replace_all_validate'):
            if hasattr(fgraph, attr):
                raise AlreadyThere("ReplaceValidate feature is already present"
                                   " or in conflict with another plugin.")
        fgraph.replace_validate = partial(self.replace_validate, fgraph)
        fgraph.replace_all_validate = partial(self.replace_all_validate, fgraph)
        fgraph.replace_all_validate_remove = partial(
            self.replace_all_validate_remove, fgraph)

    def on_detach(self, fgraph):
        History.on_detach(self, fgraph)
        Validator.on_detach(self, fgraph)
        del fgraph.replace_validate
        del fgraph.replace_all_validate
        del fgraph.replace_all_validate_remove

    def replace_validate(self, fgraph, r, new_r, reason=None):
        self.replace_all_validate(fgraph, [(r, new_r)], reason=reason)

    def replace_all_validate(self, fgraph, replacements, reason=None):
        chk = fgraph.checkpoint()
        for r, new_r in replacements:
            try:
                fgraph.replace(r, new_r, reason=reason)
            except Exception, e:
                if ('The type of the replacement must be the same' not in
                    str(e) and 'does not belong to this FunctionGraph' not in str(e)):
                    out = sys.stderr
                    print >> out, "<<!! BUG IN FGRAPH.REPLACE OR A LISTENER !!>>",
                    print >> out, type(e), e, reason
                # this might fail if the error is in a listener:
                # (fgraph.replace kinda needs better internal error handling)
                fgraph.revert(chk)
                raise
        try:
            fgraph.validate()
        except Exception, e:
            fgraph.revert(chk)
            raise
        return chk

    def replace_all_validate_remove(self, fgraph, replacements,
                                    remove, reason=None, warn=True):
        """As replace_all_validate, revert the replacement if the ops
        in the list remove are still in the graph. It also print a warning.

        """
        chk = fgraph.replace_all_validate(replacements, reason)
        for rm in remove:
            if rm in fgraph.nodes or rm in fgraph.variables:
                fgraph.revert(chk)
                if warn:
                    out = sys.stderr
                    print >> out, (
                        "WARNING: An optimization wanted to replace a Variable"
                        " in the graph, but the replacement for it doesn't"
                        " remove it. We disabled the optimization."
                        " Your function runs correctly, but it would be"
                        " appreciated if you submit this problem to the"
                        " mailing list theano-users so that we can fix it.")
                    print >> out, reason, replacements
                raise ReplacementDidntRemovedError()


class NodeFinder(dict, Bookkeeper):

    def __init__(self):
        self.fgraph = None

    def on_attach(self, fgraph):
        if self.fgraph is not None:
            raise Exception("A NodeFinder instance can only serve one FunctionGraph.")
        if hasattr(fgraph, 'get_nodes'):
            raise AlreadyThere("NodeFinder is already present or in conflict"
                               " with another plugin.")
        self.fgraph = fgraph
        fgraph.get_nodes = partial(self.query, fgraph)
        Bookkeeper.on_attach(self, fgraph)

    def on_detach(self, fgraph):
        if self.fgraph is not fgraph:
            raise Exception("This NodeFinder instance was not attached to the"
                            " provided fgraph.")
        self.fgraph = None
        del fgraph.get_nodes
        Bookkeeper.on_detach(self, fgraph)

    def on_import(self, fgraph, node):
        try:
            self.setdefault(node.op, []).append(node)
        except TypeError:  # node.op is unhashable
            return
        except Exception, e:
            print >> sys.stderr, 'OFFENDING node', type(node), type(node.op)
            try:
                print >> sys.stderr, 'OFFENDING node hash', hash(node.op)
            except Exception:
                print >> sys.stderr, 'OFFENDING node not hashable'
            raise e

    def on_prune(self, fgraph, node):
        try:
            nodes = self[node.op]
        except TypeError:  # node.op is unhashable
            return
        nodes.remove(node)
        if not nodes:
            del self[node.op]

    def query(self, fgraph, op):
        try:
            all = self.get(op, [])
        except TypeError:
            raise TypeError("%s in unhashable and cannot be queried by the"
                            " optimizer" % op)
        all = list(all)
        return all


class PrintListener(object):

    def __init__(self, active=True):
        self.active = active

    def on_attach(self, fgraph):
        if self.active:
            print "-- attaching to: ", fgraph

    def on_detach(self, fgraph):
        if self.active:
            print "-- detaching from: ", fgraph

    def on_import(self, fgraph, node):
        if self.active:
            print "-- importing: %s" % node

    def on_prune(self, fgraph, node):
        if self.active:
            print "-- pruning: %s" % node

    def on_change_input(self, fgraph, node, i, r, new_r, reason=None):
        if self.active:
            print "-- changing (%s.inputs[%s]) from %s to %s" % (
                node, i, r, new_r)


class PreserveNames:
    def on_change_input(self, fgraph, mode, i, r, new_r, reason=None):
        if r.name is not None and new_r.name is None:
            new_r.name = r.name
