# websocket_server -- WebSocket/HTTP server/client library
# https://github.com/CylonicRaider/websocket-server

"""
Support for "sessions" spanning multiple WebSocket connections.

In some (or many) cases, multiple WebSocket connections to the same URL are
(mostly) interchangeable. This module caters for that case by providing a
WebSocketSession class that transparently handles re-connecting to a WebSocket
endpoint, retrying commands whose results may have been missed, etc.

Event handling note: Some classes in this module allow their users to install
callbacks for certain "events" by setting instance attributes. An event EVT
is handled by a method called _on_EVT(); the latter invokes an instance
attribute called on_EVT (without a leading underscore) unless the attribute is
None. The externally specified callback receives the same arguments as the
method (aside from "self"). When overriding the event handler method, it is
strongly advisable to invoke the parent class' method (unless noted
otherwise). Each class lists the names of events for which this convention
applies, and allows setting them by specifying same-named keyword-only
constructor arguments; unless otherwise noted, the handler callbacks default
to None.
"""

import sys, time
import collections, uuid
import traceback
import threading

from . import client
from .constants import OP_TEXT, OP_BINARY
from .exceptions import ProtocolError, ConnectionClosedError
from .tools import spawn_daemon_thread, Scheduler, Future, EOFQueue

__all__ = ['RWST_IDLE', 'RWST_CONNECTING', 'RWST_CONNECTED',
           'RWST_INTERRUPTED', 'RWST_DISCONNECTING', 'RWST_DISCONNECTED',
           'SST_DISCONNECTED', 'SST_LOGGING_IN', 'SST_CONNECTED',
           'CST_NEW', 'CST_SENDING', 'CST_SENT', 'CST_SEND_FAILED',
           'CST_CONFIRMED', 'CST_CANCELLED',
           'ERRS_RTHREAD', 'ERRS_WS_CONNECT', 'ERRS_WS_RECV', 'ERRS_WS_SEND',
           'ERRS_SCHEDULER', 'ERRS_SERIALIZE', 'ERRS_DESERIALIZE',
           'ERRS_CALLBACK',
           'report_error',
           'backoff_constant', 'backoff_linear', 'backoff_exponential',
           'ReconnectingWebSocket', 'WebSocketSession']

RWST_IDLE          = 'IDLE'          # Fully inactive; no connection.
RWST_CONNECTING    = 'CONNECTING'    # Trying to establish a connection.
RWST_CONNECTED     = 'CONNECTED'     # Fully connected.
RWST_INTERRUPTED   = 'INTERRUPTED'   # Partially disconnected (still reading).
RWST_DISCONNECTING = 'DISCONNECTING' # Connection is being closed.
RWST_DISCONNECTED  = 'DISCONNECTED'  # Not connected; might reconnect soon.

SST_DISCONNECTED = 'DISCONNECTED' # Underlying connection not fully in place.
SST_LOGGING_IN   = 'LOGGING_IN'   # Executing early commands.
SST_CONNECTED    = 'CONNECTED'    # Connection fully established.

CST_NEW         = 'NEW'         # Command has never been sent.
CST_SENDING     = 'SENDING'     # Command is being sent.
CST_SENT        = 'SENT'        # Command has been sent but not responded to.
CST_SEND_FAILED = 'SEND_FAILED' # Sending the command failed.
CST_CONFIRMED   = 'CONFIRMED'   # Command has been responded to.
CST_CANCELLED   = 'CANCELLED'   # Command has been cancelled.

ERRS_RTHREAD     = 'RTHREAD'     # Fatal uncaught error in reader thread.
ERRS_WS_CONNECT  = 'WS_CONNECT'  # Exception raised while connecting.
ERRS_WS_RECV     = 'WS_RECV'     # Exception raised while reading from WS.
ERRS_WS_SEND     = 'WS_SEND'     # Exception while writing to or closing WS.
ERRS_SCHEDULER   = 'SCHEDULER'   # Uncaught error in Scheduler task.
ERRS_SERIALIZE   = 'SERIALIZE'   # Error while serializing a Command.
ERRS_DESERIALIZE = 'DESERIALIZE' # Error while deserializing an Event.
ERRS_CALLBACK    = 'CALLBACK'    # Uncaught error in event handler callback.

SWALLOW_CLASSES_CONNECT = (IOError, ProtocolError)
SWALLOW_CLASSES_RECV    = (IOError,)
SWALLOW_CLASSES_SEND    = (IOError, ProtocolError, ConnectionClosedError)

def run_cb(_func, *_args, **_kwds):
    """
    run_cb(func, *args, **kwds) -> any or None

    If func is not None, return the result of calling it with the given
    arguments; otherwise, do nothing and return None.
    """
    if _func is not None: return _func(*_args, **_kwds)

def report_error(exc, source, swallow, output=None):
    """
    report_error(exc, source, swallow, output=None) -> None

    Write a report about the given error to the given output stream. exc is
    the exception being handled; source indicates the origin of the error
    (in particular an ERRS_* constant); swallow indicates whether the report
    should be short (swallow is true) or extensive (swallow is false); output
    is a file-like object to write the report to, defaulting to standard error
    if None.

    This consults sys.exc_info() in order to retrieve a stack trace, but also
    works if that is not available. In order to avoid interleaving, this
    issues a single call to the output's write() method and one to its flush()
    method.
    """
    if output is None: output = sys.stderr
    buf = ['ERROR at %r:' % (source,)]
    info = sys.exc_info()
    if info[1] is not exc or swallow:
        buf.append(' %r' % (exc,))
    buf.append('\n')
    if info[1] is not None and not swallow:
        buf.append('Handler traceback (likely to provide useful context):\n')
        buf.extend(traceback.format_stack())
        buf.extend(traceback.format_exception(*info))
    output.write(''.join(buf))
    output.flush()

def backoff_constant(n):
    """
    backoff_constant(n) -> float

    Constant backoff implementation. This returns the constant 1.

    See ReconnectingWebSocket for details.
    """
    return 1

def backoff_linear(n):
    """
    backoff_linear(n) -> float

    Linear backoff implementation. This returns n.

    See ReconnectingWebSocket for details.
    """
    return n

def backoff_exponential(n):
    """
    backoff_exponential(n) -> float

    Exponential backoff implementation. This returns 2 ** n.

    See ReconnectingWebSocket for details.
    """
    return 2 ** n

# Here be concurrency dragons.

