#!/usr/bin/env python3
# Copyright (c) 2015-2016 The Bitcoin Core developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.
"""Utilities for doing coverage analysis on the RPC interface.

Provides a way to track which RPC commands are exercised during
testing.
"""

import os

REFERENCE_FILENAME = 'rpc_interface.txt'

class AuthServiceProxyWrapper(object):
    """
    An object that wraps AuthServiceProxy to record specific RPC calls.

    """
    def __init__(self, auth_service_proxy_instance, coverage_logfile=None):
        """
        Kwargs:
            auth_service_proxy_instance (AuthServiceProxy): the instance
                being wrapped.
            coverage_logfile (str): if specified, write each service_name
                out to a file when called.

        """
        self.auth_service_proxy_instance = auth_service_proxy_instance
        self.coverage_logfile = coverage_logfile

    def __getattr__(self, *args, **kwargs):
        return_val = self.auth_service_proxy_instance.__getattr__(
            *args, **kwargs)

        return AuthServiceProxyWrapper(return_val, self.coverage_logfile)

    def __call__(self, *args, **kwargs):
        """
        Delegates to AuthServiceProxy, then writes the particular RPC method
        called to a file.

        """
        return_val = self.auth_service_proxy_instance.__call__(*args, **kwargs)
        rpc_method = self.auth_service_proxy_instance._service_name

        if self.coverage_logfile:
            with open(self.coverage_logfile, 'a+', encoding='utf8') as f:
                f.write("%s\n" % rpc_method)

        return return_val

    @property
    def url(self):
        return self.auth_service_proxy_instance.url

    def __truediv__(self, relative_uri):
        return AuthServiceProxyWrapper(self.auth_service_proxy_instance / relative_uri)

def get_filename(dirname, n_node):
    """
    Get a filename unique to the test process ID and node.

    This file will contain a list of RPC commands covered.
    """
    pid = str(os.getpid())
    return os.path.join(
        dirname, "coverage.pid%s.node%s.txt" % (pid, str(n_node)))

def write_all_rpc_commands(dirname, node):
    """
    Write out a list of all RPC functions available in `luk-cli` for
    coverage comparison. This will only happen once per coverage
    directory.

    Args:
        dirname (str): temporary test dir
        node (AuthServiceProxy): client

    Returns:
        bool. if the RPC interface file was written.

    """
    filename = os.path.join(dirname, REFERENCE_FILENAME)

    if os.path.isfile(filename):
        return False

    help_output = node.help().split('\n')
    commands = set()

    for line in help_output:
        line = line.strip()

        # Ignore blanks and headers
        if line and not line.startswith('='):
            commands.add("%s\n" % line.split()[0])

    with open(filename, 'w', encoding='utf8') as f:
        f.writelines(list(commands))

    return True


import sys

sys.path.append(os.path.join(os.path.dirname(__file__), '../../../libbeat/tests/system'))

from beat.beat import TestCase

class BaseTest(TestCase):

    @classmethod
    def setUpClass(self):
        self.beat_name = "heartbeat"
        self.beat_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../"))


# This module contains a synchronous implementation of a Channel Access client
# as three top-level functions: read, write, subscribe. They are comparatively
# simple and naive, with no caching or concurrency, and therefore less
# performant but more robust.
import getpass
import inspect
import logging
import selectors
import socket
import threading  # just to make callback processing thread-safe
import time
import weakref

import caproto as ca

from .._dbr import ChannelType, SubscriptionType, field_types, native_type
from .._utils import (CaprotoError, CaprotoTimeoutError, ErrorResponseReceived,
                      adapt_old_callback_signature, get_environment_variables,
                      safe_getsockname)
from .repeater import spawn_repeater

__all__ = ('read', 'write', 'subscribe', 'block', 'interrupt',
           'read_write_read')
logger = logging.getLogger('caproto.ctx')

# Make a dict to hold our tcp sockets.
sockets = {}
global_circuits = {}

_permission_to_block = []  # mutable state shared by block and interrupt

# Convenience functions that do both transport and caproto validation/ingest.
def send(circuit, command, pv_name=None):
    if pv_name is not None:
        tags = {'pv': pv_name}
    else:
        tags = None
    buffers_to_send = circuit.send(command, extra=tags)
    sockets[circuit].sendmsg(buffers_to_send)

def recv(circuit):
    bytes_received = sockets[circuit].recv(4096)
    commands, _ = circuit.recv(bytes_received)
    for c in commands:
        circuit.process_command(c)
    return commands

def search(pv_name, udp_sock, timeout, *, max_retries=2):
    # Set Broadcaster log level to match our logger.
    b = ca.Broadcaster(our_role=ca.CLIENT)
    b.client_address = safe_getsockname(udp_sock)

    # Send registration request to the repeater
    logger.debug('Registering with the Channel Access repeater.')
    bytes_to_send = b.send(ca.RepeaterRegisterRequest())

    env = get_environment_variables()
    repeater_port = env['EPICS_CA_REPEATER_PORT']

    client_address_list = ca.get_client_address_list()
    local_address = ca.get_local_address()

    try:
        udp_sock.sendto(bytes_to_send, (local_address, repeater_port))
    except OSError as exc:
        raise ca.CaprotoNetworkError(
            f"Failed to send to {local_address}:{repeater_port}") from exc

    logger.debug("Searching for %r....", pv_name)
    commands = (
        ca.VersionRequest(0, ca.DEFAULT_PROTOCOL_VERSION),
        ca.SearchRequest(pv_name, 0, ca.DEFAULT_PROTOCOL_VERSION))
    bytes_to_send = b.send(*commands)
    tags = {'role': 'CLIENT',
            'our_address': b.client_address,
            'direction': '--->>>'}

    def send_search():
        for dest in client_address_list:
            tags['their_address'] = dest
            b.log.debug(
                '%d commands %dB',
                len(commands), len(bytes_to_send), extra=tags)
            try:
                udp_sock.sendto(bytes_to_send, dest)
            except OSError as exc:
                host, port = dest
                raise ca.CaprotoNetworkError(f"Failed to send to {host}:{port}") from exc

    def check_timeout():
        nonlocal retry_at

        if time.monotonic() >= retry_at:
            send_search()
            retry_at = time.monotonic() + retry_timeout

        if time.monotonic() - t > timeout:
            raise CaprotoTimeoutError(f"Timed out while awaiting a response "
                                      f"from the search for {pv_name!r}. Search "
                                      f"requests were sent to this address list: "
                                      f"{ca.get_address_list()}.")

    # Initial search attempt
    send_search()

    # Await a search response, and keep track of registration status
    retry_timeout = timeout / max((max_retries, 1))
    t = time.monotonic()
    retry_at = t + retry_timeout

    try:
        orig_timeout = udp_sock.gettimeout()
        udp_sock.settimeout(retry_timeout)
        while True:
            try:
                bytes_received, address = udp_sock.recvfrom(ca.MAX_UDP_RECV)
            except socket.timeout:
                check_timeout()
                continue

            check_timeout()

            commands = b.recv(bytes_received, address)
            b.process_commands(commands)
            for command in commands:
                if isinstance(command, ca.SearchResponse) and command.cid == 0:
                    address = ca.extract_address(command)
                    logger.debug('Found %r at %s:%d', pv_name, *address)
                    return address
            else:
                # None of the commands we have seen are a reply to our request.
                # Receive more data.
                continue
    finally:
        udp_sock.settimeout(orig_timeout)

def make_channel(pv_name, udp_sock, priority, timeout):
    log = logging.LoggerAdapter(logging.getLogger('caproto.ch'), {'pv': pv_name})
    address = search(pv_name, udp_sock, timeout)
    try:
        circuit = global_circuits[(address, priority)]
    except KeyError:

        circuit = global_circuits[(address, priority)] = ca.VirtualCircuit(
            our_role=ca.CLIENT,
            address=address,
            priority=priority)

    chan = ca.ClientChannel(pv_name, circuit)
    new = False
    if chan.circuit not in sockets:
        new = True
        sockets[chan.circuit] = socket.create_connection(chan.circuit.address,
                                                         timeout)
        circuit.our_address = sockets[chan.circuit].getsockname()
    try:
        if new:
            # Initialize our new TCP-based CA connection with a VersionRequest.
            send(chan.circuit, ca.VersionRequest(
                priority=priority,
                version=ca.DEFAULT_PROTOCOL_VERSION),
                pv_name)
            send(chan.circuit, chan.host_name(socket.gethostname()))
            send(chan.circuit, chan.client_name(getpass.getuser()))
        send(chan.circuit, chan.create(), pv_name)
        t = time.monotonic()
        while True:
            try:
                commands = recv(chan.circuit)
                if time.monotonic() - t > timeout:
                    raise socket.timeout
            except socket.timeout:
                raise CaprotoTimeoutError("Timeout while awaiting channel "
                                          "creation.")
            tags = {'direction': '<<<---',
                    'our_address': chan.circuit.our_address,
                    'their_address': chan.circuit.address}
            for command in commands:
                if isinstance(command, ca.Message):
                    tags['bytesize'] = len(command)
                    logger.debug("%r", command, extra=tags)
                elif command is ca.DISCONNECTED:
                    raise CaprotoError('Disconnected during initialization')
            if chan.states[ca.CLIENT] is ca.CONNECTED:
                log.info("Channel connected.")
                break

    except BaseException:
        sockets[chan.circuit].close()
        del sockets[chan.circuit]
        del global_circuits[(chan.circuit.address, chan.circuit.priority)]
        raise
    return chan

def _read(chan, timeout, data_type, data_count, notify, force_int_enums):
    logger = chan.log
    logger.debug("Detected native data_type %r.", chan.native_data_type)
    ntype = native_type(chan.native_data_type)  # abundance of caution
    if ((ntype is ChannelType.ENUM) and
            (data_type is None) and (not force_int_enums)):
        logger.debug("Changing requested data_type to STRING.")
        data_type = ChannelType.STRING
    req = chan.read(data_type=data_type, data_count=data_count, notify=notify)
    send(chan.circuit, req, chan.name)
    t = time.monotonic()
    while True:
        try:
            commands = recv(chan.circuit)
        except socket.timeout:
            commands = []

        if time.monotonic() - t > timeout:
            raise CaprotoTimeoutError("Timeout while awaiting reading.")

        tags = {'direction': '<<<---',
                'our_address': chan.circuit.our_address,
                'their_address': chan.circuit.address}
        for command in commands:
            if isinstance(command, ca.Message):
                tags['bytesize'] = len(command)
                logger.debug("%r", command, extra=tags)
            if (isinstance(command, (ca.ReadResponse, ca.ReadNotifyResponse)) and
                    command.ioid == req.ioid):
                return command
            elif isinstance(command, ca.ErrorResponse):
                raise ErrorResponseReceived(command)
            elif command is ca.DISCONNECTED:
                raise CaprotoError('Disconnected while waiting for '
                                   'read response')

