graph.py

import logging
import numpy as np
import types
from typing import Dict, Any, List, Optional
from datetime import datetime
import matplotlib.pyplot as plt
from os.path import join, split
import os
import tempfile
import webbrowser
import warnings

import networkx as nx
import autodepgraph
from autodepgraph.visualization import state_cmap
from autodepgraph import visualization as vis

# Used to find functions in modules
from importlib import import_module
try:
    # Only used for finding instrument methods.
    from qcodes.instrument.base import Instrument
except ImportError:
    Instrument = None


class AutoDepGraph_DAG(nx.DiGraph):
    """
    Attributes:
    ---------------
        node_states:
            Allowed states for the nodes
        matplotlib_edge_properties:
            Properties passed to networkx plotting of edges
        matplotlib_label_properties:
            Properties passed to networkx plotting of labels

    """
    node_states: List[str] = ['good', 'needs calibration',
                              'bad', 'unknown', 'active']
    matplotlib_edge_properties: Dict[str, Any] = {'edge_color': 'k', 'alpha': .8}
    matplotlib_label_properties: Dict[str, Any] = {'font_color': 'k'}

    def __init__(self, name, cfg_plot_mode='svg',
                 incoming_graph_data=None, **attr):
        """
        Directed Acyclic Graph used for calibrations.
        Inherits from a networkx DiGraph.

        The AutoDepGraph provides several constraints to nodes and edges
        as well as add the functionality to walk over the graph in order to
        calibrate a node and visualize the graph in real-time.
        """
        attr['name'] = name
        self.cfg_plot_mode = cfg_plot_mode
        self.cfg_plot_mode_args = {'fig': None}

        super().__init__(incoming_graph_data, **attr)

        # internal attributes
        self._DiGraphWindow = None  # used for pyqtgraph plotting
        self._graph_changed_since_plot = True  # used for pyqtgraph plotting

        # counters to count how often functions get called for debugging
        # purposes
        self._exec_cnt = 0
        self._calib_cnt = 0
        self._check_cnt = 0

    @property
    def cfg_svg_filename(self):
        """
        Default location for storing svg based visualizations of the DAG.
        """
        _path_name = split(__file__)[:-1][0]
        return join(_path_name, 'svg_viewer', 'adg_graph.svg')

    def fresh_copy(self):
        return AutoDepGraph_DAG(name=self.name,
                                cfg_plot_mode=self.cfg_plot_mode)

    def add_node(self, node_for_adding, **attr):
        """
        Adds a node to the graph, including starting attributes.

        attr:
            name        (type)  = default_value

            calibrate_function  = 'NotImplementedCalibration'
            check_functions     = 'return_fixed_value'
            tolerance   (float) = 0
            timeout     (float) = np.inf
            state       (str)   = 'unknown'

        """
        # there are set here to ensure these node attributes exist.
        # setting in this way will most likely interfere
        # with joining multiple graphs
        attr['timeout'] = attr.get('timeout',
                                   np.inf)
        attr['calibrate_function'] = attr.get(
            'calibrate_function',
            'autodepgraph.node_functions.calibration_functions' +
            '.NotImplementedCalibration')

        attr['calibrate_function_args'] = attr.get(
            'calibrate_function_args', None)

        attr['check_function'] = attr.get(
            'check_function',
            'autodepgraph.node_functions.check_functions' +
            '.return_fixed_value')

        # zero default tolerance -> always recalibrate
        attr['tolerance'] = attr.get('tolerance', 0)
        super().add_node(node_for_adding, **attr)

        self.set_node_state(node_for_adding,
                            state=attr.get('state', 'unknown'))
        # this adds a helper method to the attributes for quick access
        self._construct_maintenance_methods(nodes=[node_for_adding])

    def _construct_maintenance_methods(self, nodes):
        for n in nodes:
            self._construct_maintenance_method(node_name=n)

    def _construct_maintenance_method(self, node_name):
        node_name_no_space = node_name.replace(' ', '_').replace('-', '_')

        # This name exists so that text based storing falls back
        # on the right placeholder function
        def _construct_maintenance_method():
            self.maintain_node(node_name)

        self.__setattr__('maintain_{}'.format(node_name_no_space),
                         _construct_maintenance_method)

    def add_edge(self, u_of_edge, v_of_edge, **attr):
        """
        Adds an edge that denotes a dependency in the calibration graph.
        u_of_edge -> v_of_edge denotes that u depends on v.
        """

        # Nodes must already exist to ensure they have the right properties
        if u_of_edge not in self.nodes():
            raise KeyError('{} not in nodes'.format(u_of_edge))
        if v_of_edge not in self.nodes():
            raise KeyError('{} not in nodes'.format(v_of_edge))
        super().add_edge(u_of_edge, v_of_edge, **attr)

    def get_node_state(self, node_name):
        Delta_T = (datetime.now() -
                   self.nodes[node_name]['last_update']).total_seconds()
        if (Delta_T > self.nodes[node_name]['timeout']):
            self.nodes[node_name]['state'] = 'unknown'
        return self.nodes[node_name]['state']

    def set_node_state(self, node_name, state, update_monitor=True):
        if state not in self.node_states:
            raise IndexError(f'state {state} not in {self.node_states}')
        self.nodes[node_name]['state'] = state
        self.nodes[node_name]['last_update'] = datetime.now()
        if update_monitor:
            self.update_monitor()

    def is_manual_node(self, node_name):
        if isinstance(self.nodes[node_name]['calibrate_function'], (types.MethodType, types.FunctionType)):
            return False
        elif 'manual' in self.nodes[node_name]['calibrate_function']:
            return True
        elif 'NotImplemented' in self.nodes[node_name]['calibrate_function']:
            return True
        else:
            return False

    def maintain_node(self, node: str, verbose=False) -> str:
        """
        Maintaining a node attempts to go from any state to a good state.
            any_state -> good

        Maintain a node performs the following steps:
            1. get the state of the dependency nodes. If all are OK or unknown,
               perform a check on the node itself.
               If a node is any other state, execute it to move it to a
               good state
            2. perform the "check" experiment on the node itself. This quick
               check
            3. Perform calibration and second round of maintaining dependencies

        Returns:
            State of the node after maintaining the node
        Raises:
            Exception if the node could not be calibrated
        """
        self._exec_cnt += 1
        if verbose:
            print('Maintaining node "{}".'.format(node))

        # 1. Going over the states of all the required nodes and ensure
        # these are all in a 'Good' state.
        for req_node_name in self.adj[node]:
            req_node_state = self.nodes[req_node_name]['state']
            if req_node_state in ['good', 'unknown']:
                continue  # assume req_node is in a good state
            else:  # maintaining the node to ensure it is in a good state
                req_node_state = self.maintain_node(req_node_name,
                                                    verbose=verbose)
                if req_node_state == 'bad':
                    raise ValueError('Could not calibrate "{}"'.format(
                        req_node_name))

        # 2. Once all required nodes are OK, determine action to be taken
        state = self.nodes[node]['state']
        if state == 'needs calibration':
            # action is clear, no check required
            pass
        else:
            # determine latest status of node
            state = self.check_node(node, verbose=verbose)

        # 3. Take action based on the stae of the node
        if state == 'needs calibration':
            cal_succes = self.calibrate_node(node, verbose=verbose)
            # the calibration can still fail if dependencies that were good
            # or unknown were bad. In that case all dependencies will
            # explicitly be executed and calibration will be retried
            if not cal_succes:
                state = 'bad'
                if verbose:
                    print('Initial calibration of "{}" failed, '
                          'retrying.'.format(node))
        if state == 'bad':
            # if the state is bad it will execute *all* dependencies. Even
            # the ones that were updated before.
            if verbose:
                print('State of node "{}" is bad, maintaining all required'
                      ' nodes.'.format(node))
            for req_node_name in self.adj[node]:
                req_node_state = self.maintain_node(req_node_name,
                                                    verbose=verbose)
            cal_succes = self.calibrate_node(node, verbose=verbose)
            if not cal_succes:
                raise ValueError(
                    'Calibration of "{}" failed.'.format(node))

        state = self.nodes[node]['state']
        return state

    def check_node(self, node, verbose=False):
        """ Perform check method on specified node

        Args:
            node: Node to check
            verbose: Verbosity level
        Returns:
            Returns node state after the check
        """
        if verbose:
            print('\tChecking node {}.'.format(node))
        self.set_node_state(node, 'active')

        func = _get_function(self.nodes[node]['check_function'])
        result = func()
        if isinstance(result, float):
            if result < self.nodes[node]['tolerance']:
                self.set_node_state(node, 'good')
                if verbose:
                    print('\tNode {} is within tolerance.'.format(node))
            else:
                self.set_node_state(node, 'needs calibration')
                if verbose:
                    print('\tNode {} needs calibration.'.format(node))

        elif result == False:
            if verbose:
                print('\tNode {} is in a "bad" state.'.format(node))
            # if the function returns False it means the check is broken
            self.set_node_state(node, 'bad')
        else:
            raise ValueError('Expected float or "False", '
                             'result is: {}'.format(result))

        return self.nodes[node]['state']

    def calibrate_node(self, node: str, verbose: bool = False):
        """ Calibrate specified node

        Args:
            node: Node to calibration
            verbose: Verbosity level
        Returns:
            Returns True if the calibration was succesfull, otherwise False
        """
        if verbose:
            print('\tCalibrating node {}.'.format(node))
        self.set_node_state(node, 'active')

        func = _get_function(self.nodes[node]['calibrate_function'])
        try:
            if 'calibrate_function_args' not in self.nodes[node] or self.nodes[node]['calibrate_function_args'] is None:
                result = func()
            else:
                result = func(**self.nodes[node]['calibrate_function_args'])
        except Exception as e:
            self.set_node_state(node, 'bad')
            logging.warning(e)
            return False
        if result:
            self.set_node_state(node, 'good')
            if verbose:
                print('\tCalibration of node {} successful.'.format(node))

            return True
        else:
            self.set_node_state(node, 'bad')
            if verbose:
                print('\tCalibration of node {} failed.'.format(node))

            return False

    def set_all_node_states(self, state):
        for node_dat in self.nodes.values():
            node_dat['state'] = state

    def update_monitor(self):
        if self.cfg_plot_mode == 'matplotlib':
            self.update_monitor_mpl()
        elif self.cfg_plot_mode == 'svg':
            self.draw_svg()
        elif self.cfg_plot_mode is None or self.cfg_plot_mode == 'None':
            return
        else:
            raise ValueError('cfg_plot_mode should be in ["matplotlib",'
                             ' "svg", "None" ]')

    def update_monitor_mpl(self):
        """
        Updates a plot using the draw_graph_mpl based on matplotlib.
        """
        fig = self.cfg_plot_mode_args.get('fig', None)
        if fig is not None:
            plt.figure(fig)
        plt.clf()
        self.draw_mpl(plt.gca())
        plt.draw()
        plt.pause(.01)

    def _generate_node_positions(self, node_positions: Optional[dict] = None):
        nodes = self.nodes()

        if node_positions is None:
            node_positions = {}

        def position_generator(N=10, centre=[0, 5]):
            """ Generate circle of positions around centre """
            idx = 0
            while True:
                phi = 2*np.pi*(idx/N)+np.pi/2
                pos = 2.1*np.array([np.cos(phi), np.sin(phi)]) + centre
                yield pos
                idx = idx+1

        positions = position_generator(len(nodes))
        pos = dict([(node, node_positions.get(node, next(positions)))
                    for node in nodes])
        return pos

    def draw_mpl(self, ax=None):
        if ax is None:
            f, ax = plt.subplots()
            ax.axis('off')
        ax.set_title(self.name)
        colors_list = [state_cmap[node_dat['state']] for node_dat in
                       self.nodes.values()]
        node_positions = getattr(self, 'node_positions', None)
        if node_positions is None:
            pos = nx.nx_agraph.graphviz_layout(self, prog='dot')
        else:
            pos = self._generate_node_positions(node_positions)
        nx.draw_networkx_nodes(self, pos, ax=ax, node_color=colors_list)
        nx.draw_networkx_edges(self, pos, ax=ax, arrows=True,
                               **self.matplotlib_edge_properties)
        nx.draw_networkx_labels(
            self, pos, ax=ax, **self.matplotlib_label_properties)
        self._format_mpl_plot(ax)

    @staticmethod
    def _format_mpl_plot(ax):
        """ Method to format the generated matplotlib figure """
        ax.set_xticks([])
        ax.set_yticks([])

    def draw_svg(self, filename: str = None):
        """
        """
        if filename is None:
            filename = self.cfg_svg_filename
        self._update_drawing_attrs()
        vis.draw_graph_svg(self, filename)

    def open_html_viewer(self):
        """ Open html viewer for the file specified by the svg backend """
        template = os.path.join(os.path.split(autodepgraph.__file__)[
                                0], 'svg_viewer', 'svg_graph_viewer.html')
        with open(template, 'rt') as fid:
            x = fid.read()

        base, file = os.path.split(self.cfg_svg_filename)
        x = x.replace('adg_graph.svg', file)

        tfile = tempfile.mktemp(prefix='svgviewer-', suffix='.html', dir=base)
        with open(tfile, 'wt') as fid:
            fid.write(x)
        webbrowser.open_new_tab(tfile)
        return tfile

    def set_node_attribute(self, node, attribute, value):
        """ Set the attribute of the specified node

        Args:
            node (str): name of the node
            attribute (str): attribute to set
            value (ojbect): value to set
        """
        if attribute in ['state']:
            raise Exception('please use set_state directly')
        nx.set_node_attributes(self, {node: {attribute: value}})

    def get_node_attribute(self, node, attribute):
        """ Return the attribute of the specified node

        Args:
            node (str): name of the node
            attribute (str): attribute to get
        Returns:
            value (ojbect): attribute of the object
        """
        if attribute in ['state']:
            raise Exception('please use get_state directly')
        return self.nodes[node][attribute]

    def set_node_description(self, node, description):
        """ Set the node description field

        Args:
            node (str): name of the node
            description (str): description to set
        """
        nx.set_node_attributes(self, {node: {'description': description}})

    def calibration_state(self):
        """ Return dictionary with current calibration state """
        return dict(self.nodes)

    def _update_drawing_attrs(self):
        for node_name, node_attrs in self.nodes(True):

            state = self.get_node_state(node_name)
            color = vis.state_cmap[state]
            shape = 'hexagon' if self.is_manual_node(node_name) else 'ellipse'
            attr_dict = {'shape': shape,
                         'style': 'filled',
                         'color': color,
                         # 'fixedsize':'shape',
                         # 'fixedsize' : b"true",
                         'fixedsize': "false",
                         'fillcolor': color}
            node_attrs.update(attr_dict)


def _construct_maintenance_method():
    # This placeholder exists to allow reading and writing graphs in a graph
    # based format.
    print('Placeholder function.')
    print('Call DAG._construct_maintenance_methods() to update methods.')


def _get_function(funcStr):

    if isinstance(funcStr, (types.MethodType, types.FunctionType)):
        warnings.warn('please set function as a str', DeprecationWarning)
        f = funcStr
    elif '.' in funcStr:
        try:
            instr_name, method = funcStr.split('.')
            instr = Instrument.find_instrument(instr_name)
            f = getattr(instr, method)
        except Exception as e:
            f = get_function_from_module(funcStr)
    else:
        raise Exception('could not find function %s' % funcStr)
    return f


def get_function_from_module(funcStr):
    """
    """
    split_idx = funcStr.rfind('.')
    module_name = funcStr[:split_idx]
    mod = import_module(module_name)
    f = getattr(mod, funcStr[(split_idx+1):])
    return f


def update_node_state(graph_to_update, graph_to_update_from):
    for node_name, attrs in graph_to_update_from.nodes(True):
        if node_name in graph_to_update.nodes():
            graph_to_update.nodes[node_name]['state'] = attrs['state']
            graph_to_update.nodes[node_name]['last_update'] = attrs['last_update']