class ReconnectingWebSocket(object):
    """
    ReconnectingWebSocket(url, protos=None, cookies=None, ssl_config=None,
                          on_*=None) -> new instance

    An automatically reconnecting WebSocket wrapper. url is the WebSocket
    URL to connect to. protos is an indication on which WebSocket subprotocols
    may be used. cookies is a cookies.CookieJar instance that is used to
    store/retrieve cookies, or None for no cookie management. ssl_config is
    used for custom SSL configuration. on_* allow setting event handler
    callbacks; see below.

    Instance attributes are:
    url       : The URL to connect to. Initialized from the same-named
                constructor parameter. May be modified after instance creation
                to cause future connections to use that URL (but see
                reconnect() for a safer way of achieving that).
    protos    : Which WebSocket subprotocols may be used. May be None, a
                single string, or a list of strings. See client.connect() for
                more details. Initialized from the same-named constructor
                parameter.
    cookies   : A cookies.CookieJar instance (or anything with a compatible
                interface; see client.connect()) to use for cookie management,
                or None for no cookie management. Initialized from the
                same-named constructor parameter.
    ssl_config: A dictionary of SSL configuration as used by client.connect(),
                or None for default settings. Initialized from the same-named
                constructor parameter.
    backoff   : The connection backoff algorithm to use. Defaults to
                backoff_linear(). This is a function mapping an integer to a
                floating value; the parameter is the (zero-based) index of the
                current connection attempt (that has failed), while the return
                value is the time (in seconds) to wait until the next
                connection attempt. The index is reset when a connection
                attempt succeeds. The backoff_*() module-level functions
                provide a few ready-to-use implementations to plug into this.

    Read-only instance attributes are:
    state     : The current connection state as one of the RWST_* constants.
                Reading this attribute is not particularly useful as it might
                be changed by another thread immediately afterwards.
    state_goal: The state this ReconnectingWebSocket is trying to achieve,
                either RWST_DISCONNECTED or RWST_CONNECTED.
    conn      : The current WebSocket connection (if any) as a WebSocketFile.

    Class attributes (overridable on instances) are:
    USE_WTHREAD: Whether a writer thread should be created when connecting
                 (see below). This setting only takes effect when a writer is
                 about to be *created*; aside from that, a writer thread is
                 used if-and-only-if it is present.

    When active, a ReconnectingWebSocket uses one or two background threads.
    The *reader thread* is responsible for establishing the underlying
    connection, reading from it, and tearing it down; it is expected to spend
    most of its time waiting for the next incoming message. The optional
    *writer thread* is responsible for writing messages into the underlying
    connection and initiating shutdowns of the connection (when those are
    requested locally); if the writer thread is not used, the threads that
    request writing (or closing) operations perform them themselves,
    potentially blocking. The threads are created (automatically) whenever a
    connection is initiated (as long as they are not there); when a disconnect
    is requested, they clean themselves up after the connection is fully
    closed.

    During its life cycle, a ReconnectingWebSocket transitions between
    multiple states, which are described below (unless noted otherwise, a
    state's successor is the state described just below it):
    IDLE         : The initial state. There is no connection and none has been
                   requested.
    CONNECTING   : A connection is being established (at least, attempts are
                   made).
    CONNECTED    : A connection has been established. Two-way traffic can
                   happen. If a disconnect is requested while the connection
                   was being established, this state is transient and is
                   replaced by DISCONNECTING immediately. The next state is
                   either INTERRUPTED or DISCONNECTING.
    INTERRUPTED  : A disconnect has been requested locally. Messages can no
                   longer be sent. Waiting for the other side to confirm the
                   disconnect.
    DISCONNECTING: Received a disconnect request from the other side; closing
                   the underlying connection fully. No messages can be sent or
                   received.
    DISCONNECTED : The background worker threads are active, but there is no
                   connection. The next state is either CONNECTING or IDLE.
    Each state corresponds to a module-level RWST_* constant (see also the
    "state" attribute above).

    This class emits the following events (see the module docstring):
    connecting   : Emitted just after entering the CONNECTING state, or when a
                   connection attempt has failed and another is made (note
                   that there is no paired "connected" event in this case).
    connected    : Emitted just after entering the CONNECTED state.
    message      : Emitted when a message is received from the underlying
                   connection.
    disconnecting: Emitted just before closing the underlying connection,
                   immediately after entering the INTERRUPTED or DISCONNECTING
                   state (whichever comes first). The event handler may try to
                   send messages, but responses might not arrive.
    disconnected : Emitted just after entering the DISCONNECTED state.
    error        : Emitted when an exception is caught. The exception's
                   traceback etc. may be inspected by the callback. The
                   default is the module-level report_error() function, which
                   writes a report to standard error; set this to None to
                   ignore (swallowable) errors entirely. See also the error
                   handling note below.

    Note that this is not a drop-in replacement for the WebSocketFile class.

    Error handling note: The _on_*() event handler methods and various
    callbacks are not shielded against errors in overridden implementations;
    exceptions raised in them may bring the connection into an inconsistent
    state.
    """

    USE_WTHREAD = True

    def __init__(self, url, protos=None, cookies=None, ssl_config=None,
                 on_connecting=None, on_connected=None, on_message=None,
                 on_disconnecting=None, on_disconnected=None,
                 on_error=report_error):
        """
        __init__(url, protos=None, cookies=None, on_*=None) -> None

        Instance initializer; see the class docstring for details.
        """
        self.url = url
        self.protos = protos
        self.cookies = cookies
        self.ssl_config = ssl_config
        self.backoff = backoff_linear
        self.on_connecting = on_connecting
        self.on_connected = on_connected
        self.on_message = on_message
        self.on_disconnecting = on_disconnecting
        self.on_disconnected = on_disconnected
        self.on_error = on_error
        self.state = RWST_IDLE
        self.state_goal = RWST_DISCONNECTED
        self.conn = None
        self._cond = threading.Condition()
        self._rthread = None
        self._wthread = None
        self._wqueue = None
        self._connected = None
        self._disconnected = None
        self._disconnect_ok = False

    def __enter__(self):
        "Context manager entry; internal."
        self._cond.__enter__()

    def __exit__(self, *args):
        "Context manager exit; internal."
        self._cond.__exit__(*args)

    def _sleep(self, timeout, check=None):
        """
        _sleep(timeout, check=None) -> None

        Wait for a specified time while some condition holds. timeout is the
        time to wait for in (potentially fractional) seconds; check is a
        function taking no arguments and returning whether to continue
        sleeping, defaulting to always sleeping on.

        The sleep is based on the internal condition variable and can be
        interrupted by notifying it.
        """
        if check is None: check = lambda: True
        deadline = time.time() + timeout
        with self:
            while check():
                now = time.time()
                if now >= deadline: break
                self._cond.wait(deadline - now)

    def _cycle_connstates(self, disconnect=False, connect=False, url=None,
                          disconnect_ok=False):
        """
        _cycle_connstates(disconnect=False, connect=False, url=None,
                          disconnect_ok=False) -> Future

        Internal method backing connect() etc. If disconnect is true, performs
        a disconnect; if connect is true, performs a connect. url, if not
        None, is stored in the same-named instance attribute. disconnect_ok
        is only meaningful when disconnect is specified, and tells whether the
        disconnect is regular (True) or due to some error (False); it is
        passed on to the corresponding life cycle event handlers. Returns a
        Future that will resolve (to some unspecified value) when the
        requested operations are done (which may be immediately).

        The "state_goal" instance attribute is set to RWST_CONNECTED if
        connect is true, otherwise to RWST_DISCONNECTED if disconnect is true,
        otherwise it is unmodified.
        """

        def disconnect_task():
            "Helper function responsible for shutting down the connection."
            self._on_disconnecting(disconnect_conn.id, (not connect),
                                   disconnect_ok)
            self._do_disconnect(disconnect_conn, True)

        def on_future_error(exc, source):
            "Handler for errors raised in Future callbacks."
            self._on_error(exc, ERRS_CALLBACK, True)

        disconnect_conn = None
        with self:
            # Update attributes; prepare return value.
            if connect:
                self.state_goal = RWST_CONNECTED
                if self.state == RWST_CONNECTED:
                    ret = Future.resolved(on_error=on_future_error)
                else:
                    if self._connected is None:
                        self._connected = Future(on_error=on_future_error)
                    ret = self._connected
            elif disconnect:
                self.state_goal = RWST_DISCONNECTED
                if self.state == RWST_DISCONNECTED:
                    ret = Future.resolved(on_error=on_future_error)
                else:
                    if self._disconnected is None:
                        self._disconnected = Future(on_error=on_future_error)
                    ret = self._disconnected
            else:
                ret = Future.resolved(on_error=on_future_error)
            if url is not None:
                self.url = url
            if disconnect:
                self._disconnect_ok = disconnect_ok
            # Disconnect if told to.
            if disconnect:
                if self.state == RWST_CONNECTING:
                    # Wake up potentially sleeping reader thread.
                    self._cond.notifyAll()
                elif self.state == RWST_CONNECTED:
                    # Ensure the reader thread will be eventually woken;
                    # mark that that is going to happen.
                    self.state = RWST_INTERRUPTED
                    disconnect_conn = self.conn
            # Connect if told to; this amounts to spawning the reader thread
            # and letting it do its work.
            if connect:
                if self._rthread is None:
                    self._rthread = spawn_daemon_thread(self._rthread_main)
                if self._wthread is None and self.USE_WTHREAD:
                    self._wqueue = EOFQueue()
                    self._wthread = spawn_daemon_thread(self._wthread_main,
                                                        self._wqueue)
        if disconnect_conn is not None:
            self._run_wthread(disconnect_task)
        return ret

    def _rthread_main(self):
        """
        _rthread_main() -> None

        Internal method: Outermost code of the background thread responsible
        for reading WebSocket messages.
        """

        def detach(do_disconnect=True):
            "Clean up instance state pointing at the existence of this thread"
            with self:
                if self._rthread is this_thread:
                    self.state = RWST_IDLE
                    self._rthread = None
                    self.conn = None
                    if self._wthread is not None:
                        self._wqueue.close()
                        self._wthread = None
                        self._wqueue = None
            if conn is not None and do_disconnect:
                self._do_disconnect(conn)

        def sleep_check():
            "Callback to check whether to continue a sleep."
            # Run in a "with self:" block.
            return (self.state_goal == RWST_CONNECTED)

        this_thread = threading.current_thread()
        conn = None
        is_reconnect = False
        conn_attempt = 0
        try:
            while 1:
                # Prepare for connecting, or detach.
                with self:
                    # Detach if requested.
                    if self.state_goal != RWST_CONNECTED:
                        detach(False)
                        break
                    self.state = RWST_CONNECTING
                    params = self._conn_params()
                # Connect.
                self._on_connecting(not is_reconnect)
                try:
                    conn = self._do_connect(**params)
                except Exception as exc:
                    swallow = isinstance(exc, SWALLOW_CLASSES_CONNECT)
                    self._on_error(exc, ERRS_WS_CONNECT, swallow)
                    self._sleep(self.backoff(conn_attempt), sleep_check)
                    conn_attempt += 1
                    continue
                else:
                    conn_attempt = 0
                # Done connecting.
                with self:
                    if self.state_goal == RWST_CONNECTED:
                        self._disconnect_ok = True
                        new_params = self._conn_params()
                        do_read = (new_params == params)
                    else:
                        do_read = False
                    self.state = RWST_CONNECTED
                    self.conn = conn
                    if self._connected is not None:
                        self._connected.set(conn.id)
                        self._connected = None
                self._on_connected(conn.id, (not is_reconnect), (not do_read))
                # Read messages (unless we should disconnect immediately).
                if do_read:
                    self._do_read_loop(conn)
                # Prepare for disconnecting.
                with self:
                    run_dc_hook = (self.state != RWST_INTERRUPTED)
                    is_reconnect = (self.state_goal == RWST_CONNECTED)
                    ok = self._disconnect_ok
                    self.state = RWST_DISCONNECTING
                    self.conn = None
                # Disconnect.
                # If the disconnect is caused by a protocol violation by the
                # other side, the underlying WebSocketFile has put itself into
                # a state where it does try to read a confirmation so that we
                # do not have to take particular care of this.
                if run_dc_hook:
                    self._on_disconnecting(conn.id, (not is_reconnect), ok)
                self._do_disconnect(conn)
                old_conn_id = conn.id
                conn = None
                # Done disconnecting.
                with self:
                    self.state = RWST_DISCONNECTED
                    if self._disconnected is not None:
                        self._disconnected.set(old_conn_id)
                        self._disconnected = None
                self._on_disconnected(old_conn_id, (not is_reconnect), ok)
        except Exception as exc:
            self._on_error(exc, ERRS_RTHREAD)
        finally:
            detach()

    def _wthread_main(self, queue):
        """
        _wthread_main(queue) -> None

        Internal method: Main function of the (optional) background thread
        responsible for writing to the underlying WebSocket.

        queue is an EOFQueue yielding callbacks to be called by this thread.
        The callbacks take no arguments and their return values are ignored.
        """
        def force_disconnect():
            "Force a reconnect and/or determine that we should detach."
            dc_conn = None
            with self:
                if self._wthread is not this_thread:
                    return False
                if self.state == RWST_CONNECTED:
                    self.state = RWST_INTERRUPTED
                    dc_conn = self.conn
                    dc_final = (self.state_goal != RWST_CONNECTED)
                    queue.clear()
            if dc_conn is not None:
                # ok is always False as this is not an orderly disconnect.
                self._on_disconnecting(dc_conn.id, dc_final, False)
                self._do_disconnect(dc_conn, True)
            return True
        this_thread = threading.current_thread()
        try:
            while 1:
                try:
                    cb = queue.get()
                except EOFError:
                    break
                try:
                    cb()
                except Exception as exc:
                    swallow = isinstance(exc, SWALLOW_CLASSES_SEND)
                    self._on_error(exc, ERRS_WS_SEND, swallow)
                    if not force_disconnect(): break
        finally:
            with self:
                if self._wthread is this_thread:
                    self._wthread = None
                    self._wqueue = None

    def _conn_params(self):
        """
        _conn_params() -> dict

        Gather the parameters necessary to establish a new connection and
        return them.

        The parameters are passed as keyword arguments to _do_connect();
        subclasses overriding this method must therefore override
        _do_connect() as well.

        Concurrency note: This method is called with the internal monitor lock
        held and should therefore finish quickly.
        """
        return {'url': self.url, 'protos': self.protos,
                'cookies': self.cookies, 'ssl_config': self.ssl_config}

    def _do_connect(self, url, protos, cookies, ssl_config):
        """
        _do_connect(url, protos, cookies, ssl_config) -> WebSocketFile

        Actually establish a WebSocket connection and return it. url is the
        WebSocket URL to connect to. protos is an indication on which
        subprotocols to use (see the same-named instance attribute for
        details). cookies is a cookies.CookieJar instance or None and used
        for cookie management; see the same-named instance attribute for
        details. ssl_config is a dictionary of SSL configuration; see the
        same-named instance attribute for details. Overriding methods may
        specify additional parameters.

        The return value is expected to be a valid WebSocketFile with an "id"
        attribute containing a comparable and hashable object that uniquely
        identifies the connection and is not the None singleton. This
        implementation uses a random (version 4) UUID.

        Concurrency note: This method is called with no locks held; attribute
        accesses should be protected via "with self:" as necessary. However,
        values that require a new connection when they change (such as the
        URL) should be retrieved in _conn_params() instead; see there for
        more details.
        """
        conn = client.connect(url, protos, cookies=cookies,
                              ssl_config=ssl_config)
        conn.id = uuid.uuid4()
        return conn

    def _do_read_loop(self, conn):
        """
        _do_read_loop(conn) -> None

        Repeatedly read frames from the given WebSocket connection and pass
        them into _on_message(); on EOF, return (without calling
        _on_message()).
        """
        while 1:
            try:
                frame = conn.read_frame()
            except Exception as exc:
                swallow = isinstance(exc, SWALLOW_CLASSES_RECV)
                self._on_error(exc, ERRS_WS_RECV, swallow)
                break
            if frame is None: break
            self._on_message(frame, conn.id)

    def _do_send(self, conn, data, before_cb, after_cb):
        """
        _do_send(conn, data, before_cb, after_cb) -> None

        Backend of the send_message() method. This (synchronously) submits
        data into the WebSocketFile conn using an appropriate frame type and
        invokes the given callbacks.
        """
        if before_cb is not None:
            before_cb()
        try:
            if isinstance(data, bytes):
                conn.write_binary_frame(data)
            else:
                conn.write_text_frame(data)
        finally:
            ok = (sys.exc_info()[0] is None)
            if after_cb is not None:
                after_cb(ok)

    def _do_disconnect(self, conn, asynchronous=False):
        """
        _do_disconnect(conn, asynchronous=False) -> None

        Actually disconnect the given WebSocket connection. conn is the
        WebSocket connection to close. asynchronous tells whether this close
        is initiated from the thread responsible for reading messages (False)
        or to interrupt a connection ongoing concurrently (True).

        Concurrency note: As this method may be called to interrupt an ongoing
        connection, it should be particularly cautious about thread safety.
        """
        # Fortunately, WebSocketFile's close() takes care of the threading
        # for us (even when called asynchronously).
        conn.close(wait=(not asynchronous))

    def _run_wthread(self, cb, check_state=None):
        """
        _run_wthread(cb, check_state=None) -> bool

        Schedule the given callback to be executed on the writer thread. cb
        is the callback to run. check_state, if not None, indicates that only
        the given state may be in effect while submitting the callback; if the
        state does not match, cb is not run. Returns True if the state check
        passed, or False otherwise.

        If there is no writer thread (but check_state is satisfied), cb is
        executed synchronously. If there *is* a writer thread, note that
        queued calls may be silently discarded in the event of a sudden
        disconnect.
        """
        with self:
            if check_state is not None and self.state != check_state:
                return False
            queue = self._wqueue
            if queue is not None:
                queue.put(cb)
        if queue is None:
            cb()
        return True

    def _on_connecting(self, initial):
        """
        _on_connecting(initial) -> None

        Event handler method invoked before a connection is established.
        initial tells whether this is *not* part of a reconnect. Since the
        connection is not established yet, no ID is provided.

        The default implementation invokes the corresponding callback; see the
        class docstring for details.
        """
        run_cb(self.on_connecting, initial)

    def _on_connected(self, connid, initial, transient):
        """
        _on_connected(connid, initial, transient) -> None

        Event handler method invoked when a connection is established.
        connid is the ID of the connection; initial tells whether this is
        *not* part of a reconnect; transient tells whether the connection will
        be closed again immediately because the parameters have changed in the
        meantime.

        The default implementation invokes the corresponding callback; see the
        class docstring for details.
        """
        run_cb(self.on_connected, connid, initial, transient)

    def _on_message(self, msg, connid):
        """
        _on_message(msg, connid) -> None

        Event handler method invoked when a WebSocket message arrives. msg is
        a wsfile.Message containing the data that were received. connid is the
        ID of the connection the message was received from.

        The default implementation invokes the corresponding callback; see the
        class docstring for details.
        """
        run_cb(self.on_message, msg, connid)

    def _on_disconnecting(self, connid, final, ok):
        """
        _on_disconnecting(connid, final, ok) -> None

        Event handler method invoked when a connection is about to be closed.
        connid is the ID of the connection; final tells whether this is *not*
        part of a reconnect; ok tells whether the disconnect was "regular"
        (True) rather than caused by some sort of error (False).

        The default implementation invokes the corresponding callback; see the
        class docstring for details.
        """
        run_cb(self.on_disconnecting, connid, final, ok)

    def _on_disconnected(self, connid, final, ok):
        """
        _on_disconnect(connid, final, ok) -> None

        Event handler method invoked when a connection has been closed.
        connid is the ID of the connection; final tells whether this is *not*
        part of a reconnect; ok tells whether the disconnect was "regular"
        (True) rather than caused by some sort of error (False).

        The default implementation invokes the corresponding callback; see the
        class docstring for details.
        """
        run_cb(self.on_disconnected, connid, final, ok)

    def _on_error(self, exc, source, swallow=False):
        """
        _on_error(exc, source, swallow=False) -> None

        Event handler method invoked when an error occurs somewhere. exc is
        the exception object (sys.exc_info() can be used to retrieve more
        information about the error); source is a ERRS_* constant indicating
        in which component the error originated; swallow tells whether the
        exception should be suppressed (rather than being re-raised).

        The default implementation does nothing aside from invoking the
        corresponding callback (see the class docstring for details) and
        re-raising the exception if told to.
        Other implementations could make more fine-grained decisions based on,
        e.g., the type of the exception and the source. See also the
        module-level run_cb() convenience function for invoking the callback.
        """
        run_cb(self.on_error, exc, source, swallow)
        if not swallow: raise

    def send_message(self, data, before_cb=None, after_cb=None):
        """
        send_message(data, before_cb=None, after_cb=None) -> None

        Send a WebSocket frame containing the given data. The type of frame is
        chosen automatically depending on whether data is a byte or Unicode
        string. Depending on whether there is a dedicated background writer
        thread, the send may or may not be asynchronous. before_cb is invoked
        (if not None and with no arguments) immediately before sending the
        message, after_cb (if not None, and with a single Boolean argument
        that tells whether sending the message was successful, i.e. whether
        there was *no* exception) immediately after.

        In order to send messages, this instance must be in the CONNECTED
        state; if it is not, a ConnectionClosedError is raised. If a writer
        thread is used, all sends and callbacks (on the same
        ReconnectingWebSocket) are serialized w.r.t. each other (in
        particular, a message's after_cb is invoked before the next one's
        before_cb); otherwise, the send blocks the calling thread and the
        actual sending might be delayed (after before_cb has already been
        called).
        """
        res = self._run_wthread(lambda: self._do_send(self.conn, data,
            before_cb, after_cb), RWST_CONNECTED)
        if not res:
            raise ConnectionClosedError('Cannot send to non-connected '
                'ReconnectingWebSocket')

    def connect_async(self, url=None):
        """
        connect_async(url=None) -> Future

        Bring this instance into the "connected" state, creating a WebSocket
        connection (which will be renewed if it breaks before the next
        disconnect_async() call) as necessary.

        url, if not None, replaces the same-named instance attribute and is
        used as the URL future connections (including the one initiated by
        this method, if any) go to. If a connect is going on concurrently,
        it may or may not be redone in order to apply the new URL.

        Concurrency note: connect_async() / reconnect_async() /
        disconnect_async() initiate (as their names indicate) asynchronous
        actions and return quickly. If multiple calls are submitted in rapid
        succession, the effects of individual calls may be superseded by
        others or elided. Eventually, the last call (i.e. the last to acquire
        the internal synchronization lock) will prevail. In order to wait for
        the connection to finish, call the returned Future's wait() method.
        """
        return self._cycle_connstates(connect=True, url=url)

    def reconnect_async(self, url=None, ok=True):
        """
        reconnect_async(url=None, ok=True) -> Future

        Bring this instance into the "connected" state, forcing a
        disconnect-connect cycle if it is already connected.

        See the notes for connect_async() and disconnect_async() for more
        details.
        """
        return self._cycle_connstates(disconnect=True, connect=True, url=url,
                                      disconnect_ok=ok)

    def disconnect_async(self, ok=True):
        """
        disconnect_async(ok=True) -> Future

        Bring this instance into the "disconnected" state, closing the
        internal WebSocket connection if necessary.

        ok tells whether the close is normal (True) or caused by some sort of
        error (False).

        See the notes for connect_async() (but in reverse) for more details.
        """
        return self._cycle_connstates(disconnect=True, disconnect_ok=ok)

    def close(self):
        """
        close() -> None

        Fully shut this ReconnectingWebSocket down.

        This is a blocking alias for disconnect_async().
        """
        self.disconnect_async().wait()