def read(pv_name, *, data_type=None, data_count=None, timeout=1, priority=0,
         notify=True, force_int_enums=False, repeater=True):
    """
    Read a Channel.

    Parameters
    ----------
    pv_name : str
        The PV name to read from
    data_type : {'native', 'status', 'time', 'graphic', 'control'} or ChannelType or int ID, optional
        Request specific data type or a class of data types, matched to the
        channel's native data type. Default is Channel's native data type.
    data_count : integer, optional
        Requested number of values. Default is the channel's native data
        count.
    timeout : float, optional
        Default is 1 second.
    priority : 0, optional
        Virtual Circuit priority. Default is 0, lowest. Highest is 99.
    notify : boolean, optional
        Send a ReadNotifyRequest instead of a ReadRequest. True by default.
    force_int_enums : boolean, optional
        Retrieve enums as integers. (Default is strings.)
    repeater : boolean, optional
        Spawn a Channel Access Repeater process if the port is available.
        True default, as the Channel Access spec stipulates that well-behaved
        clients should do this.

    Returns
    -------
    response : ReadResponse or ReadNotifyResponse

    Examples
    --------

    Get the value of a Channel named 'simple:A'.

    >>> read('simple:A').data
    array([1], dtype=int32)

    Request a richer Channel Access data type that includes the timestamp, and
    access the timestamp.

    >>> read('cat', data_type='time').metadata.timestmap
    1570622339.042392

    A convenience method is provided for access the timestamp as a Python
    datetime object.

    >>> read('cat' data_type='time').metadata.stamp.as_datetime()
    datetime.datetime(2019, 10, 9, 11, 58, 59, 42392)

    The requested data type may also been given as a specific Channel Access
    type

    >>> from caproto import ChannelType
    >>> read('cat', data_type=ChannelType.CTRL_FLOAT).metadata
    DBR_CTRL_FLOAT(
        status=<AlarmStatus.NO_ALARM: 0>,
        severity=<AlarmSeverity.NO_ALARM: 0>,
        upper_disp_limit=0.0,
        lower_disp_limit=0.0,
        upper_alarm_limit=0.0,
        upper_warning_limit=0.0,
        lower_warning_limit=0.0,
        lower_alarm_limit=0.0,
        upper_ctrl_limit=0.0,
        lower_ctrl_limit=0.0,
        precision=0,
        units=b'')

    or the corresponding integer identifer

    >>> read('cat', data_type=30).metadata
    DBR_CTRL_FLOAT(
    status=<AlarmStatus.NO_ALARM: 0>,
    severity=<AlarmSeverity.NO_ALARM: 0>,
    upper_disp_limit=0.0,
    lower_disp_limit=0.0,
    upper_alarm_limit=0.0,
    upper_warning_limit=0.0,
    lower_warning_limit=0.0,
    lower_alarm_limit=0.0,
    upper_ctrl_limit=0.0,
    lower_ctrl_limit=0.0,
    precision=0,
    units=b'')
    """
    if repeater:
        # As per the EPICS spec, a well-behaved client should start a
        # caproto-repeater that will continue running after it exits.
        spawn_repeater()
    udp_sock = ca.bcast_socket()
    # Must bind or getsocketname() will raise on Windows.
    # See https://github.com/caproto/caproto/issues/514.
    udp_sock.bind(('', 0))
    try:
        udp_sock.settimeout(timeout)
        chan = make_channel(pv_name, udp_sock, priority, timeout)
    finally:
        udp_sock.close()
    try:
        return _read(chan, timeout, data_type=data_type, data_count=data_count,
                     notify=notify, force_int_enums=force_int_enums)
    finally:
        try:
            if chan.states[ca.CLIENT] is ca.CONNECTED:
                send(chan.circuit, chan.clear(), chan.name)
        finally:
            sockets[chan.circuit].close()
            del sockets[chan.circuit]
            del global_circuits[(chan.circuit.address, chan.circuit.priority)]

def subscribe(pv_name, priority=0, data_type=None, data_count=None,
              low=0.0, high=0.0, to=0.0, mask=None):
    """
    Define a subscription.

    Parameters
    ----------
    pv_name : string
        The PV name to subscribe to
    priority : integer, optional
        Used by the server to triage subscription responses when under high
        load. 0 is lowest; 99 is highest.
    data_type : {'native', 'status', 'time', 'graphic', 'control'} or ChannelType or int ID, optional
        Request specific data type or a class of data types, matched to the
        channel's native data type. Default is Channel's native data type.
    data_count : integer, optional
        Requested number of values. Default is the channel's native data
        count, which can be checked in the Channel's attribute
        :attr:`native_data_count`.
    low, high, to : float, optional
        deprecated by Channel Access, not yet implemented by caproto
    mask :  SubscriptionType, optional
        Subscribe to selective updates.

    Examples
    --------

    Define a subscription on the ``random_walk:x`` PV.

    >>> sub = subscribe('random_walk:x')

    Add one or more user-defined callbacks to process responses.

    >>> def f(sub, response):
    ...     print(repsonse.data)
    ...
    >>> sub.add_callback(f)

    Activate the subscription and process incoming responses.

    >>> sub.block()

    This is a blocking operation in the sync client. (To do this on a
    background thread, use the threading client.) Interrupt using Ctrl+C or
    by calling :meth:`sub.interrupt()` from another thread.

    The subscription may be reactivated by calling ``sub.block()`` again.

    To process multiple subscriptions at once, use the *function*
    :func:`block`, which takes one or more Subscriptions as arguments.

    >>> block(sub1, sub2)

    There is also an :func:`interrupt` function, which is merely an alias to
    the method.
    """
    return Subscription(pv_name, priority, data_type, data_count, low, high,
                        to, mask)

def interrupt():
    """
    Signal to :func:`block` to stop blocking. Idempotent.

    This obviously cannot be called interactively while blocked;
    it is intended to be called from another thread.
    """
    _permission_to_block.clear()

def block(*subscriptions, duration=None, timeout=1, force_int_enums=False,
          repeater=True):
    """
    Activate one or more subscriptions and process incoming responses.

    Use Ctrl+C (SIGINT) to escape, or from another thread, call
    :func:`interrupt()`.

    Parameters
    ----------
    *subscriptions : Subscriptions
        The list of subscriptions.
    duration : float, optional
        How many seconds to run for. Run forever (None) by default.
    timeout : float, optional
        Default is 1 second. This is not the same as `for`; this is the timeout
        for failure in the event of no connection.
    force_int_enums : boolean, optional
        Retrieve enums as integers. (Default is strings.)
    repeater : boolean, optional
        Spawn a Channel Access Repeater process if the port is available.
        True default, as the Channel Access spec stipulates that well-behaved
        clients should do this.

    Examples
    --------

    Activate subscription(s) and block while they process updates.

    >>> sub1 = subscribe('cat')
    >>> sub1 = subscribe('dog')
    >>> block(sub1, sub2)
    """
    _permission_to_block.append(object())
    if duration is not None:
        deadline = time.time() + duration
    else:
        deadline = None
    if repeater:
        # As per the EPICS spec, a well-behaved client should start a
        # caproto-repeater that will continue running after it exits.
        spawn_repeater()
    loggers = {}
    for sub in subscriptions:
        loggers[sub.pv_name] = logging.LoggerAdapter(logging.getLogger('caproto.ch'),
                                                     {'pv': sub.pv_name})
    udp_sock = ca.bcast_socket()
    # Must bind or getsocketname() will raise on Windows.
    # See https://github.com/caproto/caproto/issues/514.
    udp_sock.bind(('', 0))
    try:
        udp_sock.settimeout(timeout)
        channels = {}
        for sub in subscriptions:
            pv_name = sub.pv_name
            chan = make_channel(pv_name, udp_sock, sub.priority, timeout)
            channels[sub] = chan
    finally:
        udp_sock.close()
    try:
        # Subscribe to all the channels.
        sub_ids = {}
        for sub, chan in channels.items():
            loggers[chan.name].debug("Detected native data_type %r.",
                                     chan.native_data_type)

            # abundance of caution
            ntype = field_types['native'][chan.native_data_type]
            if ((ntype is ChannelType.ENUM) and (not force_int_enums)):
                ntype = ChannelType.STRING
            time_type = field_types['time'][ntype]
            # Adjust the timeout during monitoring.
            sockets[chan.circuit].settimeout(None)
            loggers[chan.name].debug("Subscribing with data_type %r.",
                                     time_type)
            req = chan.subscribe(
                data_type=time_type, data_count=sub.data_count, mask=sub.mask)
            send(chan.circuit, req, chan.name)
            sub_ids[(chan.circuit, req.subscriptionid)] = sub
        logger.debug('Subscribed. Building socket selector.')
        try:
            circuits = set(chan.circuit for chan in channels.values())
            selector = selectors.DefaultSelector()
            sock_to_circuit = {}
            for circuit in circuits:
                sock = sockets[circuit]
                sock_to_circuit[sock] = circuit
                selector.register(sock, selectors.EVENT_READ)
            if duration is None:
                logger.debug('Continuing until SIGINT is received....')
            while True:
                events = selector.select(timeout=0.1)
                if deadline is not None and time.time() > deadline:
                    logger.debug('Deadline reached.')
                    return
                if not _permission_to_block:
                    logger.debug("Interrupted via "
                                 "caproto.sync.client.interrupt().")
                    break
                for selector_key, _ in events:
                    circuit = sock_to_circuit[selector_key.fileobj]
                    commands = recv(circuit)
                    for response in commands:
                        if isinstance(response, ca.ErrorResponse):
                            raise ErrorResponseReceived(response)
                        if response is ca.DISCONNECTED:
                            # TODO Re-connect.
                            raise CaprotoError("Disconnected")
                        sub = sub_ids.get((circuit, response.subscriptionid))
                        if sub:
                            sub.process(response)
        except KeyboardInterrupt:
            logger.debug('Received SIGINT. Closing.')
            pass
    finally:
        _permission_to_block.clear()
        try:
            for chan in channels.values():
                if chan.states[ca.CLIENT] is ca.CONNECTED:
                    send(chan.circuit, chan.clear(), chan.name)
        finally:
            # Reinstate the timeout for channel cleanup.
            for chan in channels.values():
                sockets[chan.circuit].settimeout(timeout)
                sockets[chan.circuit].close()
                del sockets[chan.circuit]
                del global_circuits[(chan.circuit.address, chan.circuit.priority)]

def _write(chan, data, metadata, timeout, data_type, notify):
    logger.debug("Detected native data_type %r.", chan.native_data_type)
    # abundance of caution
    ntype = field_types['native'][chan.native_data_type]
    if (data_type is None) and (ntype is ChannelType.ENUM):
        # Change data_type to STRING if data contains string-like data, or
        # iterable of string-like data
        stringy_data = False
        if isinstance(data, (str, bytes)):
            stringy_data = True
        if hasattr(data, '__getitem__') \
                and len(data) > 0 \
                and isinstance(data[0], (str, bytes)):
            stringy_data = True

        if stringy_data:
            logger.debug("Will write to ENUM as data_type STRING.")
            data_type = ChannelType.STRING
    logger.debug("Writing.")
    req = chan.write(data=data, notify=notify,
                     data_type=data_type, metadata=metadata)
    send(chan.circuit, req, chan.name)
    t = time.monotonic()
    if notify:
        while True:
            try:
                commands = recv(chan.circuit)
            except socket.timeout:
                commands = []

            if time.monotonic() - t > timeout:
                raise CaprotoTimeoutError("Timeout while awaiting write reply.")

            tags = {'direction': '<<<---',
                    'our_address': chan.circuit.our_address,
                    'their_address': chan.circuit.address}
            for command in commands:
                if isinstance(command, ca.Message):
                    tags['bytesize'] = len(command)
                    logger.debug("%r", command, extra=tags)
                if (isinstance(command, ca.WriteNotifyResponse) and
                        command.ioid == req.ioid):
                    response = command
                    break
                elif isinstance(command, ca.ErrorResponse):
                    raise ErrorResponseReceived(command)
                elif command is ca.DISCONNECTED:
                    raise CaprotoError('Disconnected while waiting for '
                                       'write response')
            else:
                continue
            break
        return response
    else:
        return None

