import time
import os
import traceback
import collections
import threading

import nengo
from nengo import spa
import json

from nengo_gui.components.component import Component
from nengo_gui.components.value import Value
from nengo_gui.modal_js import infomodal
import nengo_gui.user_action
import nengo_gui.layout


class NetGraph(Component):
    config_defaults = {}
    configs = {}

    def __init__(self):
        # this component must be before all the normal graphs (so that
        # those other graphs are on top of the NetGraph), so its
        # order is between that of SimControl and the default (0)
        super(NetGraph, self).__init__(component_order=-5)

        # this lock ensures safety between check_for_reload() and update_code()
        self.code_lock = threading.Lock()
        self.new_code = None

        self.uids = {}
        self.parents = {}
        self.initialized_pan_and_zoom = False

    def attach(self, page, config, uid):
        super(NetGraph, self).attach(page, config, uid)
        self.layout = nengo_gui.layout.Layout(self.page.model)
        self.to_be_expanded = collections.deque([self.page.model])
        self.to_be_sent = collections.deque()

        self.networks_to_search = [self.page.model]

        try:
            self.last_modify_time = os.path.getmtime(self.page.filename)
        except OSError:
            self.last_modify_time = None
        except TypeError:  # happens if self.filename is None
            self.last_modify_time = None
        self.last_reload_check = time.time()

    def check_for_reload(self):
        if self.page.filename is not None:
            try:
                t = os.path.getmtime(self.page.filename)
                if self.last_modify_time is None or self.last_modify_time < t:
                    self.reload()
                    self.last_modify_time = t
            except OSError:
                pass

        with self.code_lock:
            new_code = self.new_code
            # the lock is in case update_code() is called between these lines
            self.new_code = None

        if new_code is not None:
            self.reload(code=new_code)

    def update_code(self, code):
        """Set new version of code to display."""
        with self.code_lock:
            self.new_code = code

    def reload(self, code=None):
        with self.page.lock:
            self._reload(code=code)

    def _reload(self, code=None):

        old_locals = self.page.last_good_locals
        old_default_labels = self.page.default_labels

        if code is None:
            with open(self.page.filename) as f:
                code = f.read()
            if self.page.code == code:
                # don't re-execute the identical code
                return
            else:
                # send the new code to the client
                self.page.ace_editor.update_code(code)

        self.page.execute(code)

        if self.page.error is not None:
            return

        name_finder = nengo_gui.NameFinder(self.page.locals, self.page.model)

        self.networks_to_search = [self.page.model]
        self.parents = {}

        removed_uids = {}
        rebuilt_objects = []

        # for each item in the old model, find the matching new item
        # for Nodes, Ensembles, and Networks, this means to find the item
        # with the same uid.  For Connections, we don't really have a uid,
        # so we use the uids of the pre and post objects.
        for uid, old_item in nengo.utils.compat.iteritems(dict(self.uids)):
            try:
                new_item = eval(uid, self.page.locals)
            except:
                new_item = None

            # check to make sure the new item's uid is the same as the
            # old item.  This is to catch situations where an old uid
            # happens to still refer to something in the new model, but that's
            # not the normal uid for that item.  For example, the uid
            # "ensembles[0]" might still refer to something even after that
            # ensemble is removed.
            new_uid = self.page.get_uid(new_item,
                        default_labels=name_finder.known_name)
            if new_uid != uid:
                new_item = None

            # find reasons to delete the object.  Any deleted object will
            # be recreated, so try to keep this to a minimum
            keep_object = True
            if new_item is None:
                keep_object = False
            elif not isinstance(new_item, old_item.__class__):
                # don't allow changing classes
                keep_object = False
            elif self.get_extra_info(new_item) != self.get_extra_info(old_item):
                keep_object = False

            if not keep_object:
                self.to_be_sent.append(dict(
                    type='remove', uid=uid))
                del self.uids[uid]
                removed_uids[old_item] = uid
                rebuilt_objects.append(uid)
            else:
                # fix aspects of the item that may have changed
                if self._reload_update_item(uid, old_item, new_item, 
                                            name_finder):
                    # something has changed about this object, so rebuild
                    # and components that use it
                    rebuilt_objects.append(uid)

                self.uids[uid] = new_item

        self.to_be_expanded.append(self.page.model)

        self.page.name_finder = name_finder
        self.page.default_labels = name_finder.known_name
        self.page.config = self.page.load_config()
        self.page.uid_prefix_counter = {}
        self.layout = nengo_gui.layout.Layout(self.page.model)
        self.page.code = code

        orphan_components = []
        rebuild_components = []

        # items that are shown in components, but not currently displayed
        #  in the NetGraph (i.e. things that are inside collapsed
        #  Networks, but whose values are being shown in a graph)
        collapsed_items = []

        removed_items = list(removed_uids.values())
        for c in self.page.components[:]:
            for item in c.code_python_args(old_default_labels):
                if item not in self.uids.keys() and item not in collapsed_items:

                    # item is a python string that is an argument to the
                    # constructor for the Component.  So it could be 'a',
                    # 'model.ensembles[3]', 'True', or even 'target=a'.
                    # We need to evaluate this string in the context of the
                    # locals dictionary and see what object it refers to
                    # so we can determine whether to rebuild this component.
                    #
                    # The following lambda should do this, handling both
                    # the normal argument case and the keyword argument case.
                    safe_eval = '(lambda *a, **b: list(a) + b.values())(%s)[0]'

                    # this Component depends on an item inside a collapsed
                    #  Network, so we need to check if that component has
                    #  changed or been removed
                    old_obj = eval(safe_eval % item, old_locals)

                    try:
                        new_obj = eval(safe_eval % item, self.page.locals)
                    except:
                        # the object this Component depends on no longer exists
                        new_obj = None

                    if new_obj is None:
                        removed_items.append(item)
                    elif not isinstance(new_obj, old_obj.__class__):
                        rebuilt_objects.append(item)
                    elif (self.get_extra_info(new_obj) != 
                          self.get_extra_info(old_obj)):
                        rebuilt_objects.append(item)

                    # add this to the list of collapsed items, so we
                    # don't recheck it if there's another Component that
                    # also depends on this
                    collapsed_items.append(item)

                if item in rebuilt_objects:
                    self.to_be_sent.append(dict(type='delete_graph',
                                                uid=c.original_id,
                                                notify_server=False))
                    rebuild_components.append(c.uid)
                    self.page.components.remove(c)
                    break
            else:
                for item in c.code_python_args(old_default_labels):
                    if item in removed_items:
                        self.to_be_sent.append(dict(type='delete_graph',
                                                    uid=c.original_id,
                                                    notify_server=False))
                        orphan_components.append(c)
                        break



        components = []
        # the old names for the old components
        component_uids = [c.uid for c in self.page.components]

        for k, v in list(self.page.locals.items()):
            if isinstance(v, nengo_gui.components.component.Component):
                # the object has been removed, so the Component should
                #  be removed as well
                if v in orphan_components:
                    continue

                # this is a Component that was previously removed,
                #  but is still in the config file, or it has to be
                #  rebuilt, so let's recover it
                if k not in component_uids:
                    self.page.add_component(v)
                    self.to_be_sent.append(dict(type='js', 
                                                code=v.javascript()))
                    components.append(v)
                    continue
                
                # otherwise, find the corresponding old Component
                index = component_uids.index(k)
                old_component = self.page.components[index]
                if isinstance(v, (nengo_gui.components.SimControlTemplate,
                                  nengo_gui.components.AceEditorTemplate,
                                  nengo_gui.components.NetGraphTemplate)):
                    # just keep these ones
                    components.append(old_component)
                else:
                    # replace these components with the newly generated ones
                    try:
                        self.page.add_component(v)
                        old_component.replace_with = v
                        v.original_id = old_component.original_id
                    except:
                        traceback.print_exc()
                        print('failed to recreate plot for %s' % v)
                    components.append(v)

        components.sort(key=lambda x: x.component_order)

        self.page.components = components

        self.page.changed = True

    def _reload_update_item(self, uid, old_item, new_item, new_name_finder):
        """Tell the client about changes to the item due to reload."""
        changed = False
        if isinstance(old_item, (nengo.Node,
                                 nengo.Ensemble,
                                 nengo.Network)):
            old_label = self.page.get_label(old_item)
            new_label = self.page.get_label(
                new_item, default_labels=new_name_finder.known_name)

            if old_label != new_label:
                self.to_be_sent.append(dict(
                    type='rename', uid=uid, name=new_label))
                changed = True
            if isinstance(old_item, nengo.Network):
                if self.page.config[old_item].expanded:
                    self.to_be_expanded.append(new_item)
                    changed = True

        elif isinstance(old_item, nengo.Connection):
            old_pre = old_item.pre_obj
            old_post = old_item.post_obj
            new_pre = new_item.pre_obj
            new_post = new_item.post_obj
            if isinstance(old_pre, nengo.ensemble.Neurons):
                old_pre = old_pre.ensemble
            if isinstance(old_post, nengo.ensemble.Neurons):
                old_post = old_post.ensemble
            if isinstance(new_pre, nengo.ensemble.Neurons):
                new_pre = new_pre.ensemble
            if isinstance(new_post, nengo.ensemble.Neurons):
                new_post = new_post.ensemble

            old_pre = self.page.get_uid(old_pre)
            old_post = self.page.get_uid(old_post)
            new_pre = self.page.get_uid(
                new_pre, default_labels=new_name_finder.known_name)
            new_post = self.page.get_uid(
                new_post, default_labels=new_name_finder.known_name)

            if new_pre != old_pre or new_post != old_post:
                # if the connection has changed, tell javascript
                pres = self.get_parents(
                    new_pre,
                    default_labels=new_name_finder.known_name)[:-1]
                posts = self.get_parents(
                    new_post,
                    default_labels=new_name_finder.known_name)[:-1]
                self.to_be_sent.append(dict(
                    type='reconnect', uid=uid,
                    pres=pres, posts=posts))
                changed = True
        return changed

    def get_parents(self, uid, default_labels=None):
        while uid not in self.parents:
            net = self.networks_to_search.pop(0)
            net_uid = self.page.get_uid(net, default_labels=default_labels)
            for n in net.nodes:
                n_uid = self.page.get_uid(n, default_labels=default_labels)
                self.parents[n_uid] = net_uid
            for e in net.ensembles:
                e_uid = self.page.get_uid(e, default_labels=default_labels)
                self.parents[e_uid] = net_uid
            for n in net.networks:
                n_uid = self.page.get_uid(n, default_labels=default_labels)
                self.parents[n_uid] = net_uid
                self.networks_to_search.append(n)
        parents = [uid]
        while parents[-1] in self.parents:
            parents.append(self.parents[parents[-1]])
        return parents

    def modified_config(self):
        self.page.modified_config()

    def update_client(self, client):
        now = time.time()
        if now > self.last_reload_check + 0.5:
            self.check_for_reload()
            self.last_reload_check = now

        if not self.initialized_pan_and_zoom:
            self.send_pan_and_zoom(client)
            self.initialized_pan_and_zoom = True

        while len(self.to_be_sent) > 0:
            info = self.to_be_sent.popleft()
            client.write(json.dumps(info))

        if len(self.to_be_expanded) > 0:
            with self.page.lock:
                network = self.to_be_expanded.popleft()
                self.expand_network(network, client)

    def javascript(self):
        return 'new Nengo.NetGraph(main, {uid:"%s"});' % id(self)

    def message(self, msg):
        try:
            info = json.loads(msg)
        except ValueError:
            print('invalid message', repr(msg))
            return
        action = info.get('act', None)
        undo = info.get('undo', None)
        if action is not None:
            del info['act']
            if action in ('auto_expand', 'auto_collapse'):
                getattr(self, 'act_' + action[5:])(**info)
            elif action in ('pan', 'zoom', 'create_modal'):
                # These should not use the undo stack
                getattr(self, 'act_' + action)(**info)
            else:
                act = nengo_gui.user_action.create_action(action, self, **info)
                self.page.undo_stack.append([act])
                del self.page.redo_stack[:]
        elif undo is not None:
            if undo == '1':
                self.undo()
            else:
                self.redo()
        else:
            print('received message', msg)

    def undo(self):
        if self.page.undo_stack:
            action = self.page.undo_stack.pop()
            re = []
            for act in action:
                act.undo()
                re.insert(0, act)
            self.page.redo_stack.append(re)

    def redo(self):
        if self.page.redo_stack:
            action = self.page.redo_stack.pop()
            un = []
            for act in action:
                act.apply()
                un.insert(0, act)
            self.page.undo_stack.append(un)

    def act_expand(self, uid):
        net = self.uids[uid]
        self.to_be_expanded.append(net)
        self.page.config[net].expanded = True
        self.modified_config()

    def act_collapse(self, uid):
        net = self.uids[uid]
        self.page.config[net].expanded = False
        self.remove_uids(net)
        self.modified_config()

    def remove_uids(self, net):
        for items in [net.ensembles, net.networks, net.nodes, net.connections]:
            for item in items:
                uid = self.page.get_uid(item)
                if uid in self.uids:
                    del self.uids[uid]
        for n in net.networks:
            self.remove_uids(n)

    def act_pan(self, x, y):
        self.page.config[self.page.model].pos = x, y
        self.modified_config()

    def act_zoom(self, scale, x, y):
        self.page.config[self.page.model].size = scale, scale
        self.page.config[self.page.model].pos = x, y
        self.modified_config()

    def act_create_modal(self, uid, **info):
        js = infomodal(self, uid, **info)
        self.to_be_sent.append(dict(type='js', code=js))

    def expand_network(self, network, client):
        if not self.page.config[network].has_layout:
            pos = self.layout.make_layout(network)
            for obj, layout in pos.items():
                self.page.config[obj].pos = layout['y'], layout['x']
                self.page.config[obj].size = layout['h'] / 2, layout['w'] / 2
            self.page.config[network].has_layout = True

        if network is self.page.model:
            parent = None
        else:
            parent = self.page.get_uid(network)
        for ens in network.ensembles:
            self.create_object(client, ens, type='ens', parent=parent)
        for node in network.nodes:
            self.create_object(client, node, type='node', parent=parent)
        for net in network.networks:
            self.create_object(client, net, type='net', parent=parent)
        for conn in network.connections:
            self.create_connection(client, conn, parent=parent)
        self.page.config[network].expanded = True

    def create_object(self, client, obj, type, parent):
        uid = self.page.get_uid(obj)
        if uid in self.uids:
            return

        pos = self.page.config[obj].pos
        if pos is None:
            import random
            pos = random.uniform(0, 1), random.uniform(0, 1)
            self.page.config[obj].pos = pos
        size = self.page.config[obj].size
        if size is None:
            size = (0.1, 0.1)
            self.page.config[obj].size = size
        label = self.page.get_label(obj)
        self.uids[uid] = obj
        info = dict(uid=uid, label=label, pos=pos, type=type, size=size,
                    parent=parent)
        if type == 'net':
            info['expanded'] = self.page.config[obj].expanded
        info.update(self.get_extra_info(obj))

        client.write(json.dumps(info))

    def get_extra_info(self, obj):
        '''Determine helper information for each nengo object.

        This is used by the client side to configure the display.  It is also
        used by the reload() code to determine if a NetGraph object should
        be recreated.
        '''
        info = {}
        if isinstance(obj, nengo.Node):
            if obj.output is None:
                info['passthrough'] = True
            if callable(obj.output) and hasattr(obj.output, '_nengo_html_'):
                info['html'] = True
            info['dimensions'] = int(obj.size_out)
        elif isinstance(obj, nengo.Ensemble):
            info['dimensions'] = int(obj.size_out)
            info['n_neurons'] = int(obj.n_neurons)
        elif Value.default_output(obj) is not None:
            info['default_output'] = True

        info['sp_targets'] = (
            nengo_gui.components.spa_plot.SpaPlot.applicable_targets(obj))
        return info

    def send_pan_and_zoom(self, client):
        pan = self.page.config[self.page.model].pos
        if pan is None:
            pan = 0, 0
        zoom = self.page.config[self.page.model].size
        if zoom is None:
            zoom = 1.0
        else:
            zoom = zoom[0]
        client.write(json.dumps(dict(type='pan', pan=pan)))
        client.write(json.dumps(dict(type='zoom', zoom=zoom)))

    def create_connection(self, client, conn, parent):
        uid = self.page.get_uid(conn)
        if uid in self.uids:
            return
        pre = conn.pre_obj
        if isinstance(pre, nengo.ensemble.Neurons):
            pre = pre.ensemble
        post = conn.post_obj
        if isinstance(post, nengo.ensemble.Neurons):
            post = post.ensemble
        pre = self.page.get_uid(pre)
        post = self.page.get_uid(post)
        self.uids[uid] = conn
        pres = self.get_parents(pre)[:-1]
        posts = self.get_parents(post)[:-1]
        info = dict(uid=uid, pre=pres, post=posts, type='conn', parent=parent)
        client.write(json.dumps(info))