class WebSocketSession(object):
    """
    WebSocketSession(conn, scheduler=None, on_*) -> new instance

    A wrapper around a ReconnectingWebSocket providing high-level message
    submission and reception facilities. conn is a ReconnectingWebSocket (or
    an instance of a subclass, which could theoretically employ an entirely
    different transport) providing low-level message handling. scheduler is a
    Scheduler instance that event handlers are executed by. on_* allow setting
    event handler callbacks; see below.

    Messages sent into the WebSocket are called "commands" and are represented
    by the nested Command class; messages received from the WebSocket are
    called "events" (regardless of their relation to previous commands) and
    are represented by the nested Event class.

    Read-only instance attributes are:
    conn     : The connection wrapped by this WebSocketSession, as passed to
               the constructor.
    scheduler: The Scheduler responsible for running high-level callbacks.
    state    : The current state of the WebSocketSession as a SST_* constant.
    commands : An ordered mapping from ID-s to Command instances representing
               still-live commands.
    queue    : A collections.deque of Command-s pending being sent.

    Class attributes (overridable on subclasses) are:
    Connection: The default class for the enclosed connection manager.
                Defaults to ReconnectingWebSocket and should be a subclass of
                that (or something with a compatible interface).

    Multiple background threads are involved in maintaining a
    WebSocketSession: The underlying ReconnectingWebSocket contributes one
    or two (see that class' documentation for more details); WebSocketSession
    itself executes all external callbacks using a Scheduler. As default, all
    of these are created automatically (but see the Scheduler documentation
    for exceptions); those maintained by ReconnectingWebSocket are shut down
    automatically whenever the WebSocketSession is closed, while the Scheduler
    may or may not continue running (depending on whether it is used by
    anything else).

    During its lifetime, a WebSocketSession alternates between multiple
    states:
    DISCONNECTED: The underlying connection is not present or being set up or
                  torn down.
    LOGGING_IN  : The underlying connection is established, the
                  WebSocketSession implementation is performing additional
                  setup work on it. Events may be received during this phase,
                  but commands sent in response to them will be buffered until
                  the WebSocketSession enters the CONNECTED state.
    CONNECTED   : The underlying connection is fully present; commands may
                  be submitted and events received.

    This class emits the following "events" (see the module docstring):
    connected    : An underlying connection has been established.
    logged_in    : The session is fully set up; commands will be sent.
    event        : An Event not matching any known command has been received.
                   The event handler receives the Event object as the only
                   positional argument.
    disconnecting: The connection is being torn down. Differently to
                   ReconnectingWebSocket, commands may *not* be sent.
    error        : An exception occurred. The exception's traceback etc. may
                   be inspected. The default is the module-level
                   report_error() function, which writes a report to standard
                   error; set this to None to ignore (swallowable) errors
                   entirely.
    All event handlers (except the "error" one, which must be called from the
    block that caught the exception) are executed on the Scheduler this
    instance references. Errors in event handlers are caught, forwarded to
    _on_error(), and suppressed.
    """

    # Synchronization note: WebSocketSession's lock comes outside the
    # Scheduler's lock. Remark that the Scheduler never invokes callbacks with
    # its own lock held.

    Connection = ReconnectingWebSocket

    class Command(object):
        """
        Command(data, id=None, responses=0, resendable=False, on_*=None)
            -> new instance

        A message sent from a WebSocketSession. data is the payload of the
        command; id is an identifier of the command (for recognizing
        responses); responses is the amount of responses the command expects.
        on_* allow setting event handler callbacks; see below.

        Instance attributes (all initialized from the corresponding
        constructor parameters, if those exist) are:
        data      : The payload of the command as an arbitrary
                    application-specific object. See also the serialize()
                    method.
        id        : A unique identifier of the command as some hashable
                    object. The value None is special-cased to mean that this
                    command cannot be identified (and, thus, cannot receive
                    responses).
        responses : The amount of responses this command has yet to receive.
                    WebSocketSession retains commands until their responses
                    attribute drops to zero (in particular, commands
                    initialized with responses=0 are forgotten immediately
                    after submission). None is special-cased to mean
                    arbitrarily many (such a command is never discarded).
        resendable: Whether this command may be resent safely without causing
                    unintended effects. An example of a resendable command
                    might be a non-destructive query while a creation request
                    might be non-resendable. Specific commands may choose to
                    ignore this attribute and adjust themselves when
                    on_reconnect() is invoked instead.
        state     : The processing state of this command, as a CST_* constant.
                    Managed by the WebSocketSession owing this command.

        Note that command responses are represented as instances of the Event
        class (as there is only one class for incoming frames and labelling
        frames unrelated to any command as "responses" was deemed worse than
        the current naming scheme).

        During its lifetime, a Command transitions between the following
        states (unless noted otherwise, a state's successor is the state
        defined just below it):
        NEW        : The command may have been queued but was not sent yet.
        SENDING    : The command is being sent. The next state is either SENT
                     or SEND_FAILED.
        SENT       : The command was fully sent. Note that this does not imply
                     that the other side has fully *received* the command
                     (this can happen due to the OS' internal buffering). The
                     next state is CONFIRMED.
        SEND_FAILED: An exception occurred while sending the command, perhaps
                     an I/O error. The underlying ReconnectingWebSocket is
                     about to reconnect, and this command will be resent or
                     cancelled altogether. The command might still become
                     CONFIRMED (if, e.g., the other side received enough bits
                     of the command to determine its ID, or it guessed the
                     ID).
        CONFIRMED  : A response to the command has been received. When enough
                     responses arrive, the command is dropped and there is no
                     successor state.
        CANCELLED  : The command could not be resent (or sent in the first
                     place) and was discarded. There is no successor state.

        This class emits the following events (see the module docstring):
        response : A response associated to the command has arrived. The event
                   handler receives the Event object representing the response
                   as the only positional argument.
        cancelled: The command could not finish successfully (e.g. because the
                   command is not safe-to-resend and the connection was broken
                   at an infelicitous time). The event handler may resubmit
                   the command anyway, but should be aware that it might take
                   effect twice in that case.
        Additional _on_*() methods are provided but not mapped to explicit
        events. Concrete implementations might override them, potentially
        exposing more callback attributes. All event handler methods (and thus
        callbacks) are executed by the Scheduler of the enclosing
        WebSocketSession and are (assuming a corresponding configuration)
        serialized w.r.t. each other. Exceptions raised in event handlers are
        caught, passed to the enclosing WebSocketSession's _on_error() method,
        and suppressed.
        """

        def __init__(self, data, id=None, responses=0, resendable=False,
                     on_response=None, on_cancelled=None):
            """
            __init__(data, id=None, responses=0, resendable=False, on_*)
                -> None

            Instance initializer; see the class docstring for details.
            """
            self.data = data
            self.id = id
            self.responses = responses
            self.resendable = resendable
            self.on_response = on_response
            self.on_cancelled = on_cancelled
            self.state = CST_NEW
            self._expects_response = (id is not None and
                (responses is None or responses > 0))
            self._status_tracker = None

        def serialize(self):
            """
            serialize() -> str or bytes

            Convert this Command into an application-specific on-the-wire
            format and return that as a Unicode or byte string object.

            The default implementation returns the data attribute unchanged,
            violating the method's contract unless data is already a string
            of an appropriate type.
            """
            return self.data

        def is_response_ok(self, evt):
            """
            is_response_ok(evt) -> bool

            Test whether the given Event, which may be assumed to be a
            response to this Command, does *not* consistute some sort of
            "error" response, and return that.

            The default implementation returns the event's "is_ok" attribute
            without modification.
            """
            return evt.is_ok

        def _on_sending(self):
            """
            _on_sending() -> None

            Event handler method invoked when the command has started being
            sent.

            Executed on the scheduler thread; in particular, some bits of the
            command may have already left the process; however, this is
            guaranteed to be scheduled before any related _on_response()
            (unless the server guesses the command's ID). The default
            implementation does nothing.
            """
            pass

        def _on_sent(self, ok):
            """
            _on_sent(ok) -> None

            Event handler method invoked when the command has been sent (or it
            has been tried to send it). ok tells whether the send has actually
            been successful.

            Note that a server might respond to the command before it has been
            fully received (or transmitted); in particular, _on_response()
            might be invoked before this method (see also _on_sending()).

            Executed on the scheduler thread. The default implementation
            updates internal bookkeeping data.
            """
            if self._expects_response:
                return
            elif ok:
                self._status_tracker.set()
            else:
                self._status_tracker.cancel()

        def _on_response(self, evt):
            """
            _on_response(evt) -> None

            Event handler method invoked when an event with the same ID as
            this command (i.e. a response to this command) arrives.

            Executed on the scheduler thread. The default implementation
            invokes the corresponding callback and updates internal
            bookkeeping data.
            """
            try:
                run_cb(self.on_response, evt)
            finally:
                if self._expects_response:
                    self.responses -= 1
                    if self.responses <= 0:
                        self._expects_response = False
                        if self._status_tracker is not None:
                            self._status_tracker.set()

        def _on_disconnect(self, final, ok, safe):
            """
            _on_disconnect(final, ok, safe) -> bool

            Event handler method invoked when the underlying connection has
            been closed. final and ok tell whether the close is *not*
            (presumably) temporary and *not* caused by an error, respectively;
            safe tells whether the command is in a state where it may be
            resent safely (i.e. has not been (perhaps partially) sent without
            a response). The return value indicates whether the command is
            *not* to be cancelled as obsolete.

            Executed on the scheduler thread. The default implementation
            returns True if-and-only-if any of the resendable attribute or
            safe is true.
            """
            return self.resendable or safe

        def _on_connect(self, initial):
            """
            _on_connect(initial) -> bool

            Event handler method invoked when the underlying connection has
            been established (in particular after a disconnect). initial tells
            whether the connection is *not* part of a reconnect. The return
            value indicates whether the command is to be (re-)sent.

            Implementations may mutate the internal data in order to upate
            them for the new connection.

            Executed on the scheduler thread. The default implementation
            always returns True, relying on _on_disconnect() to filter out
            not-safe-to-resend commands.
            """
            return True

        def _on_cancelled(self):
            """
            _on_cancelled() -> None

            Event handler method invoked when the command cannot be completed.

            Aside from explicit cancellation, this might also occur, e.g.,
            when the underlying connection breaks after the command has
            started being sent but before a reply to the command arrives, and
            the command is not safe-to-resend (see the resendable attribute).

            Executed on the scheduler thread. The default implementation calls
            the corresponding callback and updates internal bookkeeping data.
            Other implementations might inspect the command's state and make
            additional decisions.
            """
            try:
                run_cb(self.on_cancelled)
            finally:
                if self._status_tracker is not None:
                    self._status_tracker.cancel()

    class Event(object):
        """
        Event(data, connid, id=None, is_ok=True) -> new instance

        A message received from the underlying connection of a
        WebSocketSession. data is the payload of the event; connid is the
        ID of the connection the event arrived from; id is an identifier for
        matching events to related commands; is_ok tells whether this Event
        does *not* represent some sort of error message.

        Instance attributes (all initialized from the corresponding
        constructor parameters) are:
        data  : The payload of the event as an arbitrary application-specific
                object. See also the deserialize() class method.
        connid: The ID of the connection the event was received from.
        id    : An identifier relating this event to some command. None is
                special-cased to mean that this event cannot be identified.
                See also the Command class for more details on reply matching.
        is_ok: Whether this Event is not some sort of error message.
        """

        @classmethod
        def deserialize_msg(cls, msg, connid):
            """
            deserialize_msg(msg, connid) -> new instance

            Convert the given WebSocket message (which may have been
            reassembled from multiple frames and is represented by a
            wsfile.Message instance) into an internal format.

            The default method validates the type of the message and forwards
            its contents to the deserialize() class method.
            """
            if msg.msgtype not in (OP_TEXT, OP_BINARY):
                raise ValueError('Unrecognized WebSocket message type %r in '
                    'Event deserialization' % (msg.msgtype,))
            return cls.deserialize(msg.content, connid)

        @classmethod
        def deserialize(cls, data, connid):
            """
            deserialize(data, connid) -> new instance

            Convert the given data from an application-specific on-the-wire
            format into an internal format.

            The default implementation returns an instance with data
            unchanged, connid passed on, and an id of None. See also the
            deserialize_msg() class method for a more low-level routine.
            """
            return cls(data, connid)

        def __init__(self, data, connid, id=None, is_ok=True):
            """
            __init__(data, connid, id=None, is_ok=True) -> None

            Instance initializer; see the class docstring for details.
            """
            self.data = data
            self.connid = connid
            self.id = id
            self.is_ok = is_ok

    @classmethod
    def create(cls, url, protos=None, cookies=None, ssl_config=None, **kwds):
        """
        create(url, protos=None, cookies=None, **kwds) -> new instance

        Create a new WebSocketSession along with a new enclosed
        ReconnectingWebSocket. url, protos, cookies, and ssl_config are
        forwarded to the ReconnectingWebSocket constructor; other keyword
        arguments are forwarded to the WebSocketSession constructor. The
        actually instantiated "nested" class is defined by the "Connection"
        class attribute; unless overridden, it defaults to
        ReconnectingWebSocket.
        """
        return cls(cls.Connection(url, protos, cookies=cookies,
                                  ssl_config=ssl_config), **kwds)

    def __init__(self, conn, scheduler=None, on_connected=None,
                 on_logged_in=None, on_event=None, on_disconnecting=None,
                 on_error=report_error):
        """
        __init__(conn, scheduler=None, on_*=None) -> None

        Instance initializer; see the class docstring for details.
        """
        if scheduler is None: scheduler = Scheduler()
        self.conn = conn
        self.scheduler = scheduler
        self.state = SST_DISCONNECTED
        self.on_connected = on_connected
        self.on_logged_in = on_logged_in
        self.on_event = on_event
        self.on_disconnecting = on_disconnecting
        self.on_error = on_error
        self.commands = collections.OrderedDict()
        self.queue = collections.deque()
        self._lock = threading.RLock()
        self._send_queued = False
        self._connection_id = None
        self._connection_transient = False
        self._outer_hold = scheduler.hold()
        self._inner_hold = scheduler.hold()
        self._early_queue = collections.deque()
        self._install_callbacks()

    def __enter__(self):
        "Context manager entry; internal."
        self._lock.__enter__()

    def __exit__(self, *args):
        "Context manager exit; internal."
        self._lock.__exit__(*args)

    def _install_callbacks(self):
        """
        _install_callbacks() -> None

        Wire up callbacks for interesting events from this instance's
        connection and scheduler.
        """
        conn, sched = self.conn, self.scheduler
        conn.on_connecting = self._on_connecting
        conn.on_connected = sched.wrap(self._on_connected)
        conn.on_message = sched.wrap(self._on_raw_message)
        conn.on_disconnecting = sched.wrap(self._on_disconnecting)
        conn.on_disconnected = self._on_disconnected
        conn.on_error = self._on_error
        sched.on_error = lambda exc: self._on_error(exc, ERRS_SCHEDULER, True)

    def _run_cb(self, func, *args):
        """
        _run_cb(func, *args) -> any or None

        Invoke the given function (in particular a Command callback method, or
        None to accept any arguments and do nothing) with the given arguments
        and pass on its return value, or return None if it raises an error
        (having called the _on_error() instance method in that case).
        """
        if func is None: return
        try:
            return func(*args)
        except Exception as exc:
            self._on_error(exc, ERRS_CALLBACK, True)

    def _run_cb_async(self, func, *args):
        """
        _run_cb_async(func, *args) -> None

        Schedule the given function to be called with the given arguments via
        _run_cb(); see there for more details.
        """
        self.scheduler.add_now(lambda: self._run_cb(func, *args))

    def _do_login(self, connid, initial):
        """
        _do_login(connid, initial) -> Future or None

        Perform final setup work on an already-established underlying
        connection. connid and initial are forwarded from _on_connected().
        Returns a Future that resolves whenever the setup work is done, or
        None to indicate that it is already done.

        As the method name suggests, the "setup work" might be, e.g.,
        authenticating to other side of the underlying connection. Separating
        out this stage allows subclasses to restore state before
        commands from previous connections (which might, e.g., be privileged)
        are resent.

        This method is executed on the scheduler thread. The default
        implementation does nothing.
        """
        return None

    def _on_connecting(self, initial):
        """
        _on_connecting(initial) -> None

        Event handler method invoked when a connection is being established.

        Executed *synchronously* in a background thread. The default
        implementation acquires a hold on the underlying Scheduler to prepare
        for submitting various things to it.
        """
        self._inner_hold.acquire()

    def _on_connected(self, connid, initial, transient):
        """
        _on_connected(connid, initial, transient) -> None

        Event handler method invoked when a connection has been established.

        Executed on the scheduler thread. See also _on_logged_in(). The
        default implementation invokes the corresponding callback, and, unless
        transient is true, sets the session state to LOGGING_IN and invokes
        the _do_login() method.
        """
        def login_callback(v):
            self.scheduler.add_now(lambda: self._on_logged_in(connid,
                                                              initial))
        self._run_cb(self.on_connected, connid, initial, transient)
        with self:
            self._connection_id = connid
            self._connection_transient = transient
            if transient: return
            self.state = SST_LOGGING_IN
        res = self._do_login(connid, initial)
        if res is None:
            login_callback(None)
        else:
            res.add_done_cb(login_callback)

    def _on_logged_in(self, connid, initial):
        """
        _on_logged_in(connid, initial) -> None

        Event handler method invoked when a connection has been fully
        established (see _do_login() for additional setup).

        Executed on the scheduler thread. In contrast to _on_connected(), this
        method is only invoked for non-transient connections, so there is no
        "transient" parameter. The default implementation cancels any "early"
        commands left behind, schedules some of the regular commands to
        be resent (see the Command class for details), and invokes the
        corresponding event handling callback.
        """
        with self:
            self.state = SST_CONNECTED
            cancellist = tuple(self._early_queue)
            self._early_queue.clear()
            checklist = tuple(self.commands.values())
        for cmd in cancellist:
            self._cancel_command(cmd, True)
        cancellist = frozenset(cancellist)
        sendlist = []
        for cmd in checklist:
            if cmd in cancellist:
                pass
            elif self._run_cb(cmd._on_connect, initial):
                sendlist.append(cmd)
            else:
                self._cancel_command(cmd, True)
        with self:
            already_pending = frozenset(self.queue)
            self.queue.extend(cmd for cmd in sendlist
                              if cmd not in already_pending)
        self._do_submit()
        self._run_cb(self.on_logged_in, connid, initial)

    def _on_raw_message(self, msg, connid):
        """
        _on_raw_message(msg, connid) -> None

        Event handler method invoked when a message is received.

        Executed on the scheduler thread. This implementation constructs an
        Event from msg, dispatches it to either a command with the same ID
        (see Command._on_response()) or this instance's _on_event() (if there
        is no such command), and garbage-collects the dispatched-to command
        (if any and as necessary).
        """
        try:
            evt = self.Event.deserialize_msg(msg, connid)
        except Exception as exc:
            self._on_error(exc, ERRS_DESERIALIZE, True)
            return
        with self:
            cmd = self.commands.get(evt.id)
            # A send that looks like it failed to us (thus CST_SEND_FAILED)
            # might manage to transfer enough information to identify the
            # command nonetheless.
            if cmd is not None and cmd.state in (CST_NEW, CST_SENDING,
                                                 CST_SENT, CST_SEND_FAILED):
                cmd.state = CST_CONFIRMED
        handler = self._on_event if cmd is None else cmd._on_response
        self._run_cb(handler, evt)

    def _on_event(self, evt):
        """
        _on_event(evt) -> None

        Event handler method invoked when an Event that does not match any
        known command is received.

        Executed on the scheduler thread. The default implementation does
        nothing (aside from calling the on_event attribute, if not None).
        """
        self._run_cb(self.on_event, evt)

    def _on_disconnecting(self, connid, final, ok):
        """
        _on_disconnecting(connid, final, ok) -> None

        Event handler method invoked when a connection is about to be shut
        down.

        Executed on the scheduler thread; in particular, the connection is
        no longer usable. The default implementation updates the session
        state, cancels any "early" commands left, invokes (if the connection
        being closed is not "transient" -- see _on_connected()) the
        _on_disconnect() method of any still-pending "regular" commands, and
        runs the corresponding event handling callback.
        """
        def can_resend(cmd):
            return cmd.state not in (CST_SENDING, CST_SENT, CST_SEND_FAILED)
        with self:
            self.state = SST_DISCONNECTED
            cancellist = tuple(self._early_queue)
            self._early_queue.clear()
            if self._connection_transient:
                runlist = None
            else:
                runlist = [(cmd, can_resend(cmd))
                           for cmd in self.commands.values()]
            self._send_queued = False
            self._connection_id = None
            self._connection_transient = False
        for cmd in cancellist:
            self._cancel_command(cmd, True)
        if runlist is not None:
            cancellist = frozenset(cancellist)
            for cmd, safe in runlist:
                if cmd in cancellist:
                    pass
                elif not self._run_cb(cmd._on_disconnect, final, ok, safe):
                    self._cancel_command(cmd, True)
        self._run_cb(self.on_disconnecting, connid, final, ok)

    def _on_disconnected(self, connid, final, ok):
        """
        _on_disconnected(connid, final, ok) -> None

        Event handler method invoked when a connection has been torn down.

        Executed synchronously in a background thread. The default
        implementation releases the hold on the Scheduler acquired in
        _on_connecting().
        """
        self._inner_hold.release()

    def _on_error(self, exc, source, swallow=False):
        """
        _on_error(exc, source, swallow=False) -> None

        Event handler method invoked (synchronously) when an error occurs. See
        ReconnectingWebSocket._on_error() for details.
        """
        # Not using self._run_cb() to avoid infinite recursion.
        run_cb(self.on_error, exc, source, swallow)
        if not swallow: raise

    def _on_command_sending(self, cmd):
        """
        _on_command_sending(cmd) -> None

        Event handler method invoked when the given command is about to be
        sent. Executed synchronously just before the actual send.

        The default implementation updates the command's state and schedules
        its _on_sending() method to be run.
        """
        with self:
            if cmd.state in (CST_NEW, CST_CONFIRMED, CST_SEND_FAILED):
                cmd.state = CST_SENDING
            self._run_cb_async(cmd._on_sending)

    def _on_command_sent(self, cmd, ok):
        """
        _on_command_sent(cmd, ok) -> None

        Event handler method invoked when the given command has been
        sent. ok tels whether the send was successful.

        Executed synchronously just after the send. The default implementation
        updates the command's state (and internal command queues), schedules
        the command's _on_sent() method to be run, and initiates flushing more
        queued commands (if any) into the underlying connection.
        """
        with self:
            if cmd.state in (CST_NEW, CST_SENDING):
                cmd.state = CST_SENT if ok else CST_SEND_FAILED
            self._run_cb_async(cmd._on_sent, ok)
            self._remove_command(cmd, False)
            self._send_queued = False
        self._do_submit()

    def _remove_command(self, cmd, forget=True):
        """
        _remove_command(cmd, forget=True) -> None

        Internal: Remove the given command from the various sending queues,
        and from the ID-to-command index if forget is true.

        Note that cmd is *not* entirely removed if it is in the *middle* of
        a sending queue.
        """
        with self:
            if forget:
                self.commands.pop(cmd.id, None)
            if self.queue and self.queue[0] is cmd:
                self.queue.popleft()
            elif self._early_queue and self._early_queue[0] is cmd:
                self._early_queue.popleft()

    def _cancel_command(self, cmd, is_on_scheduler=False):
        """
        _cancel_command(cmd, is_on_scheduler=False) -> None

        Cancel the given command. is_on_scheduler tells whether the invocation
        happens on the scheduler thread; this influences whether the command's
        callback is invoked synchronously or not.

        This transitions the command into the CANCELLED state, (partially)
        removes references to it from internal indexes, and calls the
        command's _on_cancelled() handler.
        """
        with self:
            run_callback = (cmd.state != CST_CANCELLED)
            cmd.state = CST_CANCELLED
            self._remove_command(cmd)
        if not run_callback:
            pass
        elif is_on_scheduler:
            self._run_cb(cmd._on_cancelled)
        else:
            self._run_cb_async(cmd._on_cancelled)

    def _do_submit(self):
        """
        _do_submit() -> None

        Internal method: Backend of submit().

        This tries to send the first message in the send queue appropriate for
        the current session state after discarding cancelled messages; the
        message is not removed from the queue until sending finishes.
        """
        with self:
            if self._send_queued or self.state == SST_DISCONNECTED: return
            conn = self.conn
            if self.state == SST_LOGGING_IN:
                queue = self._early_queue
            else:
                queue = self.queue
            while 1:
                if not queue: return
                cmd = queue[0]
                if cmd.state != CST_CANCELLED: break
                queue.popleft()
            self._send_queued = True
        try:
            data = cmd.serialize()
        except Exception as exc:
            self._on_error(exc, ERRS_SERIALIZE, True)
            self._cancel_command(cmd)
            with self:
                self._send_queued = False
            # We expect serialization failures to be rare (*exceptional*
            # indeed) and call ourselves recursively.
            self._do_submit()
        else:
            conn.send_message(data,
                              lambda: self._on_command_sending(cmd),
                              lambda ok: self._on_command_sent(cmd, ok))

    def _submit(self, cmd, early=None):
        """
        _submit(cmd, early=None) -> Future

        Extended version of submit() for use by subclasses. cmd is the Command
        to be (eventually) sent; early is either None to indicate that this is
        a "regular" command, or a connection ID indicating that this is an
        "early" command pertaining to setting up the named connection. Returns
        a Future that resolves or fails to resolve whenever the command is
        fully processed.

        While a connection is being set up, regular commands are queued until
        the setup is finished, and early commands bearing the connection's
        name are sent immediately. Early commands submitted outside of the
        setup phase of "their" connection result in a RuntimeError being
        raised.
        """
        def on_future_error(exc, source):
            self._on_error(exc, ERRS_CALLBACK, True)
        def remove_cb(value):
            with self:
                self.commands.pop(cmd.id, None)
        with self:
            if cmd._status_tracker is not None:
                raise RuntimeError('Cannot submit Command more than once')
            if early is not None and (self.state != SST_LOGGING_IN or
                                      early != self._connection_id):
                raise RuntimeError('Submitting an early Command at an '
                    'inappropriate time')
            cmd._status_tracker = Future(on_error=on_future_error)
            if cmd._expects_response:
                self.commands[cmd.id] = cmd
                cmd._status_tracker.add_done_cb(remove_cb)
            if early is None:
                self.queue.append(cmd)
            else:
                self._early_queue.append(cmd)
        self._do_submit()
        return cmd._status_tracker

    def submit(self, cmd):
        """
        submit(cmd) -> Future

        Queue the given Command instance for being sent into the underlying
        connection and register it as waiting for responses (as necessary).
        Returns a Future that resolves or fails to resolve whenever the
        Command is fully processed (i.e. is not waiting for any responses or
        has been cancelled).

        If the session is already fully established, this also attempts to
        actually submit the command (unless there already are other commands
        being sent). It is an error to to submit the same Command instance
        more than once.
        """
        return self._submit(cmd)

    def connect(self, wait=True, **kwds):
        """
        connect(wait=True, **kwds) -> Future

        Establish an underlying connection. Returns a Future that resolves
        whenever the connection is established. wait indicates whether this
        call should block until that is the case. Additional keyword arguments
        are forwarded to ReconnectingWebSocket.connect_async(); see there for
        details.
        """
        with self:
            self._outer_hold.acquire()
            ret = self.conn.connect_async(**kwds)
        if wait: ret.wait()
        return ret

    def reconnect(self, wait=True, **kwds):
        """
        reconnect(wait=True, **kwds) -> Future

        Re-establish the underlying connection. See connect() and
        ReconnectingWebSocket.reconnect_async() for details.
        """
        with self:
            self._outer_hold.acquire()
            ret = self.conn.reconnect_async(**kwds)
        if wait: ret.wait()
        return ret

    def disconnect(self, wait=True, **kwds):
        """
        disconnect(wait=True, **kwds) -> Future

        Close the underlying connection. See connect() and
        ReconnectingWebSocket.disconnect_async() for details.
        """
        with self:
            self._outer_hold.release()
            ret = self.conn.disconnect_async(**kwds)
        if wait: ret.wait()
        return ret

    def close(self):
        """
        close() -> None

        Fully bring this WebSocketSession down.

        This is a (blocking) alias for disconnect().
        """
        self.disconnect(wait=True)