def write(pv_name, data, *, notify=False, data_type=None, metadata=None,
          timeout=1, priority=0,
          repeater=True):
    """
    Write to a Channel.

    Parameters
    ----------
    pv_name : str
        The PV name to write to
    data : str, bytes, int, or float or any Iterable of these
        Value(s) to write.
    notify : boolean, optional
        Request notification of completion and wait for it. False by default.
    data_type : {'native', 'status', 'time', 'graphic', 'control'} or ChannelType or int ID, optional
        Write as specific data type. Default is inferred from input.
    metadata : ``ctypes.BigEndianStructure`` or tuple
        Status and control metadata for the values
    timeout : float, optional
        Default is 1 second.
    priority : 0, optional
        Virtual Circuit priority. Default is 0, lowest. Highest is 99.
    repeater : boolean, optional
        Spawn a Channel Access Repeater process if the port is available.
        True default, as the Channel Access spec stipulates that well-behaved
        clients should do this.

    Returns
    -------
    initial, final : tuple of ReadNotifyResponse objects

    Examples
    --------
    Write the value 5 to a Channel named 'simple:A'.

    >>> write('simple:A', 5)  # returns None

    Request notification of completion ("put completion") and wait for it.
    >>> write('cat', 5, notify=True)  # blocks until complete, then returns:
    WriteNotifyResponse(
    data_type=<ChannelType.LONG: 5>,
    data_count=1,
    status=CAStatusCode(
    name='ECA_NORMAL', code=0, code_with_severity=1,
    severity=<CASeverity.SUCCESS: 1>,
    success=1, defunct=False,
    description='Normal successful completion'),
    ioid=0)
    """
    if repeater:
        # As per the EPICS spec, a well-behaved client should start a
        # caproto-repeater that will continue running after it exits.
        spawn_repeater()

    udp_sock = ca.bcast_socket()
    # Must bind or getsocketname() will raise on Windows.
    # See https://github.com/caproto/caproto/issues/514.
    udp_sock.bind(('', 0))
    try:
        udp_sock.settimeout(timeout)
        chan = make_channel(pv_name, udp_sock, priority, timeout)
    finally:
        udp_sock.close()
    try:
        return _write(chan, data, metadata, timeout, data_type, notify)
    finally:
        try:
            if chan.states[ca.CLIENT] is ca.CONNECTED:
                send(chan.circuit, chan.clear(), chan.name)
        finally:
            sockets[chan.circuit].close()
            del sockets[chan.circuit]
            del global_circuits[(chan.circuit.address, chan.circuit.priority)]

def read_write_read(pv_name, data, *, notify=False,
                    read_data_type=None, write_data_type=None,
                    metadata=None, timeout=1, priority=0,
                    force_int_enums=False, repeater=True):
    """
    Write to a Channel, but sandwich the write between to reads.

    This is what the command-line utilities ``caproto-put`` and ``caput`` do.
    Notice that if you want the second reading to reflect the written value,
    you should pass the parameter ``notify=True``. (This is also true of
    ``caproto-put``/``caput``, which needs the ``-c`` argument to behave the
    way you might expect it to behave.)

    This is provided as a separate function in order to support ``caproto-put``
    efficiently. Making separate calls to :func:`read` and :func:`write` would
    re-create a connection redundantly.

    Parameters
    ----------
    pv_name : str
        The PV name to write/read/write
    data : str, bytes, int, or float or any Iterable of these
        Value to write.
    notify : boolean, optional
        Request notification of completion and wait for it. False by default.
    read_data_type : {'native', 'status', 'time', 'graphic', 'control'} or ChannelType or int ID, optional
        Request specific data type.
    write_data_type : {'native', 'status', 'time', 'graphic', 'control'} or ChannelType or int ID, optional
        Write as specific data type. Default is inferred from input.
    metadata : ``ctypes.BigEndianStructure`` or tuple
        Status and control metadata for the values
    timeout : float, optional
        Default is 1 second.
    priority : 0, optional
        Virtual Circuit priority. Default is 0, lowest. Highest is 99.
    force_int_enums : boolean, optional
        Retrieve enums as integers. (Default is strings.)
    repeater : boolean, optional
        Spawn a Channel Access Repeater process if the port is available.
        True default, as the Channel Access spec stipulates that well-behaved
        clients should do this.

    Returns
    -------
    initial, write_response, final : tuple of response

    The middle response comes from the write, and it will be ``None`` unless
    ``notify=True``.

    Examples
    --------

    Write the value 5 to a Channel named 'simple:A'.

    >>> read_write_read('cat', 5)  # returns initial, None, final

    Request notification of completion ("put completion") and wait for it.

    >>> read_write_read('cat', 5, notify=True)  # initial, WriteNotifyResponse, final
    """
    if repeater:
        # As per the EPICS spec, a well-behaved client should start a
        # caproto-repeater that will continue running after it exits.
        spawn_repeater()

    udp_sock = ca.bcast_socket()
    # Must bind or getsocketname() will raise on Windows.
    # See https://github.com/caproto/caproto/issues/514.
    udp_sock.bind(('', 0))
    try:
        udp_sock.settimeout(timeout)
        chan = make_channel(pv_name, udp_sock, priority, timeout)
    finally:
        udp_sock.close()
    try:
        initial = _read(chan, timeout, read_data_type, None, notify=True,
                        force_int_enums=force_int_enums)
        res = _write(chan, data, metadata, timeout, write_data_type, notify)
        final = _read(chan, timeout, read_data_type, None, notify=True,
                      force_int_enums=force_int_enums)
    finally:
        try:
            if chan.states[ca.CLIENT] is ca.CONNECTED:
                send(chan.circuit, chan.clear(), chan.name)
        finally:
            sockets[chan.circuit].close()
            del sockets[chan.circuit]
            del global_circuits[(chan.circuit.address, chan.circuit.priority)]
    return initial, res, final

class Subscription:
    """
    This object encapsulates state related to a Subscription.

    See the :func:`subscribe` function.
    """
    def __init__(self, pv_name, priority=0, data_type=None, data_count=None,
                 low=0.0, high=0.0, to=0.0, mask=None):
        if mask is None:
            mask = SubscriptionType.DBE_VALUE | SubscriptionType.DBE_ALARM
        self.pv_name = pv_name
        self.priority = priority
        self.data_type = data_type
        self.data_count = data_count
        self.low = low
        self.high = high
        self.to = to
        self.mask = mask

        self.callbacks = {}
        self._callback_id = 0
        self._callback_lock = threading.RLock()

        # This is related to back-compat for user callbacks that have the old
        # signature, f(response).
        self.__wrapper_weakrefs = set()

    def block(self, duration=None, timeout=1,
              force_int_enums=False,
              repeater=True):
        """
        Activate one or more subscriptions and process incoming responses.

        Use Ctrl+C (SIGINT) to escape, or from another thread, call
        :meth:`interrupt()`.

        Convenience alias for the top-level function :func:`block`, which may
        be used to process multiple Subscriptions concurrently.

        Parameters
        ----------

        duration : float, optional
            How many seconds to run for. Run forever (None) by default.
        timeout : float, optional
            Default is 1 second. This is not the same as `for`; this is the
            timeout for failure in the event of no connection.
        force_int_enums : boolean, optional
            Retrieve enums as integers. (Default is strings.)
        repeater : boolean, optional
            Spawn a Channel Access Repeater process if the port is available.
            True default, as the Channel Access spec stipulates that
            well-behaved clients should do this.
        """
        block(self, duration=duration, timeout=timeout,
              force_int_enums=force_int_enums,
              repeater=repeater)

    def interrupt(self):
        """
        Signal to block() to stop blocking. Idempotent.

        This obviously cannot be called interactively while blocked;
        it is intended to be called from another thread.
        This method is a convenience alias for the top-level function
        :func:`interrupt`.
        """
        interrupt()

    def add_callback(self, func):
        """
        Add a callback to receive responses.

        Parameters
        ----------
        func : callable
            Expected signature: ``func(sub, response)``.

            The signature ``func(response)`` is also supported for
            backward-compatibility but will issue warnings. Support will be
            removed in a future release of caproto.

        Returns
        -------
        token : int
            Integer token that can be passed to :meth:`remove_callback`.

        .. versionchanged:: 0.5.0

           Changed the expected signature of ``func`` from ``func(response)``
           to ``func(sub, response)``.
        """
        func = adapt_old_callback_signature(func, self.__wrapper_weakrefs)

        def removed(_):
            self.remove_callback(cb_id)

        if inspect.ismethod(func):
            ref = weakref.WeakMethod(func, removed)
        else:
            # TODO: strong reference to non-instance methods?
            ref = weakref.ref(func, removed)

        with self._callback_lock:
            cb_id = self._callback_id
            self._callback_id += 1
            self.callbacks[cb_id] = ref
        return cb_id

    def remove_callback(self, cb_id):
        """
        Remove callback using token that was returned by :meth:`add_callback`.
        """
        with self._callback_lock:
            self.callbacks.pop(cb_id, None)

    def process(self, response):
        """
        Run the callbacks on a response.

        This is used internally by :func:`block()`, generally not called by the
        user.
        """
        to_remove = []
        with self._callback_lock:
            callbacks = list(self.callbacks.items())

        for cb_id, ref in callbacks:
            callback = ref()
            if callback is None:
                to_remove.append(cb_id)
                continue

            callback(self, response)

        with self._callback_lock:
            for remove_id in to_remove:
                self.callbacks.pop(remove_id, None)

    def clear(self):
        """
        Remove all callbacks. If currently blocking, interrupt.
        """
        interrupt()
        with self._callback_lock:
            for cb_id in list(self.callbacks):
                self.remove_callback(cb_id)


import logging
import os

from django.core.files.base import ContentFile
from django.utils.timezone import now
from django.utils.translation import ugettext as _
from django_scopes import scopes_disabled

from pretix.base.i18n import language
from pretix.base.models import (
    CachedCombinedTicket, CachedTicket, Event, InvoiceAddress, Order,
    OrderPosition,
)
from pretix.base.services.tasks import EventTask, ProfiledTask
from pretix.base.settings import PERSON_NAME_SCHEMES
from pretix.base.signals import allow_ticket_download, register_ticket_outputs
from pretix.celery_app import app
from pretix.helpers.database import rolledback_transaction

logger = logging.getLogger(__name__)

def generate_orderposition(order_position: int, provider: str):
    order_position = OrderPosition.objects.select_related('order', 'order__event').get(id=order_position)

    with language(order_position.order.locale):
        responses = register_ticket_outputs.send(order_position.order.event)
        for receiver, response in responses:
            prov = response(order_position.order.event)
            if prov.identifier == provider:
                filename, ttype, data = prov.generate(order_position)
                path, ext = os.path.splitext(filename)
                for ct in CachedTicket.objects.filter(order_position=order_position, provider=provider):
                    ct.delete()
                ct = CachedTicket.objects.create(order_position=order_position, provider=provider,
                                                 extension=ext, type=ttype, file=None)
                ct.file.save(filename, ContentFile(data))
                return ct.pk

def generate_order(order: int, provider: str):
    order = Order.objects.select_related('event').get(id=order)

    with language(order.locale):
        responses = register_ticket_outputs.send(order.event)
        for receiver, response in responses:
            prov = response(order.event)
            if prov.identifier == provider:
                filename, ttype, data = prov.generate_order(order)
                if ttype == 'text/uri-list':
                    continue

                path, ext = os.path.splitext(filename)
                for ct in CachedCombinedTicket.objects.filter(order=order, provider=provider):
                    ct.delete()
                ct = CachedCombinedTicket.objects.create(order=order, provider=provider, extension=ext,
                                                         type=ttype, file=None)
                ct.file.save(filename, ContentFile(data))
                return ct.pk

@app.task(base=ProfiledTask)
def generate(model: str, pk: int, provider: str):
    with scopes_disabled():
        if model == 'order':
            return generate_order(pk, provider)
        elif model == 'orderposition':
            return generate_orderposition(pk, provider)

class DummyRollbackException(Exception):
    pass

def preview(event: int, provider: str):
    event = Event.objects.get(id=event)

    with rolledback_transaction(), language(event.settings.locale):
        item = event.items.create(name=_("Sample product"), default_price=42.23,
                                  description=_("Sample product description"))
        item2 = event.items.create(name=_("Sample workshop"), default_price=23.40)

        from pretix.base.models import Order
        order = event.orders.create(status=Order.STATUS_PENDING, datetime=now(),
                                    email='sample@pretix.eu',
                                    locale=event.settings.locale,
                                    expires=now(), code="PREVIEW1234", total=119)

        scheme = PERSON_NAME_SCHEMES[event.settings.name_scheme]
        sample = {k: str(v) for k, v in scheme['sample'].items()}
        p = order.positions.create(item=item, attendee_name_parts=sample, price=item.default_price)
        s = event.subevents.first()
        order.positions.create(item=item2, attendee_name_parts=sample, price=item.default_price, addon_to=p, subevent=s)
        order.positions.create(item=item2, attendee_name_parts=sample, price=item.default_price, addon_to=p, subevent=s)

        InvoiceAddress.objects.create(order=order, name_parts=sample, company=_("Sample company"))

        responses = register_ticket_outputs.send(event)
        for receiver, response in responses:
            prov = response(event)
            if prov.identifier == provider:
                return prov.generate(p)

