soil/server.py

import io
import logging
import networkx as nx
import os
import threading
import tornado.ioloop
import tornado.web
import tornado.websocket
import tornado.escape
import tornado.gen
import yaml
import webbrowser
from contextlib import contextmanager
from time import sleep
from xml.etree.ElementTree import tostring

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


class PageHandler(tornado.web.RequestHandler):
    """ Handler for the HTML template which holds the visualization. """

    def get(self):
        self.render('index.html', port=self.application.port,
                    name=self.application.name)


class SocketHandler(tornado.websocket.WebSocketHandler):
    """ Handler for websocket. """

    def open(self):
        if self.application.verbose:
            logger.info('Socket opened!')

    def check_origin(self, origin):
        return True

    def on_message(self, message):
        """ Receiving a message from the websocket, parse, and act accordingly. """

        msg = tornado.escape.json_decode(message)

        if msg['type'] == 'config_file':

            if self.application.verbose:
                print(msg['data'])

            self.config = list(yaml.load_all(msg['data']))

            if len(self.config) > 1:
                error = 'Please, provide only one configuration.'
                if self.application.verbose:
                    logger.error(error)
                self.write_message({'type': 'error',
                    'error': error})
                return
            else:
                self.config = self.config[0]
                self.send_log('INFO.' + self.application.simulator.name, 'Using config: {name}'.format(name=self.config['name']))

            if 'visualization_params' in self.config:
                self.write_message({'type': 'visualization_params',
                    'data': self.config['visualization_params']})
            self.name = self.config['name']
            self.run_simulation()

            settings = []
            for key in self.config['environment_params']: 
                if type(self.config['environment_params'][key]) == float or type(self.config['environment_params'][key]) == int:
                    if self.config['environment_params'][key] <= 1:
                        setting_type = 'number'
                    else:
                        setting_type = 'great_number'
                elif type(self.config['environment_params'][key]) == bool:
                    setting_type = 'boolean'
                else:
                    setting_type = 'undefined'

                settings.append({
                    'label': key,
                    'type': setting_type,
                    'value': self.config['environment_params'][key]
                })

            self.write_message({'type': 'settings',
                'data': settings})

        elif msg['type'] == 'get_trial':
            if self.application.verbose:
                logger.info('Trial {} requested!'.format(msg['data']))
            self.send_log('INFO.' + __name__, 'Trial {} requested!'.format(msg['data']))
            self.write_message({'type': 'get_trial',
                'data': self.get_trial( int(msg['data']) ) })

        elif msg['type'] == 'run_simulation':
            if self.application.verbose:
                logger.info('Running new simulation for {name}'.format(name=self.config['name']))
            self.send_log('INFO.' + self.application.simulator.name, 'Running new simulation for {name}'.format(name=self.config['name']))
            self.config['environment_params'] = msg['data']
            self.run_simulation()

        elif msg['type'] == 'download_gexf':
            G = self.simulation[ int(msg['data']) ].history_to_graph()
            for node in G.nodes():
                if 'pos' in G.node[node]:
                    G.node[node]['viz'] = {"position": {"x": G.node[node]['pos'][0], "y": G.node[node]['pos'][1], "z": 0.0}}
                    del (G.node[node]['pos'])
            writer = nx.readwrite.gexf.GEXFWriter(version='1.2draft')
            writer.add_graph(G)
            self.write_message({'type': 'download_gexf',
                'filename': self.config['name'] + '_trial_' + str(msg['data']),
                'data': tostring(writer.xml).decode(writer.encoding) })

        elif msg['type'] == 'download_json':
            G = self.simulation[ int(msg['data']) ].history_to_graph()
            for node in G.nodes():
                if 'pos' in G.node[node]:
                    G.node[node]['viz'] = {"position": {"x": G.node[node]['pos'][0], "y": G.node[node]['pos'][1], "z": 0.0}}
                    del (G.node[node]['pos'])
            self.write_message({'type': 'download_json',
                'filename': self.config['name'] + '_trial_' + str(msg['data']),
                'data': nx.node_link_data(G) })

        else:
            if self.application.verbose:
                logger.info('Unexpected message!')

    def update_logging(self):
        try:
            if (not self.log_capture_string.closed and self.log_capture_string.getvalue()):
                for i in range(len(self.log_capture_string.getvalue().split('\n')) - 1):
                    self.send_log('INFO.' + self.application.simulator.name, self.log_capture_string.getvalue().split('\n')[i])
                self.log_capture_string.truncate(0)
                self.log_capture_string.seek(0)
        finally:
            if self.capture_logging:
                thread = threading.Timer(0.01, self.update_logging)
                thread.start()

    def on_close(self):
        if self.application.verbose:
            logger.info('Socket closed!')

    def send_log(self, logger, logging):
        self.write_message({'type': 'log',
            'logger': logger,
            'logging': logging })

    def run_simulation(self):
        # Run simulation and capture logs
        if 'visualization_params' in self.config:
            del self.config['visualization_params']
        with self.logging(self.application.simulator.name):
            try:
                self.simulation = self.application.simulator.run(self.config)
                trials = []
                for i in range(self.config['num_trials']):
                    trials.append('{}_trial_{}'.format(self.name, i))
                self.write_message({'type': 'trials',
                    'data': trials })
            except:
                error = 'Something went wrong. Please, try again.'
                self.write_message({'type': 'error',
                    'error': error})
                self.send_log('ERROR.' + self.application.simulator.name, error)

    def get_trial(self, trial):
        G = self.simulation[trial].history_to_graph()
        return nx.node_link_data(G)

    @contextmanager
    def logging(self, logger):
        self.capture_logging = True
        self.logger_application = logging.getLogger(logger)
        self.log_capture_string = io.StringIO()
        ch = logging.StreamHandler(self.log_capture_string)
        self.logger_application.addHandler(ch)
        self.update_logging()
        yield self.capture_logging

        sleep(0.2)
        self.log_capture_string.close()
        self.logger_application.removeHandler(ch)
        self.capture_logging = False
        return self.capture_logging
    

class ModularServer(tornado.web.Application):
    """ Main visualization application. """

    port = 8001
    page_handler = (r'/', PageHandler)
    socket_handler = (r'/ws', SocketHandler)
    static_handler = (r'/(.*)', tornado.web.StaticFileHandler,
                      {'path': 'templates'})
    local_handler = (r'/local/(.*)', tornado.web.StaticFileHandler,
                     {'path': ''})

    handlers = [page_handler, socket_handler, static_handler, local_handler]
    settings = {'debug': True,
                'template_path': os.path.dirname(__file__) + '/templates'}

    def __init__(self, simulator, name='SOIL', verbose=True, *args, **kwargs):
        
        self.verbose = verbose
        self.name = name
        self.simulator = simulator

        # Initializing the application itself:
        super().__init__(self.handlers, **self.settings)

    def launch(self, port=None):
        """ Run the app. """
        
        if port is not None:
            self.port = port
        url = 'http://127.0.0.1:{PORT}'.format(PORT=self.port)
        print('Interface starting at {url}'.format(url=url))
        self.listen(self.port)
        # webbrowser.open(url)
        tornado.ioloop.IOLoop.instance().start()