'''
Train convolutional graph netwok with initial features from mol2v3c (1-vectors for non-drug)
Unsupervised - Loss function according to cui2vec distances

Calssification - classification by indication (for limited drugs)

'''
import pandas as pn
import networkx as nx
import numpy as np
#First, read graph and leave only drugs that we have mol2vec embeddings for.

#create lables (cui2vec distances)
cui2vec = pn.read_csv('/mnt/galiasn/cui2vec_pretrained/cui2vec_pretrained.csv')
cuisnames = cui2vec['Unnamed: 0']
cui2vec = cui2vec.rename(columns={'Unnamed: 0':'CUI'})

drugGraphDF = pn.read_csv('/mnt/galiasn/DrugPredicates/predicate_PSHU_RCPT_DSYN_10.csv')
drugGraphDF =drugGraphDF.append(pn.read_csv('/mnt/galiasn/DrugPredicates/predicate_PSHU_RCPT_DSYN_30.csv'))
drugGraphDF =  drugGraphDF.append(pn.read_csv('/mnt/galiasn/DrugPredicates/predicate_PSHU_RCPT_DSYN_50.csv'))

drugGraphDF = drugGraphDF.append(pn.read_csv('/mnt/galiasn/DrugPredicates/predicate_PSHU_RCPT_DSYN_70.csv'))
drugGraphDF = drugGraphDF.append(pn.read_csv('/mnt/galiasn/DrugPredicates/predicate_PSHU_RCPT_DSYN_100.csv'))
drugGraphDF = drugGraphDF.append(pn.read_csv('/mnt/galiasn/DrugPredicates/predicate_PSHU_RCPT_DSYN_over100.csv'))