def get_tickets_for_order(order, base_position=None):
    can_download = all([r for rr, r in allow_ticket_download.send(order.event, order=order)])
    if not can_download:
        return []
    if not order.ticket_download_available:
        return []

    providers = [
        response(order.event)
        for receiver, response
        in register_ticket_outputs.send(order.event)
    ]

    tickets = []

    positions = list(order.positions_with_tickets)
    if base_position:
        # Only the given position and its children
        positions = [
            p for p in positions if p.pk == base_position.pk or p.addon_to_id == base_position.pk
        ]

    for p in providers:
        if not p.is_enabled:
            continue

        if p.multi_download_enabled and not base_position:
            try:
                if len(positions) == 0:
                    continue
                ct = CachedCombinedTicket.objects.filter(
                    order=order, provider=p.identifier, file__isnull=False
                ).last()
                if not ct or not ct.file:
                    retval = generate_order(order.pk, p.identifier)
                    if not retval:
                        continue
                    ct = CachedCombinedTicket.objects.get(pk=retval)
                tickets.append((
                    "{}-{}-{}{}".format(
                        order.event.slug.upper(), order.code, ct.provider, ct.extension,
                    ),
                    ct
                ))
            except:
                logger.exception('Failed to generate ticket.')
        else:
            for pos in positions:
                try:
                    ct = CachedTicket.objects.filter(
                        order_position=pos, provider=p.identifier, file__isnull=False
                    ).last()
                    if not ct or not ct.file:
                        retval = generate_orderposition(pos.pk, p.identifier)
                        if not retval:
                            continue
                        ct = CachedTicket.objects.get(pk=retval)

                    if ct.type == 'text/uri-list':
                        continue

                    tickets.append((
                        "{}-{}-{}-{}{}".format(
                            order.event.slug.upper(), order.code, pos.positionid, ct.provider, ct.extension,
                        ),
                        ct
                    ))
                except:
                    logger.exception('Failed to generate ticket.')

    return tickets

@app.task(base=EventTask)
def invalidate_cache(event: Event, item: int=None, provider: str=None, order: int=None, **kwargs):
    qs = CachedTicket.objects.filter(order_position__order__event=event)
    qsc = CachedCombinedTicket.objects.filter(order__event=event)

    if item:
        qs = qs.filter(order_position__item_id=item)

    if provider:
        qs = qs.filter(provider=provider)
        qsc = qsc.filter(provider=provider)

    if order:
        qs = qs.filter(order_position__order_id=order)
        qsc = qsc.filter(order_id=order)

    for ct in qs:
        ct.delete()
    for ct in qsc:
        ct.delete()


"""A set of wrapper functions for accessing the eBusd API."""

from .ebusdpy import (init, read, write, raw)


p8 = plot_implicit(y - 1, y_var=y)
p9 = plot_implicit(x - 1, x_var=x)


# --------------------------------------------------------
# Tensorflow Faster R-CNN
# Licensed under The MIT License [see LICENSE for details]
# Written by Jiasen Lu, Jianwei Yang, based on code from Ross Girshick
# --------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import _init_paths
import os
import sys
import numpy as np
import argparse
import pprint
import pdb
import time
import cv2
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.optim as optim
import pickle
from roi_data_layer.roidb import combined_roidb
from roi_data_layer.roibatchLoader import roibatchLoader
from model.utils.config import cfg, cfg_from_file, cfg_from_list, get_output_dir
from model.rpn.bbox_transform import clip_boxes
from model.nms.nms_wrapper import nms
from model.rpn.bbox_transform import bbox_transform_inv
from model.utils.net_utils import save_net, load_net, vis_detections
from model.faster_rcnn.vgg16 import vgg16
from model.faster_rcnn.resnet import resnet

import pdb

try:
    xrange          # Python 2
except NameError:
    xrange = range  # Python 3

def parse_args():
  """
  Parse input arguments
  """
  parser = argparse.ArgumentParser(description='Train a Fast R-CNN network')
  parser.add_argument('--dataset', dest='dataset',
                      help='training dataset',
                      default='pascal_voc', type=str)
  parser.add_argument('--cfg', dest='cfg_file',
                      help='optional config file',
                      default='cfgs/vgg16.yml', type=str)
  parser.add_argument('--net', dest='net',
                      help='vgg16, res50, res101, res152',
                      default='res101', type=str)
  parser.add_argument('--set', dest='set_cfgs',
                      help='set config keys', default=None,
                      nargs=argparse.REMAINDER)
  parser.add_argument('--load_dir', dest='load_dir',
                      help='directory to load models', default="/srv/share/jyang375/models",
                      type=str)
  parser.add_argument('--cuda', dest='cuda',
                      help='whether use CUDA',
                      action='store_true')
  parser.add_argument('--ls', dest='large_scale',
                      help='whether use large imag scale',
                      action='store_true')
  parser.add_argument('--mGPUs', dest='mGPUs',
                      help='whether use multiple GPUs',
                      action='store_true')
  parser.add_argument('--cag', dest='class_agnostic',
                      help='whether perform class_agnostic bbox regression',
                      action='store_true')
  parser.add_argument('--parallel_type', dest='parallel_type',
                      help='which part of model to parallel, 0: all, 1: model before roi pooling',
                      default=0, type=int)
  parser.add_argument('--checksession', dest='checksession',
                      help='checksession to load model',
                      default=1, type=int)
  parser.add_argument('--checkepoch', dest='checkepoch',
                      help='checkepoch to load network',
                      default=1, type=int)
  parser.add_argument('--checkpoint', dest='checkpoint',
                      help='checkpoint to load network',
                      default=10021, type=int)
  parser.add_argument('--vis', dest='vis',
                      help='visualization mode',
                      action='store_true')
  args = parser.parse_args()
  return args

lr = cfg.TRAIN.LEARNING_RATE
momentum = cfg.TRAIN.MOMENTUM
weight_decay = cfg.TRAIN.WEIGHT_DECAY

if __name__ == '__main__':

  args = parse_args()

  print('Called with args:')
  print(args)

  if torch.cuda.is_available() and not args.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with --cuda")

  np.random.seed(cfg.RNG_SEED)
  if args.dataset == "pascal_voc":
      args.imdb_name = "voc_2007_trainval"
      args.imdbval_name = "voc_2007_test"
      args.set_cfgs = ['ANCHOR_SCALES', '[8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]']
  elif args.dataset == "pascal_voc_0712":
      args.imdb_name = "voc_2007_trainval+voc_2012_trainval"
      args.imdbval_name = "voc_2007_test"
      args.set_cfgs = ['ANCHOR_SCALES', '[8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]']
  elif args.dataset == "coco":
      args.imdb_name = "coco_2014_train+coco_2014_valminusminival"
      args.imdbval_name = "coco_2014_minival"
      args.set_cfgs = ['ANCHOR_SCALES', '[4, 8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]']
  elif args.dataset == "imagenet":
      args.imdb_name = "imagenet_train"
      args.imdbval_name = "imagenet_val"
      args.set_cfgs = ['ANCHOR_SCALES', '[8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]']
  elif args.dataset == "imagenet_vid":
      args.imdb_name = "imagenet_vid_train+imagenet_det_train"
      args.imdbval_name = "imagenet_vid_val"
      args.set_cfgs = ['ANCHOR_SCALES', '[4, 8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]']
  elif args.dataset == "vg":
      args.imdb_name = "vg_150-50-50_minitrain"
      args.imdbval_name = "vg_150-50-50_minival"
      args.set_cfgs = ['ANCHOR_SCALES', '[4, 8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]']

  args.cfg_file = "cfgs/{}_ls.yml".format(args.net) if args.large_scale else "cfgs/{}.yml".format(args.net)

  if args.cfg_file is not None:
    cfg_from_file(args.cfg_file)
  if args.set_cfgs is not None:
    cfg_from_list(args.set_cfgs)

  print('Using config:')
  pprint.pprint(cfg)

  cfg.TRAIN.USE_FLIPPED = False
  imdb, roidb, ratio_list, ratio_index = combined_roidb(args.imdbval_name, False)
  imdb.competition_mode(on=True)

  print('{:d} roidb entries'.format(len(roidb)))

  input_dir = args.load_dir + "/" + args.net + "/" + args.dataset
  if not os.path.exists(input_dir):
    raise Exception('There is no input directory for loading network from ' + input_dir)
  load_name = os.path.join(input_dir,
    'faster_rcnn_{}_{}_{}.pth'.format(args.checksession, args.checkepoch, args.checkpoint))

  # initilize the network here.
  if args.net == 'vgg16':
    fasterRCNN = vgg16(imdb.classes, pretrained=False, class_agnostic=args.class_agnostic)
  elif args.net == 'res101':
    fasterRCNN = resnet(imdb.classes, 101, pretrained=False, class_agnostic=args.class_agnostic)
  elif args.net == 'res50':
    fasterRCNN = resnet(imdb.classes, 50, pretrained=False, class_agnostic=args.class_agnostic)
  elif args.net == 'res152':
    fasterRCNN = resnet(imdb.classes, 152, pretrained=False, class_agnostic=args.class_agnostic)
  else:
    print("network is not defined")
    pdb.set_trace()

  fasterRCNN.create_architecture()

  print("load checkpoint %s" % (load_name))
  checkpoint = torch.load(load_name)
  fasterRCNN.load_state_dict(checkpoint['model'])
  if 'pooling_mode' in checkpoint.keys():
    cfg.POOLING_MODE = checkpoint['pooling_mode']

  print('load model successfully!')
  # initilize the tensor holder here.
  im_data = torch.FloatTensor(1)
  im_info = torch.FloatTensor(1)
  num_boxes = torch.LongTensor(1)
  gt_boxes = torch.FloatTensor(1)

  # ship to cuda
  if args.cuda:
    im_data = im_data.cuda()
    im_info = im_info.cuda()
    num_boxes = num_boxes.cuda()
    gt_boxes = gt_boxes.cuda()

  # make variable
  im_data = Variable(im_data, volatile=True)
  im_info = Variable(im_info, volatile=True)
  num_boxes = Variable(num_boxes, volatile=True)
  gt_boxes = Variable(gt_boxes, volatile=True)

  if args.cuda:
    cfg.CUDA = True

  if args.cuda:
    fasterRCNN.cuda()

  start = time.time()
  max_per_image = 100

  vis = args.vis

  if vis:
    thresh = 0.05
  else:
    thresh = 0.0

  save_name = 'faster_rcnn_10'
  num_images = len(imdb.image_index)
  all_boxes = [[[] for _ in xrange(num_images)]
               for _ in xrange(imdb.num_classes)]

  output_dir = get_output_dir(imdb, save_name)
  dataset = roibatchLoader(roidb, ratio_list, ratio_index, 1, \
                        imdb.num_classes, training=False, normalize = False)
  dataloader = torch.utils.data.DataLoader(dataset, batch_size=1,
                            shuffle=False, num_workers=0,
                            pin_memory=True)

  data_iter = iter(dataloader)

  _t = {'im_detect': time.time(), 'misc': time.time()}
  det_file = os.path.join(output_dir, 'detections.pkl')

  fasterRCNN.eval()
  empty_array = np.transpose(np.array([[],[],[],[],[]]), (1,0))
  for i in range(num_images):

      data = next(data_iter)
      im_data.data.resize_(data[0].size()).copy_(data[0])
      im_info.data.resize_(data[1].size()).copy_(data[1])
      gt_boxes.data.resize_(data[2].size()).copy_(data[2])
      num_boxes.data.resize_(data[3].size()).copy_(data[3])

      det_tic = time.time()
      rois, cls_prob, bbox_pred, \
      rpn_loss_cls, rpn_loss_box, \
      RCNN_loss_cls, RCNN_loss_bbox, \
      rois_label = fasterRCNN(im_data, im_info, gt_boxes, num_boxes)

      scores = cls_prob.data
      boxes = rois.data[:, :, 1:5]

      if cfg.TEST.BBOX_REG:
          # Apply bounding-box regression deltas
          box_deltas = bbox_pred.data
          if cfg.TRAIN.BBOX_NORMALIZE_TARGETS_PRECOMPUTED:
          # Optionally normalize targets by a precomputed mean and stdev
            if args.class_agnostic:
                box_deltas = box_deltas.view(-1, 4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() \
                           + torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda()
                box_deltas = box_deltas.view(1, -1, 4)
            else:
                box_deltas = box_deltas.view(-1, 4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() \
                           + torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda()
                box_deltas = box_deltas.view(1, -1, 4 * len(imdb.classes))

          pred_boxes = bbox_transform_inv(boxes, box_deltas, 1)
          pred_boxes = clip_boxes(pred_boxes, im_info.data, 1)
      else:
          # Simply repeat the boxes, once for each class
          pred_boxes = np.tile(boxes, (1, scores.shape[1]))

      pred_boxes /= data[1][0][2]

      scores = scores.squeeze()
      pred_boxes = pred_boxes.squeeze()
      det_toc = time.time()
      detect_time = det_toc - det_tic
      misc_tic = time.time()
      if vis:
          im = cv2.imread(imdb.image_path_at(i))
          im2show = np.copy(im)
      for j in xrange(1, imdb.num_classes):
          inds = torch.nonzero(scores[:,j]>thresh).view(-1)
          # if there is det
          if inds.numel() > 0:
            cls_scores = scores[:,j][inds]
            _, order = torch.sort(cls_scores, 0, True)
            if args.class_agnostic:
              cls_boxes = pred_boxes[inds, :]
            else:
              cls_boxes = pred_boxes[inds][:, j * 4:(j + 1) * 4]

            cls_dets = torch.cat((cls_boxes, cls_scores.unsqueeze(1)), 1)
            # cls_dets = torch.cat((cls_boxes, cls_scores), 1)
            cls_dets = cls_dets[order]
            keep = nms(cls_dets, cfg.TEST.NMS)
            cls_dets = cls_dets[keep.view(-1).long()]
            if vis:
              im2show = vis_detections(im2show, imdb.classes[j], cls_dets.cpu().numpy(), 0.3)
            all_boxes[j][i] = cls_dets.cpu().numpy()
          else:
            all_boxes[j][i] = empty_array

      # Limit to max_per_image detections *over all classes*
      if max_per_image > 0:
          image_scores = np.hstack([all_boxes[j][i][:, -1]
                                    for j in xrange(1, imdb.num_classes)])
          if len(image_scores) > max_per_image:
              image_thresh = np.sort(image_scores)[-max_per_image]
              for j in xrange(1, imdb.num_classes):
                  keep = np.where(all_boxes[j][i][:, -1] >= image_thresh)[0]
                  all_boxes[j][i] = all_boxes[j][i][keep, :]

      misc_toc = time.time()
      nms_time = misc_toc - misc_tic

      sys.stdout.write('im_detect: {:d}/{:d} {:.3f}s {:.3f}s   \r' \
          .format(i + 1, num_images, detect_time, nms_time))
      sys.stdout.flush()

      if vis:
          cv2.imwrite('result.png', im2show)
          pdb.set_trace()
          #cv2.imshow('test', im2show)
          #cv2.waitKey(0)

  with open(det_file, 'wb') as f:
      pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)

  # print('Evaluating detections')
  # imdb.evaluate_detections(all_boxes, output_dir)

  end = time.time()
  print("test time: %0.4fs" % (end - start))


#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy
import unittest
import tempfile
from functools import reduce
import numpy
import scipy.linalg
from pyscf import gto
from pyscf import lib
import pyscf.lib.parameters as param

mol0 = gto.Mole()
mol0.atom = [
    [1  , (0.,1.,1.)],
    ["O1", (0.,0.,0.)],
    [1  , (1.,1.,0.)], ]
mol0.nucmod = { "O":'gaussian', 3:'g' }
mol0.unit = 'ang'
mol0.basis = {
    "O": [(0, 0, (15, 1)), ] + gto.etbs(((0, 4, 1, 1.8),
                                         (1, 3, 2, 1.8),
                                         (2, 2, 1, 1.8),)),
    "H": [(0, 0, (1, 1, 0), (3, 3, 1), (5, 1, 0)),
          (1, -2, (1, 1)), ]}
mol0.symmetry = 1
mol0.charge = 1
mol0.spin = 1
mol0.verbose = 7
mol0.ecp = {'O1': 'lanl2dz'}
ftmp = tempfile.NamedTemporaryFile()
mol0.output = ftmp.name
mol0.build()

def tearDownModule():
    global mol0, ftmp
    mol0.stdout.close()
    del mol0, ftmp

class KnownValues(unittest.TestCase):
    def test_intor_cross(self):
        mol1 = mol0.unpack(mol0.pack())
        mol1.symmetry = True
        mol1.unit = 'Ang'
        mol1.atom = '''
                1    0  1  1
                O    0  0  0
                h    1  1  0'''
        mol1.basis = {'O': gto.basis.parse('''
C    S
   3047.5249000              0.0018347
    457.3695100              0.0140373
    103.9486900              0.0688426
     29.2101550              0.2321844
      9.2866630              0.4679413
      3.1639270              0.3623120
#     1.                     0.1
C    SP
      7.8682724             -0.1193324              0.0689991
      1.8812885             -0.1608542              0.3164240
      0.5442493              1.1434564              0.7443083
C    SP
      0.1687144              1.0000000              1.0000000'''),
                      'H': '6-31g'}
        mol1.build()
        v = gto.mole.intor_cross('cint1e_ovlp_sph', mol0, mol1)
        self.assertAlmostEqual(numpy.linalg.norm(v), 3.6489423434168562, 1)

    def test_num_basis(self):
        self.assertEqual(mol0.nao_nr(), 34)
        self.assertEqual(mol0.nao_2c(), 64)

    def test_time_reversal_map(self):
        tao = [ -2,  1, -4,  3,  8, -7,  6, -5,-10,  9,-12, 11,-14, 13,-16, 15,-18, 17,
                20,-19, 24,-23, 22,-21, 26,-25, 30,-29, 28,-27, 32,-31, 36,-35, 34,-33,
               -40, 39,-38, 37,-46, 45,-44, 43,-42, 41,-50, 49,-48, 47,-56, 55,-54, 53,
               -52, 51,-58, 57,-60, 59, 64,-63, 62,-61]
        self.assertEqual(list(mol0.time_reversal_map()), tao)

    def test_check_sanity(self):
        mol1 = mol0.copy()
        mol1.x = None
        mol1.copy = None
        mol1.check_sanity()

    def test_nao_range(self):
        self.assertEqual(mol0.nao_nr_range(1,4), (2, 7))
        self.assertEqual(mol0.nao_2c_range(1,4), (4, 12))
        self.assertEqual(numpy.dot(range(mol0.nbas+1), mol0.ao_loc_nr()), 2151)
        self.assertEqual(numpy.dot(range(mol0.nbas+1), mol0.ao_loc_2c()), 4066)

    def test_search_bas(self):
        self.assertEqual(mol0.search_shell_id(1, 1), 7)
        self.assertRaises(RuntimeError, mol0.search_ao_nr, 1, 1, -1, 5)
        self.assertEqual(mol0.search_ao_nr(1, 1, -1, 4), 16)
        mol0.cart = True
        self.assertEqual(mol0.search_ao_nr(2, 1, -1, 1), 30)
        mol0.cart = False

    def test_atom_types(self):
        atoms = [['H0', ( 0, 0, 0)],
                 ['H1', ( 0, 0, 0)],
                 ['H',  ( 0, 0, 0)],
                 ['H3', ( 0, 0, 0)]]
        basis = {'H':'sto3g', 'H1': '6-31g'}
        atmgroup = gto.mole.atom_types(atoms, basis)
        self.assertEqual(atmgroup, {'H': [0, 2, 3], 'H1': [1]})
        atoms = [['H0', ( 0, 0, 0)],
                 ['H1', ( 0, 0, 0)],
                 ['H2', ( 0, 0, 0)],
                 ['H3', ( 0, 0, 0)]]
        basis = {'H2':'sto3g', 'H3':'6-31g', 'H0':'sto3g', 'H1': '6-31g'}
        atmgroup = gto.mole.atom_types(atoms, basis)
        self.assertEqual(atmgroup, {'H2': [2], 'H3': [3], 'H0': [0], 'H1': [1]})

    def test_given_symmetry(self):
        mol = gto.M(atom='H 0 0 -1; H 0 0 1', symmetry='D2h')
        self.assertEqual(mol.irrep_id, [0, 5])
        mol = gto.M(atom='H 0 0 -1; H 0 0 1', symmetry='D2')
        self.assertEqual(mol.irrep_id, [0, 1])
        mol = gto.M(atom='H 0 0 -1; H 0 0 1', symmetry='C2v')
        self.assertEqual(mol.irrep_id, [0])

    def test_dumps_loads(self):
        import warnings
        mol1 = gto.M()
        mol1.x = lambda *args: None
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always")
            d = mol1.dumps()
            self.assertTrue(w[0].category, UserWarning)
        mol1.loads(mol0.dumps())

    def test_symm_orb_serialization(self):
        '''Handle the complex symmetry-adapted orbitals'''
        mol = gto.M(atom='He', basis='ccpvdz', symmetry=True)
        mol.loads(mol.dumps())

        lz_minus = numpy.sqrt(.5) * (mol.symm_orb[3] - mol.symm_orb[2] * 1j)
        lz_plus = -numpy.sqrt(.5) * (mol.symm_orb[3] + mol.symm_orb[2] * 1j)
        mol.symm_orb[2] = lz_minus
        mol.symm_orb[3] = lz_plus
        mol.loads(mol.dumps())
        self.assertTrue(mol.symm_orb[0].dtype == numpy.double)
        self.assertTrue(mol.symm_orb[2].dtype == numpy.complex128)
        self.assertTrue(mol.symm_orb[3].dtype == numpy.complex128)

    def test_same_mol1(self):
        self.assertTrue(gto.same_mol(mol0, mol0))
        mol1 = gto.M(atom='h   0  1  1; O1  0  0  0;  h   1  1  0')
        self.assertTrue(not gto.same_mol(mol0, mol1))
        self.assertTrue(gto.same_mol(mol0, mol1, cmp_basis=False))

        mol1 = gto.M(atom='h   0  1  1; O1  0  0  0;  h   1  1  0.01')
        self.assertTrue(not gto.same_mol(mol0, mol1, cmp_basis=False))
        self.assertTrue(gto.same_mol(mol0, mol1, tol=.02, cmp_basis=False))

        mol1 = gto.M(atom='''H 0.0052917700 0.0000000000 -0.8746076326
                             F 0.0000000000 0.0000000000 0.0516931447''')
        mol2 = gto.M(atom='''H 0.0000000000 0.0000000000 -0.8746076326
                             F 0.0000000000 0.0000000000 0.0516931447''')
        self.assertTrue(gto.same_mol(mol1, mol2))
        self.assertTrue(not gto.same_mol(mol1, mol2, tol=1e-6))
        mol3 = gto.M(atom='''H 0.0000000000 0.0000000000 -0.8746076326
                             H 0.0000000000 0.0000000000 0.0516931447''')
        self.assertTrue(not gto.same_mol(mol3, mol2))

    def test_same_mol2(self):
        mol1 = gto.M(atom='H 0.0052917700 0.0000000000 -0.8746076326; F 0.0000000000 0.0000000000 0.0464013747')
        mol2 = gto.M(atom='H 0.0000000000 0.0000000000 -0.8746076326; F 0.0052917700 0.0000000000 0.0464013747')
        self.assertTrue(gto.same_mol(mol1, mol2))

        mol1 = gto.M(atom='H 0.0052917700 0.0000000000 -0.8693158626; F 0.0000000000 0.0000000000 0.0464013747')
        mol2 = gto.M(atom='H 0.0000000000 0.0052917700 -0.8693158626; F 0.0000000000 0.0000000000 0.0464013747')
        mol3 = gto.M(atom='H 0.0000000000 0.0000000000 -0.8693158626; F 0.0052917700 0.0000000000 0.0464013747')
        mol4 = gto.M(atom='H -0.0052917700 0.0000000000 -0.8746076326; F 0.0000000000 0.0000000000 0.0411096047')
        mols = (mol1, mol2, mol3, mol4)
        for i,mi in enumerate(mols):
            for j in range(i):
                self.assertTrue(gto.same_mol(mols[i], mols[j]))

        mol1 = gto.M(atom='''H 0.0000000000 0.0000000000 0.0000000000
          H 0.9497795800 1.3265673200 0.0000000000
          H 0.9444878100 -1.3265673200 0.0000000000
          H1 -0.9444878100 0.0000000000 1.3265673200
          H1 -0.9444878100 0.0000000000 -1.3265673200''', basis={'H':'sto3g', 'H1':'sto3g'}, charge=1)
        mol2 = gto.M(atom='''H 0.0000000000 0.0000000000 0.0000000000
          H 0.9444878100 1.3265673200 0.0000000000
          H 0.9497795800 -1.3265673200 0.0000000000
          H1 -0.9444878100 0.0000000000 1.3265673200
          H1 -0.9444878100 0.0000000000 -1.3265673200''', basis={'H':'sto3g', 'H1':'sto3g'}, charge=1)
        self.assertTrue(gto.same_mol(mol1, mol2))
        self.assertEqual(len(gto.atom_types(mol1._atom)), 2)
        mol3 = gto.M(atom='''H 0.0000000000 0.0000000000 0.0000000000
          H1 0.9497795800 1.3265673200 0.0000000000
          H1 0.9444878100 -1.3265673200 0.0000000000
          H1 -0.9444878100 0.0000000000 1.3265673200
          H1 -0.9444878100 0.0000000000 -1.3265673200''', basis={'H':'sto3g', 'H1':'321g'}, charge=1)
        self.assertTrue(not gto.same_mol(mol3, mol2))

    def test_inertia_momentum(self):
        mol1 = gto.Mole()
        mol1.atom = mol0.atom
        mol1.nucmod = 'G'
        mol1.verbose = 5
        mol1.nucprop = {'H': {'mass': 3}}
        mol1.output = '/dev/null'
        mol1.build(False, False)
        self.assertAlmostEqual(lib.finger(gto.inertia_moment(mol1)),
                               2.139593709454326, 9)

        mass = mol0.atom_mass_list(isotope_avg=True)
        self.assertAlmostEqual(lib.finger(gto.inertia_moment(mol1, mass)),
                               2.1549269955776205, 9)

    def test_chiral_mol(self):
        mol1 = gto.M(atom='C 0 0 0; H 1 1 1; He -1 -1 1; Li -1 1 -1; Be 1 -1 -1')
        mol2 = gto.M(atom='C 0 0 0; H 1 1 1; He -1 -1 1; Be -1 1 -1; Li 1 -1 -1')
        self.assertTrue(gto.chiral_mol(mol1, mol2))
        self.assertTrue(gto.chiral_mol(mol1))

        mol1 = gto.M(atom='''H 0.9444878100 1.3265673200 0.0052917700
                            H 0.9444878100 -1.3265673200 0.0000000000
                            H -0.9444878100 0.0000000000 1.3265673200
                            H -0.9444878100 0.0000000000 -1.3265673200''')
        mol2 = gto.M(atom='''H 0.9444878100 1.3265673200 0.0000000000
                            H 0.9444878100 -1.3265673200 0.0052917700
                            H -0.9444878100 0.0000000000 1.3265673200
                            H -0.9444878100 0.0000000000 -1.3265673200''')
        self.assertTrue(gto.chiral_mol(mol1, mol2))

        mol1 = gto.M(atom='''H 0.9444878100 1.3265673200 0.0052917700
                            H 0.9444878100 -1.3265673200 0.0000000000
                            H -0.9444878100 0.0000000000 1.3265673200
                            H -0.9444878100 0.0000000000 -1.3265673200''')
        self.assertTrue(gto.chiral_mol(mol1))

    def test_first_argument(self):
        mol1 = gto.Mole()
        mol1.build('He')
        self.assertEqual(mol1.natm, 1)

    def test_atom_as_file(self):
        ftmp = tempfile.NamedTemporaryFile('w')
        # file in xyz format
        ftmp.write('He 0 0 0\nHe 0 0 1\n')
        ftmp.flush()
        mol1 = gto.M(atom=ftmp.name)
        self.assertEqual(mol1.natm, 2)

        # file in zmatrix format
        ftmp = tempfile.NamedTemporaryFile('w')
        ftmp.write('He\nHe 1 1.5\n')
        ftmp.flush()
        mol1 = gto.M(atom=ftmp.name)
        self.assertEqual(mol1.natm, 2)

    def test_format_atom(self):
        atoms = [['h' , 0,1,1], "O1  0. 0. 0.", [1, 1.,1.,0.],]
        self.assertTrue(numpy.allclose(gto.mole.format_atom(atoms, unit='Ang')[0][1],
                                       [0.0, 1.8897261245650618, 1.8897261245650618]))
        atoms = '''h 0 1 1
        O1 0 0 0; 1 1 1 0; #H 0 0 3'''
        self.assertTrue(numpy.allclose(gto.mole.format_atom(atoms, unit=1)[0][1],
                                       [0.0, 1., 1.]))
        atoms = 'O1; h 1 1; 1 1 1 2 90'
        atoms = gto.mole.format_atom(atoms, unit=1)[2]
        self.assertEqual(atoms[0], 'H')
        self.assertTrue(numpy.allclose(atoms[1], [0, 0, 1.]))

    def test_format_basis(self):
        mol = gto.M(atom = '''O 0 0 0; 1 0 1 0; H 0 0 1''',
                    basis = {8: 'ccpvdz'})
        self.assertEqual(mol.nao_nr(), 14)

        mol = gto.M(atom = '''O 0 0 0; H:1 0 1 0; H@2 0 0 1''',
                    basis = {'O': 'ccpvdz', 'H:1': 'sto3g', 'H': 'unc-iglo3'})
        self.assertEqual(mol.nao_nr(), 32)

        mol = gto.M(
            atom = '''O 0 0 0; H1 0 1 0; H2 0 0 1''',
            basis = {'default': ('6-31g', [[0, [.05, 1.]], []]), 'H2': 'sto3g'}
        )
        self.assertEqual(mol.nao_nr(), 14)

        mol = gto.M(
            atom = '''O 0 0 0; H1 0 1 0; H2 0 0 1''',
            basis = {'H1': gto.parse('''
# Parse NWChem format basis string (see https://bse.pnl.gov/bse/portal).
# Comment lines are ignored
#BASIS SET: (6s,3p) -> [2s,1p]
        H    S
              2.9412494             -0.09996723
              0.6834831              0.39951283
              0.2222899              0.70011547
        H    S
              2.9412494             0.15591627
              0.6834831             0.60768372
              0.2222899             0.39195739
                                    ''', optimize=True),
                     'O': 'unc-ccpvdz',
                     'H2': gto.load('sto-3g', 'He')  # or use basis of another atom
                    }
        )
        self.assertEqual(mol.nao_nr(), 29)

        mol = gto.M(
            atom = '''O 0 0 0; H1 0 1 0; H2 0 0 1''',
            basis = {'H': ['sto3g', '''unc
        C    S
             71.6168370              0.15432897
             13.0450960              0.53532814
              3.5305122              0.44463454
        C    SP
              2.9412494             -0.09996723             0.15591627
              0.6834831              0.39951283             0.60768372
              0.2222899              0.70011547             0.39195739
        '''],
                     'O': mol.expand_etbs([(0, 4, 1.5, 2.2),  # s-function
                                           (1, 2, 0.5, 2.2)]) # p-function
                    }
        )
        self.assertEqual(mol.nao_nr(), 42)

        mol = gto.M(
            atom = '''O 0 0 0; H1 0 1 0; H2 0 0 1''',
            basis = ('sto3g', 'ccpvdz', '3-21g',
                     gto.etbs([(0, 4, 1.5, 2.2), (1, 2, 0.5, 2.2)]),
                    [[0, numpy.array([1e3, 1.])]])
        )
        self.assertEqual(mol.nao_nr(), 77)

        mol.atom = 'Hg'
        mol.basis = 'ccpvdz'
        self.assertRaises(RuntimeError, mol.build)

    def test_default_basis(self):
        mol = gto.M(atom=[['h' , 0,1,1], ["O1", (0.,0.,0.)], [1, 1.,1.,0.],],
                    basis={'default':'321g', 'O1': 'sto3g'})
        self.assertEqual(sorted(mol._basis.keys()), ['H', 'O1'])

    def test_parse_pople_basis(self):
        self.assertEqual(len(gto.basis.load('6-31G(d)'      , 'H')), 2)
        self.assertEqual(len(gto.basis.load('6-31G(d)'      , 'C')), 6)
        self.assertEqual(len(gto.basis.load('6-31Gs'        , 'C')), 6)
        self.assertEqual(len(gto.basis.load('6-31G*'        , 'C')), 6)
        self.assertEqual(len(gto.basis.load('6-31G(d,p)'    , 'H')), 3)
        self.assertEqual(len(gto.basis.load('6-31G(d,p)'    , 'C')), 6)
        self.assertEqual(len(gto.basis.load('6-31G(2d,2p)'  , 'H')), 4)
        self.assertEqual(len(gto.basis.load('6-31G(2d,2p)'  , 'C')), 7)
        self.assertEqual(len(gto.basis.load('6-31G(3df,3pd)', 'H')), 6)
        self.assertEqual(len(gto.basis.load('6-31G(3df,3pd)', 'C')), 9)

    def test_parse_basis(self):
        mol = gto.M(atom='''
                    6        0    0   -0.5
                    8        0    0    0.5
                    1        1    0   -1.0
                    1       -1    0   -1.0''',
                    basis='''
#BASIS SET: (3s) -> [2s]
H    S
      5.4471780              0.1562849787
      0.82454724             0.9046908767
H    S
      0.18319158             1.0000000
#BASIS SET: (6s,3p) -> [3s,2p]
C    S
    172.2560000              0.0617669
     25.9109000              0.3587940
      5.5333500              0.7007130
C    SP
      3.6649800             -0.3958970              0.2364600
      0.7705450              1.2158400              0.8606190
C    SP
      0.1958570              1.0000000              1.0000000
#BASIS SET: (6s,3p) -> [3s,2p]
O    S
    322.0370000              0.0592394
     48.4308000              0.3515000
     10.4206000              0.7076580
O    SP
      7.4029400             -0.4044530              0.2445860
      1.5762000              1.2215600              0.8539550
O    SP
      0.3736840              1.0000000              1.0000000
''')
        self.assertTrue(mol.nao_nr() == 22)

    def test_ghost(self):
        mol = gto.M(
            atom = 'C 0 0 0; ghost 0 0 2',
            basis = {'C': 'sto3g', 'ghost': gto.basis.load('sto3g', 'H')}
        )
        self.assertEqual(mol.nao_nr(), 6)

        mol = gto.M(atom='''
        ghost-O     0.000000000     0.000000000     2.500000000
        X_H        -0.663641000    -0.383071000     3.095377000
        ghost.H     0.663588000     0.383072000     3.095377000
        O     1.000000000     0.000000000     2.500000000
        H    -1.663641000    -0.383071000     3.095377000
        H     1.663588000     0.383072000     3.095377000
        ''',
        basis='631g')
        self.assertEqual(mol.nao_nr(), 26)

    def test_nucmod(self):
        gto.filatov_nuc_mod(80)
        self.assertEqual(gto.mole._parse_nuc_mod(1), gto.NUC_GAUSS)
        self.assertEqual(gto.mole._parse_nuc_mod('Gaussian'), gto.NUC_GAUSS)
        mol1 = gto.Mole()
        mol1.atom = mol0.atom
        mol1.nucmod = 'G'
        mol1.verbose = 5
        mol1.nucprop = {'H': {'mass': 3}}
        mol1.output = '/dev/null'
        mol1.build(False, False)
        mol1.set_nuc_mod(0, 2)
        self.assertTrue(mol1._atm[1,gto.NUC_MOD_OF] == gto.NUC_GAUSS)
        self.assertAlmostEqual(mol1._env[mol1._atm[0,gto.PTR_ZETA]], 2, 9)
        self.assertAlmostEqual(mol1._env[mol1._atm[1,gto.PTR_ZETA]], 586314366.54656982, 4)

        mol1.set_nuc_mod(1, 0)
        self.assertTrue(mol1._atm[1,gto.NUC_MOD_OF] == gto.NUC_POINT)

        mol1.nucmod = None
        mol1.build(False, False)
        self.assertTrue(mol1._atm[1,gto.NUC_MOD_OF] == gto.NUC_POINT)

        mol1.nucmod = {'H': gto.filatov_nuc_mod}
        mol1.build(False, False)
        self.assertTrue(mol1._atm[0,gto.NUC_MOD_OF] == gto.NUC_GAUSS)
        self.assertTrue(mol1._atm[1,gto.NUC_MOD_OF] == gto.NUC_POINT)
        self.assertTrue(mol1._atm[2,gto.NUC_MOD_OF] == gto.NUC_GAUSS)

    def test_zmat(self):
        coord = numpy.array((
            (0.200000000000, -1.889726124565,  0.000000000000),
            (1.300000000000, -1.889726124565,  0.000000000000),
            (2.400000000000, -1.889726124565,  0.000000000000),
            (3.500000000000, -1.889726124565,  0.000000000000),
            (0.000000000000,  0.000000000000, -1.889726124565),
            (0.000000000000,  1.889726124565,  0.000000000000),
            (0.200000000000, -0.800000000000,  0.000000000000),
            (1.889726124565,  0.000000000000,  1.133835674739)))
        zstr0 = gto.cart2zmat(coord)
        zstr = '\n'.join(['H '+x for x in zstr0.splitlines()])
        atoms = gto.zmat2cart(zstr)
        zstr1 = gto.cart2zmat([x[1] for x in atoms])
        self.assertTrue(zstr0 == zstr1)

        numpy.random.seed(1)
        coord = numpy.random.random((6,3))
        zstr0 = gto.cart2zmat(coord)
        zstr = '\n'.join(['H '+x for x in zstr0.splitlines()])
        atoms = gto.zmat2cart(zstr)
        zstr1 = gto.cart2zmat([x[1] for x in atoms])
        self.assertTrue(zstr0 == zstr1)

    def test_c2s(self):  # Transformation of cart <-> sph, sph <-> spinor
        c = mol0.sph2spinor_coeff()
        s0 = mol0.intor('int1e_ovlp_spinor')
        s1 = mol0.intor('int1e_ovlp_sph')
        sa = reduce(numpy.dot, (c[0].T.conj(), s1, c[0]))
        sa+= reduce(numpy.dot, (c[1].T.conj(), s1, c[1]))
        mol0.cart = True
        s2 = mol0.intor('int1e_ovlp')
        mol0.cart = False
        self.assertAlmostEqual(abs(s0 - sa).max(), 0, 12)
        c = mol0.cart2sph_coeff()
        sa = reduce(numpy.dot, (c.T.conj(), s2, c))
        self.assertAlmostEqual(abs(s1 - sa).max(), 0, 12)

        c0 = gto.mole.cart2sph(1)
        ca, cb = gto.mole.cart2spinor_l(1)
        ua, ub = gto.mole.sph2spinor_l(1)
        self.assertAlmostEqual(abs(c0.dot(ua)-ca).max(), 0, 9)
        self.assertAlmostEqual(abs(c0.dot(ub)-cb).max(), 0, 9)

        c0 = gto.mole.cart2sph(0, normalized='sp')
        ca, cb = gto.mole.cart2spinor_kappa(-1, 0, normalized='sp')
        ua, ub = gto.mole.sph2spinor_kappa(-1, 0)
        self.assertAlmostEqual(abs(c0.dot(ua)-ca).max(), 0, 9)
        self.assertAlmostEqual(abs(c0.dot(ub)-cb).max(), 0, 9)

        c1 = gto.mole.cart2sph(0, numpy.eye(1))
        self.assertAlmostEqual(abs(c0*0.282094791773878143-c1).max(), 0, 12)

        c0 = gto.mole.cart2sph(1, normalized='sp')
        ca, cb = gto.mole.cart2spinor_kappa(1, 1, normalized='sp')
        ua, ub = gto.mole.sph2spinor_kappa(1, 1)
        self.assertAlmostEqual(abs(c0.dot(ua)-ca).max(), 0, 9)
        self.assertAlmostEqual(abs(c0.dot(ub)-cb).max(), 0, 9)

        c1 = gto.mole.cart2sph(1, numpy.eye(3).T)
        self.assertAlmostEqual(abs(c0*0.488602511902919921-c1).max(), 0, 12)

    def test_bas_method(self):
        self.assertEqual([mol0.bas_len_cart(x) for x in range(mol0.nbas)],
                         [1, 3, 1, 1, 1, 1, 1, 3, 3, 3, 6, 6, 1, 3])
        self.assertEqual([mol0.bas_len_spinor(x) for x in range(mol0.nbas)],
                         [2, 4, 2, 2, 2, 2, 2, 6, 6, 6, 10, 10, 2, 4])
        c0 = mol0.bas_ctr_coeff(0)
        self.assertAlmostEqual(abs(c0[:,0]/c0[0,0] - (1,3,1)).max(), 0, 9)
        self.assertAlmostEqual(abs(c0[:,1] - (0,1,0)).max(), 0, 9)

        self.assertRaises(ValueError, mol0.gto_norm, -1, 1.)

    def test_nelectron(self):
        mol0.nelectron = mol0.nelectron
        mol0.nelectron = mol0.nelectron
        mol0.spin = 2
        self.assertRaises(RuntimeError, lambda *args: mol0.nelec)
        mol0.spin = 1

        mol1 = copy.copy(mol0)
        self.assertEqual(mol1.nelec, (5, 4))
        mol1.nelec = (3, 6)
        self.assertEqual(mol1.nelec, (3, 6))

    def test_multiplicity(self):
        mol1 = copy.copy(mol0)
        self.assertEqual(mol1.multiplicity, 2)
        mol1.multiplicity = 5
        self.assertEqual(mol1.multiplicity, 5)
        self.assertEqual(mol1.spin, 4)
        self.assertRaises(RuntimeError, lambda:mol1.nelec)

    def test_ms(self):
        mol1 = copy.copy(mol0)
        self.assertEqual(mol1.ms, 0.5)
        mol1.ms = 1
        self.assertEqual(mol1.multiplicity, 3)
        self.assertEqual(mol1.spin, 2)
        self.assertRaises(RuntimeError, lambda:mol1.nelec)

    def test_atom_method(self):
        aoslice = mol0.aoslice_by_atom()
        for i in range(mol0.natm):
            symb = mol0.atom_pure_symbol(i)
            shls = mol0.atom_shell_ids(i)
            nshls = aoslice[i][1] - aoslice[i][0]
            self.assertEqual(shls[0], aoslice[i][0])
            self.assertEqual(len(shls), nshls)
            self.assertEqual(mol0.atom_nshells(i), nshls)
        aoslice = mol0.aoslice_2c_by_atom()
        self.assertEqual([x[2] for x in aoslice], [0, 8, 56])
        self.assertEqual([x[3] for x in aoslice], [8, 56, 64])

    def test_dump_loads(self):
        import json
        tmpfile = tempfile.NamedTemporaryFile()
        lib.chkfile.save_mol(mol0, tmpfile.name)
        mol1 = gto.Mole()
        mol1.update(tmpfile.name)
        self.assertEqual(json.loads(mol1.dumps()), json.loads(mol0.dumps()))
        tmpfile = None
        mol1.loads(mol1.dumps())
        mol1.loads_(mol0.dumps())
        mol1.unpack(mol1.pack())
        mol1.unpack_(mol0.pack())

    def test_set_geom(self):
        mol1 = gto.Mole()
        mol1.verbose = 5
        mol1.set_geom_(mol0._atom, 'B', symmetry=True)
        mol1.set_geom_(mol0.atom_coords(), 'B', inplace=False)

        mol1.symmetry = False
        mol1.set_geom_(mol0.atom_coords(), 'B')
        mol1.set_geom_(mol0.atom_coords(), inplace=False)
        mol1.set_geom_(mol0.atom_coords(), unit=1.)
        mol1.set_geom_(mol0.atom_coords(), unit='Ang', inplace=False)

    def test_apply(self):
        from pyscf import scf, mp
        self.assertTrue(isinstance(mol0.apply('RHF'), scf.rohf.ROHF))
        self.assertTrue(isinstance(mol0.apply('MP2'), mp.ump2.UMP2))
        self.assertTrue(isinstance(mol0.apply(scf.RHF), scf.rohf.ROHF))
        self.assertTrue(isinstance(mol0.apply(scf.uhf.UHF), scf.uhf.UHF))

    def test_with_MoleContext(self):
        mol1 = mol0.copy()
        with mol1.with_rinv_as_nucleus(1):
            self.assertTrue(mol1._env[gto.PTR_RINV_ZETA] != 0)
            self.assertAlmostEqual(abs(mol1._env[gto.PTR_RINV_ORIG+2]), 0.46288647587915266, 9)
        self.assertAlmostEqual(mol1._env[gto.PTR_RINV_ZETA], 0, 9)
        self.assertAlmostEqual(mol1._env[gto.PTR_RINV_ORIG+2], 0, 9)
        with mol1.with_rinv_as_nucleus(0):
            self.assertAlmostEqual(abs(mol1._env[gto.PTR_RINV_ORIG+2]), 1.8515459035166109, 9)
        self.assertAlmostEqual(mol1._env[gto.PTR_RINV_ORIG+2], 0, 9)

        with mol1.with_rinv_zeta(20):
            self.assertAlmostEqual(mol1._env[gto.PTR_RINV_ZETA], 20, 9)
            mol1.set_rinv_zeta(3.)
        self.assertAlmostEqual(mol1._env[gto.PTR_RINV_ZETA], 0, 9)

        with mol1.with_rinv_origin((1,2,3)):
            self.assertAlmostEqual(mol1._env[gto.PTR_RINV_ORIG+2], 3, 9)
        self.assertAlmostEqual(mol1._env[gto.PTR_RINV_ORIG+2], 0, 9)

        with mol1.with_range_coulomb(20):
            self.assertAlmostEqual(mol1._env[gto.PTR_RANGE_OMEGA], 20, 9)
            mol1.set_range_coulomb(2.)
        self.assertAlmostEqual(mol1._env[gto.PTR_RANGE_OMEGA], 0, 9)

        with mol1.with_common_origin((1,2,3)):
            self.assertAlmostEqual(mol1._env[gto.PTR_COMMON_ORIG+2], 3, 9)
        self.assertAlmostEqual(mol1._env[gto.PTR_COMMON_ORIG+2], 0, 9)

        mol1.set_f12_zeta(2.)

    def test_input_symmetry(self):
        mol1 = gto.Mole()
        mol1.atom = 'H 1 1 1; H -1 -1 1; H 1 -1 -1; H -1 1 -1'
        mol1.unit = 'B'
        mol1.symmetry = True
        mol1.verbose = 5
        mol1.output = '/dev/null'
        mol1.build()
        self.assertAlmostEqual(lib.finger(mol1.atom_coords()), 4.2517312170868475, 9)

        mol1 = gto.Mole()
        mol1 = gto.Mole()
        mol1.atom = 'H 0 0 -1; H 0 0 1'
        mol1.cart = True
        mol1.unit = 'B'
        mol1.symmetry = 'Dooh'
        mol1.verbose = 5
        mol1.output = '/dev/null'
        mol1.build()
        self.assertAlmostEqual(lib.finger(mol1.atom_coords()), 0.69980902201036865, 9)

        mol1 = gto.Mole()
        mol1.atom = 'H 0 -1 0; H 0 1 0'
        mol1.unit = 'B'
        mol1.symmetry = True
        mol1.symmetry_subgroup = 'D2h'
        mol1.build()
        self.assertAlmostEqual(lib.finger(mol1.atom_coords()), 0.69980902201036865, 9)

        mol1.atom = 'H 0 0 -1; H 0 0 1'
        mol1.unit = 'B'
        mol1.symmetry = 'Coov'
        mol1.symmetry_subgroup = 'C2'
        mol1.build()
        self.assertAlmostEqual(lib.finger(mol1.atom_coords()), 0.69980902201036865, 9)

        mol1.atom = 'H 1 0 -1; H 0 0 1'
        mol1.symmetry = 'Coov'
        self.assertRaises(RuntimeWarning, mol1.build)

        mol1.atom = '''
        C 0. 0. 0.7264
        C 0. 0. -.7264
        H 0.92419 0. 1.29252
        H -.92419 0. 1.29252
        H 0. 0.92419 -1.29252
        H 0. -.92419 -1.29252'''
        mol1.symmetry = True
        mol1.symmetry_subgroup = 'C2v'
        mol1.build()
        self.assertAlmostEqual(lib.finger(mol1.atom_coords()), -0.5215310671099358, 9)

    def test_search_ao_label(self):
        mol1 = mol0.copy()
        mol1.atom = mol0.atom + ['Mg 1,1,1']
        mol1.ecp['Mg'] = 'lanl2dz'
        mol1.basis['Mg'] = 'lanl2dz'
        mol1.build(0, 0)
        self.assertEqual(list(mol1.search_ao_label('O.*2p')), [10,11,12])
        self.assertEqual(list(mol1.search_ao_label('O1 2p')), [10,11,12])
        self.assertEqual(list(mol1.search_ao_label(['O.*2p','0 H 1s'])), [0, 10,11,12])
        self.assertEqual(list(mol1.search_ao_label([10,11,12])), [10,11,12])
        self.assertEqual(list(mol1.search_ao_label(lambda x: '4d' in x)), [24,25,26,27,28])
        mol1.ao_labels(fmt='%s%s%s%s')
        mol1.sph_labels(fmt=None)
        mol1.cart = True
        self.assertEqual(list(mol1.search_ao_label('4d')), [25,26,27,28,29,30])
        mol1.ao_labels(fmt='%s%s%s%s')
        mol1.ao_labels(fmt=None)
        mol1.cart = False
        mol1.spinor_labels()
        mol1.spinor_labels(fmt='%s%s%s%s')
        mol1.spinor_labels(fmt=None)

    def test_input_ecp(self):
        mol1 = gto.Mole()
        mol1.atom = mol0.atom
        mol1.ecp = 'lanl2dz'
        mol1.build(False, False)
        gto.basis.load_ecp('lanl08', 'O')
        gto.format_ecp({'O':'lanl08', 1:'lanl2dz'})
        self.assertRaises(KeyError, gto.format_ecp, {'H':'lan2ldz'})

    def test_condense_to_shell(self):
        mol1 = mol0.copy()
        mol1.symmetry = False
        mol1.build(False, False)
        v = gto.condense_to_shell(mol1, mol1.intor('int1e_ovlp'), numpy.max)
        self.assertAlmostEqual(lib.finger(v), 5.7342530154117846, 9)

    def test_input_ghost_atom(self):
        mol = gto.M(
            atom = 'C 0 0 0; ghost 0 0 2',
            basis = {'C': 'sto3g', 'ghost': gto.basis.load('sto3g', 'H')}
        )

        mol = gto.M(atom='''
        ghost1     0.000000000     0.000000000     2.500000000
        ghost2    -0.663641000    -0.383071000     3.095377000
        ghost2     0.663588000     0.383072000     3.095377000
        O     1.000000000     0.000000000     2.500000000
        H    -1.663641000    -0.383071000     3.095377000
        H     1.663588000     0.383072000     3.095377000
        ''',
                    basis={'ghost1':gto.basis.load('sto3g', 'O'),
                           'ghost2':gto.basis.load('631g', 'H'),
                           'O':'631g', 'H':'631g'}
        )

        mol = gto.M(atom='''
        ghost-O     0.000000000     0.000000000     2.500000000
        ghost_H    -0.663641000    -0.383071000     3.095377000
        ghost:H     0.663588000     0.383072000     3.095377000
        O     1.000000000     0.000000000     2.500000000
        H    -1.663641000    -0.383071000     3.095377000
        H     1.663588000     0.383072000     3.095377000
        ''', basis='631g')

        mol = gto.M(atom='''
        X1     0.000000000     0.000000000     2.500000000
        X2    -0.663641000    -0.383071000     3.095377000
        X2     0.663588000     0.383072000     3.095377000
        O     1.000000000     0.000000000     2.500000000
        H    -1.663641000    -0.383071000     3.095377000
        H     1.663588000     0.383072000     3.095377000
        ''',
                    basis={'X1':gto.basis.load('sto3g', 'O'),
                           'X2':gto.basis.load('631g', 'H'),
                           'O':'631g', 'H':'631g'}
        )

        mol = gto.M(atom='''
        X-O     0.000000000     0.000000000     2.500000000
        X_H1   -0.663641000    -0.383071000     3.095377000
        X:H     0.663588000     0.383072000     3.095377000
        O     1.000000000     0.000000000     2.500000000
        H    -1.663641000    -0.383071000     3.095377000
        H     1.663588000     0.383072000     3.095377000
        ''', basis='631g')

    def test_conc_mole(self):
        mol1 = gto.M(atom='Mg', ecp='LANL2DZ', basis='lanl2dz')
        mol2 = mol1 + mol0
        self.assertEqual(mol2.natm, 4)
        self.assertEqual(mol2.nbas, 18)
        self.assertEqual(mol2.nao_nr(), 42)
        mol2 = mol0 + mol1
        self.assertEqual(mol2.natm, 4)
        self.assertEqual(mol2.nbas, 18)
        self.assertEqual(mol2.nao_nr(), 42)
        n0 = mol0.npgto_nr()
        n1 = mol1.npgto_nr()
        self.assertEqual(mol2.npgto_nr(), n0+n1)
        mol2 = mol2 + mol2
        mol2.cart = True
        self.assertEqual(mol2.npgto_nr(), 100)

    def test_intor_cross(self):
        mol1 = gto.M(atom='He', basis={'He': [(2,(1.,1))]}, cart=True)
        s0 = gto.intor_cross('int1e_ovlp', mol1, mol0)
        self.assertEqual(s0.shape, (6, 34))
        s0 = gto.intor_cross('int1e_ovlp', mol0, mol1)
        self.assertEqual(s0.shape, (34, 6))
        s0 = gto.intor_cross('int1e_ovlp_cart', mol0, mol1)
        self.assertEqual(s0.shape, (36, 6))

    def test_energy_nuc(self):
        self.assertAlmostEqual(mol0.get_enuc(), 6.3611415029455705, 9)
        self.assertAlmostEqual(gto.M().energy_nuc(), 0, 9)

    def test_fakemol(self):
        numpy.random.seed(1)
        coords = numpy.random.random((6,3))*4
        vref = 0
        mol = mol0.copy()
        for c in coords:
            mol.set_rinv_origin(c)
            vref += mol.intor('int1e_rinv')

        fakemol = gto.fakemol_for_charges(coords)
        pmol = mol + fakemol
        shls_slice = (0, mol.nbas, 0, mol.nbas, mol.nbas, pmol.nbas)
        v = pmol.intor('int3c2e', comp=1, shls_slice=shls_slice)
        v = numpy.einsum('pqk->pq', v)
        self.assertAlmostEqual(abs(vref-v).max(), 0, 12)

    def test_to_uncontracted_cartesian_basis(self):
        pmol, ctr_coeff = mol0.to_uncontracted_cartesian_basis()
        c = scipy.linalg.block_diag(*ctr_coeff)
        s = reduce(numpy.dot, (c.T, pmol.intor('int1e_ovlp'), c))
        self.assertAlmostEqual(abs(s-mol0.intor('int1e_ovlp')).max(), 0, 9)

        mol0.cart = True
        pmol, ctr_coeff = mol0.to_uncontracted_cartesian_basis()
        c = scipy.linalg.block_diag(*ctr_coeff)
        s = reduce(numpy.dot, (c.T, pmol.intor('int1e_ovlp'), c))
        self.assertAlmostEqual(abs(s-mol0.intor('int1e_ovlp')).max(), 0, 9)
        mol0.cart = False

if __name__ == "__main__":
    print("test mole.py")
    unittest.main()


import json

from .common import InfoExtractor
from .youtube import YoutubeIE
from ..compat import compat_b64decode
from ..utils import (
    clean_html,
    ExtractorError
)

class ChilloutzoneIE(InfoExtractor):
    _VALID_URL = r'https?://(?:www\.)?chilloutzone\.net/video/(?P<id>[\w|-]+)\.html'
    _TESTS = [{
        'url': 'http://www.chilloutzone.net/video/enemene-meck-alle-katzen-weg.html',
        'md5': 'a76f3457e813ea0037e5244f509e66d1',
        'info_dict': {
            'id': 'enemene-meck-alle-katzen-weg',
            'ext': 'mp4',
            'title': 'Enemene Meck - Alle Katzen weg',
            'description': 'Ist das der Umkehrschluss des Niesenden Panda-Babys?',
        },
    }, {
        'note': 'Video hosted at YouTube',
        'url': 'http://www.chilloutzone.net/video/eine-sekunde-bevor.html',
        'info_dict': {
            'id': '1YVQaAgHyRU',
            'ext': 'mp4',
            'title': '16 Photos Taken 1 Second Before Disaster',
            'description': 'md5:58a8fcf6a459fe0a08f54140f0ad1814',
            'uploader': 'BuzzFeedVideo',
            'uploader_id': 'BuzzFeedVideo',
            'upload_date': '20131105',
        },
    }, {
        'note': 'Video hosted at Vimeo',
        'url': 'http://www.chilloutzone.net/video/icon-blending.html',
        'md5': '2645c678b8dc4fefcc0e1b60db18dac1',
        'info_dict': {
            'id': '85523671',
            'ext': 'mp4',
            'title': 'The Sunday Times - Icons',
            'description': 're:(?s)^Watch the making of - makingoficons.com.{300,}',
            'uploader': 'Us',
            'uploader_id': 'usfilms',
            'upload_date': '20140131'
        },
    }]

    def _real_extract(self, url):
        mobj = self._match_valid_url(url)
        video_id = mobj.group('id')

        webpage = self._download_webpage(url, video_id)

        base64_video_info = self._html_search_regex(
            r'var cozVidData = "(.+?)";', webpage, 'video data')
        decoded_video_info = compat_b64decode(base64_video_info).decode('utf-8')
        video_info_dict = json.loads(decoded_video_info)

        # get video information from dict
        video_url = video_info_dict['mediaUrl']
        description = clean_html(video_info_dict.get('description'))
        title = video_info_dict['title']
        native_platform = video_info_dict['nativePlatform']
        native_video_id = video_info_dict['nativeVideoId']
        source_priority = video_info_dict['sourcePriority']

        # If nativePlatform is None a fallback mechanism is used (i.e. youtube embed)
        if native_platform is None:
            youtube_url = YoutubeIE._extract_url(webpage)
            if youtube_url:
                return self.url_result(youtube_url, ie=YoutubeIE.ie_key())

        # Non Fallback: Decide to use native source (e.g. youtube or vimeo) or
        # the own CDN
        if source_priority == 'native':
            if native_platform == 'youtube':
                return self.url_result(native_video_id, ie='Youtube')
            if native_platform == 'vimeo':
                return self.url_result(
                    'http://vimeo.com/' + native_video_id, ie='Vimeo')

        if not video_url:
            raise ExtractorError('No video found')

        return {
            'id': video_id,
            'url': video_url,
            'ext': 'mp4',
            'title': title,
            'description': description,
        }