#this graph hold all the semrep predicates for drugs, disesaes and receptors

drugGraphDF = drugGraphDF[[u'OBJECT_CUI', u'SUBJECT_CUI',
      u'SUBJECT_NAME', u'SUBJECT_SEMTYPE',
      u'OBJECT_NAME', u'OBJECT_SEMTYPE']]

#create cui mapping

cuimap1 = drugGraphDF[['SUBJECT_NAME','SUBJECT_CUI']]
cuimap2 = drugGraphDF[['OBJECT_NAME','OBJECT_CUI']]
cuimap1 = cuimap1.rename(columns={'SUBJECT_NAME':'NAME','SUBJECT_CUI':'CUI'})
cuimap2 = cuimap2.rename(columns={'OBJECT_NAME':'NAME','OBJECT_CUI':'CUI'})
cuimap1 = cuimap1.append(cuimap2)
cuimap1 = cuimap1.drop_duplicates()
cuimap1 = pn.merge(cuimap1,cui2vec,on='CUI',how='inner')

#Leave only drugs

drugGraphDF = drugGraphDF[drugGraphDF['SUBJECT_SEMTYPE']=='phsu']
drugGraphDF = drugGraphDF[drugGraphDF['OBJECT_SEMTYPE']=='phsu']

#leave only cuis that are in cui2vec
#cui2vec = cui2vec.rename(columns={'Unnamed: 0':'SUBJECT_CUI'})
#drugGraphDF = pn.merge(cui2vec[['SUBJECT_CUI']],drugGraphDF,on='SUBJECT_CUI',how='inner')
#cui2vec = cui2vec.rename(columns={'SUBJECT_CUI':'OBJECT_CUI'})
#drugGraphDF = pn.merge(cui2vec[['OBJECT_CUI']],drugGraphDF,on='OBJECT_CUI',how='inner')
drugGraphDF = drugGraphDF[drugGraphDF['SUBJECT_CUI'].isin(cui2vec['CUI'].unique())]
drugGraphDF = drugGraphDF[drugGraphDF['OBJECT_CUI'].isin(cui2vec['CUI'].unique())]

tmp = drugGraphDF[drugGraphDF['SUBJECT_SEMTYPE'].isin(['phsu'])]
tmp = tmp[tmp['SUBJECT_SEMTYPE'].isin(['phsu'])]
drugNames = set(tmp['OBJECT_NAME'].unique()).union(set(tmp['SUBJECT_NAME'].unique()))

#add self loops
selfLoopDF = pn.DataFrame()
selfLoopDF['SUBJECT_NAME'] = list(drugGraphDF['SUBJECT_NAME'].unique())
selfLoopDF['OBJECT_NAME'] = list(drugGraphDF['SUBJECT_NAME'].unique())
#dummies
selfLoopDF['OBJECT_CUI'] = 'C0'*len(selfLoopDF)
selfLoopDF['SUBJECT_CUI'] = 'C0'*len(selfLoopDF)
selfLoopDF['SUBJECT_SEMTYPE'] = 'phsu'*len(selfLoopDF)
selfLoopDF['OBJECT_SEMTYPE'] = 'phsu'*len(selfLoopDF)

drugGraphDF = drugGraphDF.append(selfLoopDF)
#drugGraphDF = selfLoopDF
drugGraphDF['weight'] = [1]*len(drugGraphDF)
drugGraphDFgb = drugGraphDF.groupby(['SUBJECT_NAME','OBJECT_NAME'],as_index=False).sum()

fullGraph = nx.from_pandas_edgelist(drugGraphDFgb,'SUBJECT_NAME','OBJECT_NAME')

#adjacency matrix
AdjMat = nx.adjacency_matrix(fullGraph)
#add identity matrix
#Idnt = np.identity(len(fullGraph.nodes))
#AdjMat_Idnt = AdjMat+Idnt
'''
Create features. The features for each node is the ml2vec embedding.
If the node does not have mol2vec embeddings the features are a vector of ones
'''
#drugStructure holds all the drugs that we have a chemichal structure for. Compute mol2vec for each of them

modelMol2Vec = word2vec.Word2Vec.load('/home/galiasn/DATA/MechanismBasedRepurposing/Data/model_300dim.pkl')

drugStructure = pn.read_csv("/home/galiasn/DATA/MechanismBasedRepurposing/Data/structure links.csv",',')

drugStructure['noSmiles'] = drugStructure['SMILES'].isnull()
drugStructure = drugStructure[drugStructure['noSmiles']==False]

#drugStructure['mol'] = drugStructure['SMILES'].apply(Chem.MolFromSmiles)
ind = 0
lst=[]
while ind<len(drugStructure):
    smiles = drugStructure['SMILES'].iloc[0]
    mol = Chem.MolFromSmiles(smiles)
    if type(mol) is Chem.rdchem.Mol:
        lst.append(mol)
    else:
        lst.append(float('Nan'))
    ind+=1
    print(ind)
drugStructure['mol']= lst
drugStructure['noSmiles'] = drugStructure['mol'].isnull()
drugStructure = drugStructure[drugStructure['noSmiles']==False]
drugStructure['mol-sentences'] = drugStructure.apply(lambda x: MolSentence(mol2alt_sentence(x['mol'], 1)), axis=1)
drugStructure['mol2vec'] = [x for x in sentences2vec(drugStructure['mol-sentences'], modelMol2Vec, unseen='UNK')]

def translateToVec(val):
    if type(val)==type(1):
        return [1]*300
    else:
        return val
nodecuisnodes = list(fullGraph.nodes())
featureDF = pn.DataFrame()
featureDF['Name'] = nodecuisnodes
featureDF = pn.merge(featureDF,drugStructure[['Name','mol2vec']],on='Name',how='left')
featureDF['mol2vec']=featureDF['mol2vec'].fillna(1)
featureDF['mol2vec'] = featureDF['mol2vec'].apply(translateToVec)

#lables
lablesDF = pn.DataFrame()
lablesDF['NAME'] = list(fullGraph.nodes())
lablesDF = pn.merge(lablesDF,cuimap1,on='NAME',how='inner')

lables = lablesDF.loc[:,'V1':'V500']

#nodecuis = set(tmpDrugs['OBJECT_CUI'].unique()).union(set(tmpDrugs['SUBJECT_CUI'].unique()))
#lables = cui2vec[cui2vec['OBJECT_CUI'].isin(nodecuis)]


from structure.Structure import *
from parserfunc import *

# penalized features: number of empty caches (penaltyCaches), number of requested videos that are not in caches (penaltyVideos)

# define glsEvaluation: evualtuates solution and penalizes certain features if verified
def glsEvaluation(data, sol, penaltyCaches, penaltyVideos, penaltyFactor):
    evaluation = data.evaluation(sol)

    evaluation += penaltyFactor * (sol.emptyCaches() * penaltyCaches + sol.notCachedVideos(data) * penaltyVideos)

    return evaluation

# define glsUtility: calculates utility of feature in given solution

def glsUtility(data, sol, penaltyCaches, penaltyVideos, costCaches, costVideos):
    utilityCaches = sol.emptyCaches() * (costCaches / (1 + penaltyCaches))
    utilityVideos = sol.notCachedVideos(
        data) * (costVideos / (1 + penaltyVideos))

    return utilityCaches, utilityVideos

def gls(data, stoppingIterations, penaltyFactor):

    t0 = time.perf_counter()

    currSolution = data.generateRandomSol()
    bestSolution = currSolution

    penaltyCaches = 0
    penaltyVideos = 0

    costCaches = 1000
    costVideos = 500

    currEval = data.evaluation(bestSolution)
    counter = 0
    lastEval = 0
    iter = 0

    while(1):

        # search for bestCandidate in neighbourhood according to glsEvaluation function
        neighbourhood = data.neighbourhood(currSolution)
        bestCandidate = neighbourhood[0]
        for candidate in neighbourhood:
            if glsEvaluation(data, candidate, penaltyCaches, penaltyVideos, penaltyFactor) > glsEvaluation(data, bestCandidate, penaltyCaches, penaltyVideos, penaltyFactor):
                bestCandidate = candidate

        # select better solution
        if glsEvaluation(data, bestCandidate, penaltyCaches, penaltyVideos, penaltyFactor) > glsEvaluation(data, currSolution, penaltyCaches, penaltyVideos, penaltyFactor):
            currSolution = bestCandidate
            if data.evaluation(currSolution) > data.evaluation(bestSolution):
                bestSolution = currSolution
        else:
            utilityCaches, utilityVideos = glsUtility(
                data, currSolution, penaltyCaches, penaltyVideos, costCaches, costVideos)
            if(utilityCaches > utilityVideos):
                penaltyCaches += 1
            else:
                penaltyVideos += 1

        currEval = data.evaluation(bestSolution)
        print("Iteration nº", iter, ": Best Solution ->", currEval)
        iter += 1

        # stopping
        if(lastEval == currEval):
            counter += 1
        else:
            counter = 0
            lastEval = currEval

        if(counter == stoppingIterations):
            print("No evolution detected after",
                  stoppingIterations, "iterations. Stopping")
            break

    t1 = time.perf_counter()
    return bestSolution, t1-t0


from django.shortcuts import render
from .models import *
from  ui_project.settings import WEST_BANK_URL, EAST_BANK_URL
import requests
from django.views.decorators.csrf import csrf_exempt
# Create your views here.
def homeView(request):
    if request.method == "GET":
        r = requests.get(WEST_BANK_URL + "/bank/availability")
        availability = r.json()
        context = {}
        tuitions = Tuition.objects.all()
        json_tuitions = [{"codigo": tuition.id, "semestre": tuition.semester, "monto": tuition.amount, "carrera": tuition.major} for tuition in tuitions]
        context.update({"tuitions": json_tuitions, "payment": availability["Endpoint de pago"], "balance": availability["Endpoint de balance"]})

        return render(request, "home.html", context=context)

def payTuitionView(request, tuition_id):
    if request.method == "GET":
        return render(request, "pago.html", context = {"id_matricula": tuition_id})
    elif request.method == "POST":
        print(request.POST)
        bank_url = WEST_BANK_URL if request.POST["method"] == "WEST BANK" else EAST_BANK_URL
        tuition = Tuition.objects.get(id = tuition_id)
        data = {
            "card_number": int(request.POST["card-number"]),
            "id_number": int(request.POST["cc"]),
            "amount": tuition.amount,
            "callback": "http://127.0.0.1:8000/callback",
            "transaction_number": tuition_id,
            "key": int(request.POST["key"])
        }
        print(data)
        r = requests.post(bank_url + "/bank/pay", json = data, data=data)
        print(r.json())
        tuition_payment = TuitionPayment(payer_name = request.POST["name"],
            tuition = tuition,
            email = request.POST["email"],
            concept = request.POST["concept"],
            sede = request.POST["sede"],
            card_number = int(request.POST["card-number"]))

        tuition_payment.save()
        return render(request, "pago.html", context = {"id_matricula": tuition_id, "mensaje": "pago en proceso"})

@csrf_exempt
def callbackView(request):
    print("received callback")
    if request.method == "POST":
        tuition = Tuition.objects.filter(id = request.POST["transaction_number"])
        tuition_payment = TuitionPayment.objects.get(tuition = tuition)
        tuition_payment.state = "pagado"
        tuition_payment.save()
        print('finished callback')

def checkBalanceView(request):
    if request.method == "GET":
        return render(request, "saldo.html", context = {})
    elif request.method == "POST":
        number = request.POST["input-numero"]
        r = requests.get(WEST_BANK_URL + "/bank/balance/" + str(number))
        if r.status_code == 200:
            context = r.json()
            return render(request, "saldo.html", context = r.json())
        else:
            r = requests.get(EAST_BANK_URL + "/bank/balance/" + str(number))
            if r.status_code == 200:
                return render(request, "saldo.html", context = r.json())
            return render(request, "saldo.html", context = {"error": True})


import argparse
import importlib
import pkgutil

from api_builder import configuration, subcommands

def get_module_parser(mod, modname, parents=[], add_help=True):
    """
    Returns an argument parser for the sub-command's CLI.

     :param mod: the sub-command's python module
     :param modnam: the string name of the python module
     :return: ArgumentParser
    """
    return argparse.ArgumentParser(
        usage=configuration.EXECUTABLE_NAME + ' ' + modname + ' [options]',
        description=mod.get_description(), parents=parents,
        add_help=add_help)

def get_application_parser(commands):
    """
    Builds an argument parser for the application's CLI.

    :param commands:
    :return: ArgumentParser
    """

    parser = argparse.ArgumentParser(
        description=configuration.APPLICATION_DESCRIPTION,
        usage =configuration.EXECUTABLE_NAME + ' [sub-command] [options]',
        add_help=False)

    parser.add_argument(
        'sub_command',
        choices=[name for name in commands],
        nargs="?")

    parser.add_argument("-h", "--help", action="store_true")
    return parser

def get_module(name):
    """
    Convenience method for importing a module (i.e. sub-command) from a string

    :param name: module name to import
    :return: the module object
    """
    return importlib.import_module("api_builder.subcommands." + name)

def main():

    sub_commands = [m.name for m in pkgutil.iter_modules(subcommands.__path__)]
    application_parser = get_application_parser(sub_commands)
    args = application_parser.parse_known_args()[0]

    # Print help message
    if not args.sub_command:
        application_parser.print_help()

    # Delegate to sub-command help message
    if args.sub_command and args.help:
        command = get_module(args.sub_command)
        module_parser = get_module_parser(command, args.sub_command)
        command.set_args(module_parser)
        module_parser.print_help()

    # Execute the sub-command
    if args.sub_command and not args.help:
        command = get_module(args.sub_command)
        module_parser = get_module_parser(command, args.sub_command, parents=[application_parser], add_help=False)
        command.set_args(module_parser)
        module_args = module_parser.parse_args()
        command.execute(module_args)


from .data import *
from torch.nn import Conv2d, Sequential, Module

def adapt_conv(conv: Conv2d, n_channels: int, pretrained: bool = False,
               init: Optional[Callable] = None, padding_mode: str = 'zeros'):
    '''Create a new layer that adapts `conv` to accept `n_channels` inputs.
       Copies existing weights if `pretrained` or initialises them with `init`.'''
    if conv.in_channels == n_channels: return conv  # No need to adapt
    args = {n: getattr(conv, n) for n in [
                       'kernel_size', 'stride', 'padding', 'dilation', 'groups']}
    bias = conv.bias is not None
    if 'padding_mode' in Conv2d.__constants__:  # Padding mode added in PyTorch 1.1
        args['padding_mode'] = ifnone(padding_mode, conv.padding_mode)
    new_conv = Conv2d(n_channels, conv.out_channels, bias=bias, **args)
    if pretrained:
        exp_shape = (conv.out_channels, conv.in_channels, *conv.kernel_size)
        assert conv.weight.shape == exp_shape, f"Unexpected weights shape, expected {exp_shape}, got {conv.weight.shape}."
        new_conv.weight.data[...] = conv.weight.data[:, 0:1, :, :]
        if bias: new_conv.bias.data = conv.bias.data
    elif init: init_default(new_conv, init)
    new_conv.to(conv.weight.device)
    return new_conv

def adapt_model(model: Union[Module, Sequential], n_channels: int, name: str = 'conv1',
                   pretrained: bool = False, init: Optional[Callable] = None, padding_mode: str = 'zeros'):
    '''Adapt a convolutional model to `n_channels` inputs and copy weights if `pretrained` or initialise with `init`.'''
    # Find direct parent of first conv layer. Could be either a Sequential or a custom Module (but not the Conv itself)
    while (isinstance(model, Sequential) and
           isinstance(model[0], (Sequential, Module)) and
           not isinstance(model[0], Conv2d)):
        model = model[0]
    if isinstance(model, Sequential) and isinstance(model[0], Conv2d):
        conv1 = model[0]
        def update(conv): model[0] = conv
    elif isinstance(model, Module) and hasattr(model, name):
        conv1 = getattr(model, name)
        update = partial(setattr, model, name)
    else: raise TypeError(f"Could not locate first convolution layer. If it is a named layer then pass it's name, otherwise use adapt_conv.")
    update(adapt_conv(conv1, n_channels, pretrained=pretrained,
           init=init, padding_mode=padding_mode))

# Thanks to github.com/thomasbrandon/ for the audio_learner, adapt_conv_and adapt_model functions

def audio_learner(data: AudioDataBunch, base_arch: Callable = models.resnet18, metrics=accuracy,
                  cut: Union[int, Callable] = None, pretrained: bool = False, lin_ftrs: Optional[Collection[int]] = None,
                  ps: Floats = 0.5, custom_head: Optional[nn.Module] = None, split_on: Optional[SplitFuncOrIdxList] = None,
                  bn_final: bool = False, init=nn.init.kaiming_normal_, concat_pool: bool = True,
                  padding_mode: str = 'zeros', **kwargs: Any) -> Learner:
    '''Create a learner to apply a CNN model to audio spectrograms.'''
    learn = cnn_learner(data, base_arch, cut=cut, metrics=metrics, pretrained=pretrained, lin_ftrs=lin_ftrs, ps=ps,
                        custom_head=custom_head, split_on=split_on, bn_final=bn_final, init=init,
                        concat_pool=concat_pool, **kwargs)
    channels = _calc_channels(data)
    adapt_model(learn.model, channels, pretrained=pretrained,
                init=init, padding_mode=padding_mode)
    learn.unfreeze()  # Model shouldn't be frozen, unlike vision
    return learn

# This will have to be updated in future for multichannel but is a quick fix for now

def _calc_channels(data: AudioDataBunch):
    channels = data.train_ds[0][0].nchannels * \
        3 if data.config.delta else data.train_ds[0][0].nchannels
    return channels

def audio_predict(learn, item: AudioItem):
    '''Applies preprocessing to an AudioItem before predicting its class'''
    if isinstance(item, AudioItem):
        item = item.path
    print("Debugging", item)
    al = AudioList([item], path=item, config=learn.data.x.config)
    ai = AudioList.open(al, item)
    return learn.predict(ai)


import requests

def nameAndStausDroplet(token, url):
    # get info about all droplets
    droplets = requests.get(
        url, headers={'Authorization': f'Bearer {token}'}).json()['droplets']
    listOfDroplets = []
    # loop through all droplets
    for item in range(len(droplets)):
        # create list with name and status of each droplet
        listOfDroplets.append(
            [droplets[item]['name'], droplets[item]['status']])
    return(listOfDroplets)


"""
	Make message details a bit easier to work with.

	ContextManagedMessage means that a socket can be given and, if set,
	the context management bits will automatically send the message on exit.

	For example, here's how the status update is sent:

		with kernel.iopub_broadcast('status', message) as update:
			update.content.execution_state = status

	AdHocObjects are what let the content of the message get set in an arbitrary way.
"""
logger = shared.tools.jupyter.logging.Logger()

__all__ = ['KernelMessagingMixin']

from shared.tools.jupyter.base import JupyterKernelBaseMixin
from shared.tools.jupyter.wire import WireMessage
from shared.tools.jupyter.zmq import SocketType, ZMsg
from shared.tools.jupyter.status import declare_busy, declare_idle

from uuid import uuid4

class ContextManagedMessage(WireMessage):

	def __init__(self, zMessage=None, key='', signature_scheme='sha256',
		# initial value overrides, useful as kwargs
		ids=None, header=None, parent_header=None,
		metadata=None, content=None, raw_data=None,

		topic_prefix='',       # can take the place of ids on broadcast
		topic_broadcast=False, # broadcast instead of target socket IDs
		socket=None,
		):
		super(ContextManagedMessage, self).__init__(zMessage, key, signature_scheme,
				ids, header, parent_header, metadata, content, raw_data,
			)

		self.target_socket = socket

		# for use in broadcast, prepends message type with this
		self.topic_prefix = topic_prefix
		self.topic_broadcast = topic_broadcast

	def _add_ids_to_zMessage(self, zMessage):
		"""Override so topics may be broadcast (like for IOPub)"""
		if self.topic_broadcast:
			zMessage.add(self.topic_prefix + self.header.msg_type)
		else:
			if self.ids:
				for entry in self.ids:
					zMessage.add(entry)
		zMessage.add(self._MESSAGE_SPLITTING_DELIMITER_KEY_BETWEEN_IDS_AND_MESSAGE_PROPER)

	def __enter__(self):
		return self

	def __exit__(self, exc_type, exc_val, exc_tb):
		if self.target_socket:
			self.send()

	def send(self, socket=None):
		"""Send message to target"""
		if socket is None:
			socket = self.target_socket
		assert socket

		zMessage = self.package()

		try:
			assert zMessage.send(socket), 'failed to send message on %r: %r' % (socket, zMessage,)
		except Exception as error:
			raise error
		finally:
			zMessage.destroy()

class KernelMessagingMixin(JupyterKernelBaseMixin):
	"""
	Consolidate the messaging bits for the kernel in one place.

	First are the bits for sending messages, and then the message
	handling bits.
	"""
	# SENDING

	def _new_message(self, msg_type, target_socket=None, origin_message=None, topic_broadcast=False):
		return ContextManagedMessage(
			key = self.key,
			header = {
				'date': self.now,
				'msg_id': str(uuid4()),
				'session': self.session_id,
				'username': self.username,
				'msg_type': msg_type,
				'version': WireMessage.WIRE_PROTOCOL_VERSION,
			},
			parent_header = origin_message.header if origin_message else None,
			# default to a direct reply
			ids = origin_message.ids if origin_message else None,
			# in case message is for broadcast on a topic (like IOPub)
			topic_prefix = 'kernel.%(kernel_id)s.' % self,
			topic_broadcast = topic_broadcast,
			socket = target_socket,
		)

	# most messages on IOPub will be broadcast based on the topic
	def iopub_broadcast(self, msg_type, origin_message=None):
		return self._new_message(msg_type, self.iopub_socket, origin_message, topic_broadcast=True)

	# ensure replies for messages are handled and sent back to the requester
	def iopub_message(self, msg_type, origin_message=None):
		return self._new_message(msg_type, self.iopub_socket, origin_message, topic_broadcast=False)

	def shell_message(self, msg_type, origin_message=None):
		return self._new_message(msg_type, self.shell_socket, origin_message)

	def control_message(self, msg_type, origin_message=None):
		return self._new_message(msg_type, self.control_socket, origin_message)

	def stdin_message(self, msg_type, origin_message=None):
		return self._new_message(msg_type, self.stdin_socket, origin_message)

	# RECIEVING

	def _handle_zmessage(self, role, socket):
		"""
		Consume a ZMQ message off the socket, if any, and then handle it based on the role.
		"""
		zMessage = ZMsg.recvMsg(socket, self.ZMQ_DONTWAIT)
		if zMessage is not None:
			message = WireMessage(zMessage,
					  key=self.key,
					  signature_scheme=self.signature_scheme
				)
			try:
				declare_busy(self, message)
				if self.ACTIVE_HANDLER_RELOAD:
					reload_function(self[role + '_handler'])(self, message)
				else:
					self[role + '_handler'](self, message)
				zMessage.destroy()
			finally:
				declare_idle(self, message)

	def _handle_zbytes(self, role, socket):
		"""
		Handle socket data as raw bytes.

		This is only used by the heartbeat role.
		"""
		payload = socket.recv(self.ZMQ_DONTWAIT)
		self[role + '_handler'](self, payload)