"""Python API for the lab streaming layer.

The lab streaming layer provides a set of functions to make instrument data
accessible in real time within a lab network. From there, streams can be
picked up by recording programs, viewing programs or custom experiment
applications that access data streams in real time.

The API covers two areas:
- The "push API" allows to create stream outlets and to push data (regular
  or irregular measurement time series, event data, coded audio/video frames,
  etc.) into them.
- The "pull API" allows to create stream inlets and read time-synched
  experiment data from them (for recording, viewing or experiment control).

pylsl has been tested with Python 2.7 and 3.4.

"""

import os
import platform
import struct
from ctypes import CDLL, util, byref, c_char_p, c_void_p, c_double, c_int, \
    c_long, c_float, c_short, c_byte, c_longlong, cast, POINTER

__all__ = ['IRREGULAR_RATE', 'DEDUCED_TIMESTAMP', 'FOREVER', 'cf_float32',
           'cf_double64', 'cf_string', 'cf_int32', 'cf_int16', 'cf_int8',
           'cf_int64', 'cf_undefined', 'protocol_version', 'library_version',
           'library_info',
           'local_clock', 'StreamInfo', 'StreamOutlet', 'resolve_streams',
           'resolve_byprop', 'resolve_bypred', 'StreamInlet', 'XMLElement',
           'ContinuousResolver', 'TimeoutError', 'LostError',
           'InvalidArgumentError', 'InternalError', 'stream_info',
           'stream_outlet', 'stream_inlet', 'xml_element', 'timeout_error',
           'lost_error', 'vectorf', 'vectord', 'vectorl', 'vectori',
           'vectors', 'vectorc', 'vectorstr', 'resolve_stream']

# =================
# === Constants ===
# =================

# Constant to indicate that a stream has variable sampling rate.
IRREGULAR_RATE = 0.0

# Constant to indicate that a sample has the next successive time stamp
# according to the stream's defined sampling rate. Optional optimization to
# transmit less data per sample.
DEDUCED_TIMESTAMP = -1.0

# A very large time value (ca. 1 year); can be used in timeouts.
FOREVER = 32000000.0

# Value formats supported by LSL. LSL data streams are sequences of samples,
# each of which is a same-size vector of values with one of the below types.

# For up to 24-bit precision measurements in the appropriate physical unit (
# e.g., microvolts). Integers from -16777216 to 16777216 are represented
# accurately.
cf_float32 = 1
# For universal numeric data as long as permitted by network and disk budget.
#  The largest representable integer is 53-bit.
cf_double64 = 2
# For variable-length ASCII strings or data blobs, such as video frames,
# complex event descriptions, etc.
cf_string = 3
# For high-rate digitized formats that require 32-bit precision. Depends
# critically on meta-data to represent meaningful units. Useful for
# application event codes or other coded data.
cf_int32 = 4
# For very high bandwidth signals or CD quality audio (for professional audio
#  float is recommended).
cf_int16 = 5
# For binary signals or other coded data.
cf_int8 = 6
# For now only for future compatibility. Support for this type is not
# available on all languages and platforms.
cf_int64 = 7
# Can not be transmitted.
cf_undefined = 0

# Post processing flags
proc_none = 0
proc_clocksync = 1
proc_dejitter = 2
proc_monotonize = 4
proc_threadsafe = 8
proc_ALL = proc_none | proc_clocksync | proc_dejitter | proc_monotonize | proc_threadsafe

# ==========================================================
# === Free Functions provided by the lab streaming layer ===
# ==========================================================

def protocol_version():
    """Protocol version.

    The major version is protocol_version() / 100;
    The minor version is protocol_version() % 100;

    Clients with different minor versions are protocol-compatible with each
    other while clients with different major versions will refuse to work
    together.

    """
    return lib.lsl_protocol_version()

def library_version():
    """Version of the underlying liblsl library.

    The major version is library_version() / 100;
    The minor version is library_version() % 100;

    """
    return lib.lsl_library_version()

def library_info():
    """Get a string containing library information. The format of the string shouldn't be used
    for anything important except giving a a debugging person a good idea which exact library
    version is used."""
    return lib.lsl_library_info().decode('utf-8')

def local_clock():
    """Obtain a local system time stamp in seconds.

    The resolution is better than a milisecond. This reading can be used to
    assign time stamps to samples as they are being acquired.

    If the "age" of a sample is known at a particular time (e.g., from USB
    transmission delays), it can be used as an offset to lsl_local_clock() to
    obtain a better estimate of when a sample was actually captured. See
    StreamOutlet.push_sample() for a use case.

    """
    return lib.lsl_local_clock()

# ==========================
# === Stream Declaration ===
# ==========================

class StreamInfo:
    """The StreamInfo object stores the declaration of a data stream.

    Represents the following information:
     a) stream data format (#channels, channel format)
     b) core information (stream name, content type, sampling rate)
     c) optional meta-data about the stream content (channel labels,
        measurement units, etc.)

    Whenever a program wants to provide a new stream on the lab network it will
    typically first create a StreamInfo to describe its properties and then
    construct a StreamOutlet with it to create the stream on the network.
    Recipients who discover the outlet can query the StreamInfo; it is also
    written to disk when recording the stream (playing a similar role as a file
    header).

    """

    def __init__(self, name='untitled', type='', channel_count=1,
                 nominal_srate=IRREGULAR_RATE, channel_format=cf_float32,
                 source_id='', handle=None):
        """Construct a new StreamInfo object.

        Core stream information is specified here. Any remaining meta-data can
        be added later.

        Keyword arguments:
        name -- Name of the stream. Describes the device (or product series)
                that this stream makes available (for use by programs,
                experimenters or data analysts). Cannot be empty.
        type -- Content type of the stream. By convention LSL uses the content
                types defined in the XDF file format specification where
                applicable (https://github.com/sccn/xdf). The content type is the
                preferred way to find streams (as opposed to searching by name).
        channel_count -- Number of channels per sample. This stays constant for
                         the lifetime of the stream. (default 1)
        nominal_srate -- The sampling rate (in Hz) as advertised by the data
                         source, regular (otherwise set to IRREGULAR_RATE).
                         (default IRREGULAR_RATE)
        channel_format -- Format/type of each channel. If your channels have
                          different formats, consider supplying multiple
                          streams or use the largest type that can hold
                          them all (such as cf_double64). It is also allowed
                          to pass this as a string, without the cf_ prefix,
                          e.g., 'float32' (default cf_float32)
        source_id -- Unique identifier of the device or source of the data, if
                     available (such as the serial number). This is critical
                     for system robustness since it allows recipients to
                     recover from failure even after the serving app, device or
                     computer crashes (just by finding a stream with the same
                     source id on the network again). Therefore, it is highly
                     recommended to always try to provide whatever information
                     can uniquely identify the data source itself.
                     (default '')

        """
        if handle is not None:
            self.obj = c_void_p(handle)
        else:
            if isinstance(channel_format, str):
                channel_format = string2fmt[channel_format]
            self.obj = lib.lsl_create_streaminfo(c_char_p(str.encode(name)),
                                                 c_char_p(str.encode(type)),
                                                 channel_count,
                                                 c_double(nominal_srate),
                                                 channel_format,
                                                 c_char_p(str.encode(source_id)))
            self.obj = c_void_p(self.obj)
            if not self.obj:
                raise RuntimeError("could not create stream description "
                                   "object.")

    def __del__(self):
        """ Destroy a previously created StreamInfo object. """
        # noinspection PyBroadException
        try:
            lib.lsl_destroy_streaminfo(self.obj)
        except:
            pass

    # === Core Information (assigned at construction) ===

    def name(self):
        """Name of the stream.

        This is a human-readable name. For streams offered by device modules,
        it refers to the type of device or product series that is generating
        the data of the stream. If the source is an application, the name may
        be a more generic or specific identifier. Multiple streams with the
        same name can coexist, though potentially at the cost of ambiguity (for
        the recording app or experimenter).

        """
        return lib.lsl_get_name(self.obj).decode('utf-8')

    def type(self):
        """Content type of the stream.

        The content type is a short string such as "EEG", "Gaze" which
        describes the content carried by the channel (if known). If a stream
        contains mixed content this value need not be assigned but may instead
        be stored in the description of channel types. To be useful to
        applications and automated processing systems using the recommended
        content types is preferred.

        """
        return lib.lsl_get_type(self.obj).decode('utf-8')

    def channel_count(self):
        """Number of channels of the stream.

        A stream has at least one channel; the channel count stays constant for
        all samples.

        """
        return lib.lsl_get_channel_count(self.obj)

    def nominal_srate(self):
        """Sampling rate of the stream, according to the source (in Hz).

        If a stream is irregularly sampled, this should be set to
        IRREGULAR_RATE.

        Note that no data will be lost even if this sampling rate is incorrect
        or if a device has temporary hiccups, since all samples will be
        transmitted anyway (except for those dropped by the device itself).
        However, when the recording is imported into an application, a good
        data importer may correct such errors more accurately if the advertised
        sampling rate was close to the specs of the device.

        """
        return lib.lsl_get_nominal_srate(self.obj)

    def channel_format(self):
        """Channel format of the stream.

        All channels in a stream have the same format. However, a device might
        offer multiple time-synched streams each with its own format.

        """
        return lib.lsl_get_channel_format(self.obj)

    def source_id(self):
        """Unique identifier of the stream's source, if available.

        The unique source (or device) identifier is an optional piece of
        information that, if available, allows that endpoints (such as the
        recording program) can re-acquire a stream automatically once it is
        back online.

        """
        return lib.lsl_get_source_id(self.obj).decode('utf-8')

    # === Hosting Information (assigned when bound to an outlet/inlet) ===

    def version(self):
        """Protocol version used to deliver the stream."""
        return lib.lsl_get_version(self.obj)

    def created_at(self):
        """Creation time stamp of the stream.

        This is the time stamp when the stream was first created
        (as determined via local_clock() on the providing machine).

        """
        return lib.lsl_get_created_at(self.obj)

    def uid(self):
        """Unique ID of the stream outlet instance (once assigned).

        This is a unique identifier of the stream outlet, and is guaranteed to
        be different across multiple instantiations of the same outlet (e.g.,
        after a re-start).

        """
        return lib.lsl_get_uid(self.obj).decode('utf-8')

    def session_id(self):
        """Session ID for the given stream.

        The session id is an optional human-assigned identifier of the
        recording session. While it is rarely used, it can be used to prevent
        concurrent recording activitites on the same sub-network (e.g., in
        multiple experiment areas) from seeing each other's streams
        (can be assigned in a configuration file read by liblsl, see also
        Network Connectivity in the LSL wiki).

        """
        return lib.lsl_get_session_id(self.obj).decode('utf-8')

    def hostname(self):
        """Hostname of the providing machine."""
        return lib.lsl_get_hostname(self.obj).decode('utf-8')

    # === Data Description (can be modified) ===

    def desc(self):
        """Extended description of the stream.

        It is highly recommended that at least the channel labels are described
        here. See code examples on the LSL wiki. Other information, such
        as amplifier settings, measurement units if deviating from defaults,
        setup information, subject information, etc., can be specified here, as
        well. Meta-data recommendations follow the XDF file format project
        (github.com/sccn/xdf/wiki/Meta-Data or web search for: XDF meta-data).

        Important: if you use a stream content type for which meta-data
        recommendations exist, please try to lay out your meta-data in
        agreement with these recommendations for compatibility with other
        applications.

        """
        return XMLElement(lib.lsl_get_desc(self.obj))

    def as_xml(self):
        """Retrieve the entire stream_info in XML format.

        This yields an XML document (in string form) whose top-level element is
        <description>. The description element contains one element for each
        field of the stream_info class, including:
        a) the core elements <name>, <type>, <channel_count>, <nominal_srate>,
           <channel_format>, <source_id>
        b) the misc elements <version>, <created_at>, <uid>, <session_id>,
           <v4address>, <v4data_port>, <v4service_port>, <v6address>,
           <v6data_port>, <v6service_port>
        c) the extended description element <desc> with user-defined
           sub-elements.

        """
        return lib.lsl_get_xml(self.obj).decode('utf-8')

# =====================
# === Stream Outlet ===
# =====================

class StreamOutlet:
    """A stream outlet.

    Outlets are used to make streaming data (and the meta-data) available on
    the lab network.

    """

    def __init__(self, info, chunk_size=0, max_buffered=360):
        """Establish a new stream outlet. This makes the stream discoverable.

        Keyword arguments:
        description -- The StreamInfo object to describe this stream. Stays
                constant over the lifetime of the outlet.
        chunk_size --- Optionally the desired chunk granularity (in samples)
                       for transmission. If unspecified, each push operation
                       yields one chunk. Inlets can override this setting.
                       (default 0)
        max_buffered -- Optionally the maximum amount of data to buffer (in
                        seconds if there is a nominal sampling rate, otherwise
                        x100 in samples). The default is 6 minutes of data.
                        Note that, for high-bandwidth data, you will want to
                        use a lower value here to avoid running out of RAM.
                        (default 360)

        """
        self.obj = lib.lsl_create_outlet(info.obj, chunk_size, max_buffered)
        self.obj = c_void_p(self.obj)
        if not self.obj:
            raise RuntimeError("could not create stream outlet.")
        self.channel_format = info.channel_format()
        self.channel_count = info.channel_count()
        self.do_push_sample = fmt2push_sample[self.channel_format]
        self.do_push_chunk = fmt2push_chunk[self.channel_format]
        self.value_type = fmt2type[self.channel_format]
        self.sample_type = self.value_type*self.channel_count

    def __del__(self):
        """Destroy an outlet.

        The outlet will no longer be discoverable after destruction and all
        connected inlets will stop delivering data.

        """
        # noinspection PyBroadException
        try:
            lib.lsl_destroy_outlet(self.obj)
        except:
            pass

    def push_sample(self, x, timestamp=0.0, pushthrough=True):
        """Push a sample into the outlet.

        Each entry in the list corresponds to one channel.

        Keyword arguments:
        x -- A list of values to push (one per channel).
        timestamp -- Optionally the capture time of the sample, in agreement
                     with local_clock(); if omitted, the current
                     time is used. (default 0.0)
        pushthrough -- Whether to push the sample through to the receivers
                       instead of buffering it with subsequent samples.
                       Note that the chunk_size, if specified at outlet
                       construction, takes precedence over the pushthrough flag.
                       (default True)

        """
        if len(x) == self.channel_count:
            if self.channel_format == cf_string:
                x = [v.encode('utf-8') for v in x]
            handle_error(self.do_push_sample(self.obj, self.sample_type(*x),
                                             c_double(timestamp),
                                             c_int(pushthrough)))
        else:
            raise ValueError("length of the data must correspond to the "
                             "stream's channel count.")

    def push_chunk(self, x, timestamp=0.0, pushthrough=True):
        """Push a list of samples into the outlet.

        samples -- A list of samples, either as a list of lists or a list of
                   multiplexed values.
        timestamp -- Optionally the capture time of the most recent sample, in
                     agreement with local_clock(); if omitted, the current
                     time is used. The time stamps of other samples are
                     automatically derived according to the sampling rate of
                     the stream. (default 0.0)
        pushthrough Whether to push the chunk through to the receivers instead
                    of buffering it with subsequent samples. Note that the
                    chunk_size, if specified at outlet construction, takes
                    precedence over the pushthrough flag. (default True)

        """
        try:
            n_values = self.channel_count * len(x)
            data_buff = (self.value_type * n_values).from_buffer(x)
            handle_error(self.do_push_chunk(self.obj, data_buff,
                                            c_long(n_values),
                                            c_double(timestamp),
                                            c_int(pushthrough)))
        except TypeError:
            if len(x):
                if type(x[0]) is list:
                    x = [v for sample in x for v in sample]
                if self.channel_format == cf_string:
                    x = [v.encode('utf-8') for v in x]
                if len(x) % self.channel_count == 0:
                    constructor = self.value_type*len(x)
                    # noinspection PyCallingNonCallable
                    handle_error(self.do_push_chunk(self.obj, constructor(*x),
                                                    c_long(len(x)),
                                                    c_double(timestamp),
                                                    c_int(pushthrough)))
                else:
                    raise ValueError("each sample must have the same number of "
                                     "channels.")

    def have_consumers(self):
        """Check whether consumers are currently registered.

        While it does not hurt, there is technically no reason to push samples
        if there is no consumer.

        """
        return bool(lib.lsl_have_consumers(self.obj))

    def wait_for_consumers(self, timeout):
        """Wait until some consumer shows up (without wasting resources).

        Returns True if the wait was successful, False if the timeout expired.

        """
        return bool(lib.lsl_wait_for_consumers(self.obj, c_double(timeout)))

# =========================
# === Resolve Functions ===
# =========================

def resolve_streams(wait_time=1.0):
    """Resolve all streams on the network.

    This function returns all currently available streams from any outlet on
    the network. The network is usually the subnet specified at the local
    router, but may also include a group of machines visible to each other via
    multicast packets (given that the network supports it), or list of
    hostnames. These details may optionally be customized by the experimenter
    in a configuration file (see Network Connectivity in the LSL wiki).

    Keyword arguments:
    wait_time -- The waiting time for the operation, in seconds, to search for
                 streams. Warning: If this is too short (<0.5s) only a subset
                 (or none) of the outlets that are present on the network may
                 be returned. (default 1.0)

    Returns a list of StreamInfo objects (with empty desc field), any of which
    can subsequently be used to open an inlet. The full description can be
    retrieved from the inlet.

    """
    # noinspection PyCallingNonCallable
    buffer = (c_void_p*1024)()
    num_found = lib.lsl_resolve_all(byref(buffer), 1024, c_double(wait_time))
    return [StreamInfo(handle=buffer[k]) for k in range(num_found)]

def resolve_byprop(prop, value, minimum=1, timeout=FOREVER):
    """Resolve all streams with a specific value for a given property.

    If the goal is to resolve a specific stream, this method is preferred over
    resolving all streams and then selecting the desired one.

    Keyword arguments:
    prop -- The StreamInfo property that should have a specific value (e.g.,
            "name", "type", "source_id", or "desc/manufaturer").
    value -- The string value that the property should have (e.g., "EEG" as
             the type property).
    minimum -- Return at least this many streams. (default 1)
    timeout -- Optionally a timeout of the operation, in seconds. If the
               timeout expires, less than the desired number of streams
               (possibly none) will be returned. (default FOREVER)

    Returns a list of matching StreamInfo objects (with empty desc field), any
    of which can subsequently be used to open an inlet.

    Example: results = resolve_Stream_byprop("type","EEG")

    """
    # noinspection PyCallingNonCallable
    buffer = (c_void_p*1024)()
    num_found = lib.lsl_resolve_byprop(byref(buffer), 1024,
                                       c_char_p(str.encode(prop)),
                                       c_char_p(str.encode(value)),
                                       minimum,
                                       c_double(timeout))
    return [StreamInfo(handle=buffer[k]) for k in range(num_found)]

def resolve_bypred(predicate, minimum=1, timeout=FOREVER):
    """Resolve all streams that match a given predicate.

    Advanced query that allows to impose more conditions on the retrieved
    streams; the given string is an XPath 1.0 predicate for the <description>
    node (omitting the surrounding []'s), see also
    http://en.wikipedia.org/w/index.php?title=XPath_1.0&oldid=474981951.

    Keyword arguments:
    predicate -- The predicate string, e.g. "name='BioSemi'" or
                "type='EEG' and starts-with(name,'BioSemi') and
                 count(description/desc/channels/channel)=32"
    minimum -- Return at least this many streams. (default 1)
    timeout -- Optionally a timeout of the operation, in seconds. If the
               timeout expires, less than the desired number of streams
               (possibly none) will be returned. (default FOREVER)

    Returns a list of matching StreamInfo objects (with empty desc field), any
    of which can subsequently be used to open an inlet.

    """
    # noinspection PyCallingNonCallable
    buffer = (c_void_p*1024)()
    num_found = lib.lsl_resolve_bypred(byref(buffer), 1024,
                                       c_char_p(str.encode(predicate)),
                                       minimum,
                                       c_double(timeout))
    return [StreamInfo(handle=buffer[k]) for k in range(num_found)]

# ====================
# === Memory functions
# ====================
def free_char_p_array_memory(char_p_array,num_elements):
    pointers = cast(char_p_array, POINTER(c_void_p))
    for p in range(num_elements):
        if pointers[p] is not None:  # only free initialized pointers
            lib.lsl_destroy_string(pointers[p])

# ====================
# === Stream Inlet ===
# ====================

class StreamInlet:
    """A stream inlet.

    Inlets are used to receive streaming data (and meta-data) from the lab
    network.

    """

    def __init__(self, info, max_buflen=360, max_chunklen=0, recover=True, processing_flags=0):
        """Construct a new stream inlet from a resolved stream description.

        Keyword arguments:
        description -- A resolved stream description object (as coming from one
                of the resolver functions). Note: the stream_inlet may also be
                constructed with a fully-specified stream_info, if the desired
                channel format and count is already known up-front, but this is
                strongly discouraged and should only ever be done if there is
                no time to resolve the stream up-front (e.g., due to
                limitations in the client program).
        max_buflen -- Optionally the maximum amount of data to buffer (in
                      seconds if there is a nominal sampling rate, otherwise
                      x100 in samples). Recording applications want to use a
                      fairly large buffer size here, while real-time
                      applications would only buffer as much as they need to
                      perform their next calculation. (default 360)
        max_chunklen -- Optionally the maximum size, in samples, at which
                        chunks are transmitted (the default corresponds to the
                        chunk sizes used by the sender). Recording programs
                        can use a generous size here (leaving it to the network
                        how to pack things), while real-time applications may
                        want a finer (perhaps 1-sample) granularity. If left
                        unspecified (=0), the sender determines the chunk
                        granularity. (default 0)
        recover -- Try to silently recover lost streams that are recoverable
                   (=those that that have a source_id set). In all other cases
                   (recover is False or the stream is not recoverable)
                   functions may throw a lost_error if the stream's source is
                   lost (e.g., due to an app or computer crash). (default True)

        """
        if type(info) is list:
            raise TypeError("description needs to be of type StreamInfo, "
                            "got a list.")
        self.obj = lib.lsl_create_inlet(info.obj, max_buflen, max_chunklen,
                                        recover)
        self.obj = c_void_p(self.obj)
        if not self.obj:
            raise RuntimeError("could not create stream inlet.")
        if processing_flags > 0:
            handle_error(lib.lsl_set_postprocessing(self.obj, processing_flags))
        self.channel_format = info.channel_format()
        self.channel_count = info.channel_count()
        self.do_pull_sample = fmt2pull_sample[self.channel_format]
        self.do_pull_chunk = fmt2pull_chunk[self.channel_format]
        self.value_type = fmt2type[self.channel_format]
        self.sample_type = self.value_type*self.channel_count
        self.sample = self.sample_type()
        self.buffers = {}

    def __del__(self):
        """Destructor. The inlet will automatically disconnect if destroyed."""
        # noinspection PyBroadException
        try:
            lib.lsl_destroy_inlet(self.obj)
        except:
            pass

    def info(self, timeout=FOREVER):
        """Retrieve the complete information of the given stream.

        This includes the extended description. Can be invoked at any time of
        the stream's lifetime.

        Keyword arguments:
        timeout -- Timeout of the operation. (default FOREVER)

        Throws a TimeoutError (if the timeout expires), or LostError (if the
        stream source has been lost).

        """
        errcode = c_int()
        result = lib.lsl_get_fullinfo(self.obj, c_double(timeout),
                                      byref(errcode))
        handle_error(errcode)
        return StreamInfo(handle=result)

    def open_stream(self, timeout=FOREVER):
        """Subscribe to the data stream.

        All samples pushed in at the other end from this moment onwards will be
        queued and eventually be delivered in response to pull_sample() or
        pull_chunk() calls. Pulling a sample without some preceding open_stream
        is permitted (the stream will then be opened implicitly).

        Keyword arguments:
        timeout -- Optional timeout of the operation (default FOREVER).

        Throws a TimeoutError (if the timeout expires), or LostError (if the
        stream source has been lost).

        """
        errcode = c_int()
        lib.lsl_open_stream(self.obj, c_double(timeout), byref(errcode))
        handle_error(errcode)

    def close_stream(self):
        """Drop the current data stream.

        All samples that are still buffered or in flight will be dropped and
        transmission and buffering of data for this inlet will be stopped. If
        an application stops being interested in data from a source
        (temporarily or not) but keeps the outlet alive, it should call
        lsl_close_stream() to not waste unnecessary system and network
        resources.

        """
        lib.lsl_close_stream(self.obj)

    def time_correction(self, timeout=FOREVER):
        """Retrieve an estimated time correction offset for the given stream.

        The first call to this function takes several miliseconds until a
        reliable first estimate is obtained. Subsequent calls are instantaneous
        (and rely on periodic background updates). The precision of these
        estimates should be below 1 ms (empirically within +/-0.2 ms).

        Keyword arguments:
        timeout -- Timeout to acquire the first time-correction estimate
                   (default FOREVER).

        Returns the current time correction estimate. This is the number that
        needs to be added to a time stamp that was remotely generated via
        local_clock() to map it into the local clock domain of this
        machine.

        Throws a TimeoutError (if the timeout expires), or LostError (if the
        stream source has been lost).

        """
        errcode = c_int()
        result = lib.lsl_time_correction(self.obj, c_double(timeout),
                                         byref(errcode))
        handle_error(errcode)
        return result

    def pull_sample(self, timeout=FOREVER, sample=None):
        """Pull a sample from the inlet and return it.

        Keyword arguments:
        timeout -- The timeout for this operation, if any. (default FOREVER)
                   If this is passed as 0.0, then the function returns only a
                   sample if one is buffered for immediate pickup.

        Returns a tuple (sample,timestamp) where sample is a list of channel
        values and timestamp is the capture time of the sample on the remote
        machine, or (None,None) if no new sample was available. To remap this
        time stamp to the local clock, add the value returned by
        .time_correction() to it.

        Throws a LostError if the stream source has been lost. Note that, if
        the timeout expires, no TimeoutError is thrown (because this case is
        not considered an error).

        """

        # support for the legacy API
        if type(timeout) is list:
            assign_to = timeout
            timeout = sample if type(sample) is float else 0.0
        else:
            assign_to = None

        errcode = c_int()
        timestamp = self.do_pull_sample(self.obj, byref(self.sample),
                                        self.channel_count, c_double(timeout),
                                        byref(errcode))
        handle_error(errcode)
        if timestamp:
            sample = [v for v in self.sample]
            if self.channel_format == cf_string:
                sample = [v.decode('utf-8') for v in sample]
            if assign_to is not None:
                assign_to[:] = sample
            return sample, timestamp
        else:
            return None, None

    def pull_chunk(self, timeout=0.0, max_samples=1024, dest_obj=None):
        """Pull a chunk of samples from the inlet.

        Keyword arguments:
        timeout -- The timeout of the operation; if passed as 0.0, then only
                   samples available for immediate pickup will be returned.
                   (default 0.0)
        max_samples -- Maximum number of samples to return. (default
                       1024)
        dest_obj -- A Python object that supports the buffer interface.
                    If this is provided then the dest_obj will be updated in place
                    and the samples list returned by this method will be empty.
                    It is up to the caller to trim the buffer to the appropriate
                    number of samples.
                    A numpy buffer must be order='C'
                    (default None)

        Returns a tuple (samples,timestamps) where samples is a list of samples
        (each itself a list of values), and timestamps is a list of time-stamps.

        Throws a LostError if the stream source has been lost.

        """
        # look up a pre-allocated buffer of appropriate length
        num_channels = self.channel_count
        max_values = max_samples*num_channels

        if max_samples not in self.buffers:
            # noinspection PyCallingNonCallable
            self.buffers[max_samples] = ((self.value_type*max_values)(),
                                         (c_double*max_samples)())
        if dest_obj is not None:
            data_buff = (self.value_type * max_values).from_buffer(dest_obj)
        else:
            data_buff = self.buffers[max_samples][0]
        ts_buff = self.buffers[max_samples][1]

        # read data into it
        errcode = c_int()
        # noinspection PyCallingNonCallable
        num_elements = self.do_pull_chunk(self.obj, byref(data_buff),
                                          byref(ts_buff), max_values,
                                          max_samples, c_double(timeout),
                                          byref(errcode))
        handle_error(errcode)
        # return results (note: could offer a more efficient format in the
        # future, e.g., a numpy array)
        num_samples = num_elements/num_channels
        if dest_obj is None:
            samples = [[data_buff[s*num_channels+c] for c in range(num_channels)]
                       for s in range(int(num_samples))]
            if self.channel_format == cf_string:
                samples = [[v.decode('utf-8') for v in s] for s in samples]
                free_char_p_array_memory(data_buff, max_values)
        else:
            samples = None
        timestamps = [ts_buff[s] for s in range(int(num_samples))]
        return samples, timestamps

    def samples_available(self):
        """Query whether samples are currently available for immediate pickup.

        Note that it is not a good idea to use samples_available() to determine
        whether a pull_*() call would block: to be sure, set the pull timeout
        to 0.0 or an acceptably low value. If the underlying implementation
        supports it, the value will be the number of samples available
        (otherwise it will be 1 or 0).

        """
        return lib.lsl_samples_available(self.obj)

    def was_clock_reset(self):
        """Query whether the clock was potentially reset since the last call.

        This is rarely-used function is only needed for applications that
        combine multiple time_correction values to estimate precise clock
        drift if they should tolerate cases where the source machine was
        hot-swapped or restarted.

        """
        return bool(lib.lsl_was_clock_reset(self.obj))

# ===================
# === XML Element ===
# ===================

class XMLElement:
    """A lightweight XML element tree modeling the .desc() field of StreamInfo.

    Has a name and can have multiple named children or have text content as
    value; attributes are omitted. Insider note: The interface is modeled after
    a subset of pugixml's node type and is compatible with it. See also
    http://pugixml.googlecode.com/svn/tags/latest/docs/manual/access.html for
    additional documentation.

    """

    def __init__(self, handle):
        """Construct new XML element from existing handle."""
        self.e = c_void_p(handle)

    # === Tree Navigation ===

    def first_child(self):
        """Get the first child of the element."""
        return XMLElement(lib.lsl_first_child(self.e))

    def last_child(self):
        """Get the last child of the element."""
        return XMLElement(lib.lsl_last_child(self.e))

    def child(self, name):
        """Get a child with a specified name."""
        return XMLElement(lib.lsl_child(self.e, str.encode(name)))

    def next_sibling(self, name=None):
        """Get the next sibling in the children list of the parent node.

        If a name is provided, the next sibling with the given name is returned.

        """
        if name is None:
            return XMLElement(lib.lsl_next_sibling(self.e))
        else:
            return XMLElement(lib.lsl_next_sibling_n(self.e, str.encode(name)))

    def previous_sibling(self, name=None):
        """Get the previous sibling in the children list of the parent node.

        If a name is provided, the previous sibling with the given name is
        returned.

        """
        if name is None:
            return XMLElement(lib.lsl_previous_sibling(self.e))
        else:
            return XMLElement(lib.lsl_previous_sibling_n(self.e,
                                                         str.encode(name)))

    def parent(self):
        """Get the parent node."""
        return XMLElement(lib.lsl_parent(self.e))

    # === Content Queries ===

    def empty(self):
        """Whether this node is empty."""
        return bool(lib.lsl_empty(self.e))

    def is_text(self):
        """Whether this is a text body (instead of an XML element).

        True both for plain char data and CData.

        """
        return bool(lib.lsl_is_text(self.e))

    def name(self):
        """Name of the element."""
        return lib.lsl_name(self.e).decode('utf-8')

    def value(self):
        """Value of the element."""
        return lib.lsl_value(self.e).decode('utf-8')

    def child_value(self, name=None):
        """Get child value (value of the first child that is text).

        If a name is provided, then the value of the first child with the
        given name is returned.

        """
        if name is None:
            res = lib.lsl_child_value(self.e)
        else:
            res = lib.lsl_child_value_n(self.e, str.encode(name))
        return res.decode('utf-8')

    # === Modification ===

    def append_child_value(self, name, value):
        """Append a child node with a given name, which has a (nameless)
        plain-text child with the given text value."""
        return XMLElement(lib.lsl_append_child_value(self.e,
                                                     str.encode(name),
                                                     str.encode(value)))

    def prepend_child_value(self, name, value):
        """Prepend a child node with a given name, which has a (nameless)
        plain-text child with the given text value."""
        return XMLElement(lib.lsl_prepend_child_value(self.e,
                                                      str.encode(name),
                                                      str.encode(value)))

    def set_child_value(self, name, value):
        """Set the text value of the (nameless) plain-text child of a named
        child node."""
        return XMLElement(lib.lsl_set_child_value(self.e,
                                                  str.encode(name),
                                                  str.encode(value)))

    def set_name(self, name):
        """Set the element's name. Returns False if the node is empty."""
        return bool(lib.lsl_set_name(self.e, str.encode(name)))

    def set_value(self, value):
        """Set the element's value. Returns False if the node is empty."""
        return bool(lib.lsl_set_value(self.e, str.encode(value)))

    def append_child(self, name):
        """Append a child element with the specified name."""
        return XMLElement(lib.lsl_append_child(self.e, str.encode(name)))

    def prepend_child(self, name):
        """Prepend a child element with the specified name."""
        return XMLElement(lib.lsl_prepend_child(self.e, str.encode(name)))

    def append_copy(self, elem):
        """Append a copy of the specified element as a child."""
        return XMLElement(lib.lsl_append_copy(self.e, elem.e))

    def prepend_copy(self, elem):
        """Prepend a copy of the specified element as a child."""
        return XMLElement(lib.lsl_prepend_copy(self.e, elem.e))

    def remove_child(self, rhs):
        """Remove a given child element, specified by name or as element."""
        if type(rhs) is XMLElement:
            lib.lsl_remove_child(self.e, rhs.e)
        else:
            lib.lsl_remove_child_n(self.e, rhs)

# ==========================
# === ContinuousResolver ===
# ==========================

class ContinuousResolver:
    """A convenience class resolving streams continuously in the background.

    This object can be queried at any time for the set of streams that are
    currently visible on the network.

    """

    def __init__(self, prop=None, value=None, pred=None, forget_after=5.0):
        """Construct a new continuous_resolver.

        Keyword arguments:
        forget_after -- When a stream is no longer visible on the network
                        (e.g., because it was shut down), this is the time in
                        seconds after which it is no longer reported by the
                        resolver.

        """
        if pred is not None:
            if prop is not None or value is not None:
                raise ValueError("you can only either pass the prop/value "
                                 "argument or the pred argument, but not "
                                 "both.")
            self.obj = lib.lsl_create_continuous_resolver_bypred(str.encode(pred),
                                                                 c_double(forget_after))
        elif prop is not None and value is not None:
            self.obj = lib.lsl_create_continuous_resolver_byprop(str.encode(prop),
                                                                 str.encode(value),
                                                                 c_double(forget_after))
        elif prop is not None or value is not None:
            raise ValueError("if prop is specified, then value must be "
                             "specified, too, and vice versa.")
        else:
            self.obj = lib.lsl_create_continuous_resolver(c_double(forget_after))
        self.obj = c_void_p(self.obj)
        if not self.obj:
            raise RuntimeError("could not create continuous resolver.")

    def __del__(self):
        """Destructor for the continuous resolver."""
        # noinspection PyBroadException
        try:
            lib.lsl_destroy_continuous_resolver(self.obj)
        except:
            pass

    def results(self):
        """Obtain the set of currently present streams on the network.

        Returns a list of matching StreamInfo objects (with empty desc
        field), any of which can subsequently be used to open an inlet.

        """
        # noinspection PyCallingNonCallable
        buffer = (c_void_p*1024)()
        num_found = lib.lsl_resolver_results(self.obj, byref(buffer), 1024)
        return [StreamInfo(handle=buffer[k]) for k in range(num_found)]

# =========================
# === Error Definitions ===
# =========================

# noinspection PyShadowingBuiltins
class TimeoutError(RuntimeError):
    # note: although this overrides the name of a built-in exception,
    #       this API is retained here for compatiblity with the Python 2.x
    #       version of pylsl
    pass

class LostError(RuntimeError):
    pass

class InvalidArgumentError(RuntimeError):
    pass

class InternalError(RuntimeError):
    pass

def handle_error(errcode):
    """Error handler function. Translates an error code into an exception."""
    if type(errcode) is c_int:
        errcode = errcode.value
    if errcode == 0:
        pass  # no error
    elif errcode == -1:
        raise TimeoutError("the operation failed due to a timeout.")
    elif errcode == -2:
        raise LostError("the stream has been lost.")
    elif errcode == -3:
        raise InvalidArgumentError("an argument was incorrectly specified.")
    elif errcode == -4:
        raise InternalError("an internal error has occurred.")
    elif errcode < 0:
        raise RuntimeError("an unknown error has occurred.")

# =================================================
# === Compatibility Interface for old pylsl API ===
# =================================================

# set class aliases
stream_info = StreamInfo
stream_outlet = StreamOutlet
stream_inlet = StreamInlet
xml_element = XMLElement
timeout_error = TimeoutError
lost_error = LostError
vectorf = vectord = vectorl = vectori = vectors = vectorc = vectorstr = list

def resolve_stream(*args):
    if len(args) == 0:
        return resolve_streams()
    elif type(args[0]) in [int, float]:
        return resolve_streams(args[0])
    elif type(args[0]) is str:
        if len(args) == 1:
            return resolve_bypred(args[0])
        elif type(args[1]) in [int, float]:
            return resolve_bypred(args[0], args[1])
        else:
            if len(args) == 2:
                return resolve_byprop(args[0], args[1])
            else:
                return resolve_byprop(args[0], args[1], args[2])

# ==================================
# === Module Initialization Code ===
# ==================================

# find and load library
os_name = platform.system()
bitness = 8 * struct.calcsize("P")
if os_name in ['Windows', 'Microsoft']:
    libname = 'liblsl32.dll' if bitness == 32 else 'liblsl64.dll'
elif os_name == 'Darwin':
    libname = 'liblsl32.dylib' if bitness == 32 else 'liblsl64.dylib'
elif os_name == 'Linux':
    libname = 'liblsl32.so' if bitness == 32 else 'liblsl64.so'
else:
    raise RuntimeError("unrecognized operating system:", os_name)
libpath = os.path.join(os.path.dirname(__file__), 'lib', libname)
if not os.path.isfile(libpath):
    libpath = util.find_library('lsl' + str(bitness))
if not libpath:
    raise RuntimeError("library " + libname + " was not found - make sure "
                       "that it is on the search path (e.g., in the same "
                       "folder as pylsl.py).")
lib = CDLL(libpath)

# set function return types where necessary
lib.lsl_local_clock.restype = c_double
lib.lsl_create_streaminfo.restype = c_void_p
lib.lsl_library_info.restype = c_char_p
lib.lsl_get_name.restype = c_char_p
lib.lsl_get_type.restype = c_char_p
lib.lsl_get_nominal_srate.restype = c_double
lib.lsl_get_source_id.restype = c_char_p
lib.lsl_get_created_at.restype = c_double
lib.lsl_get_uid.restype = c_char_p
lib.lsl_get_session_id.restype = c_char_p
lib.lsl_get_hostname.restype = c_char_p
lib.lsl_get_desc.restype = c_void_p
lib.lsl_get_xml.restype = c_char_p
lib.lsl_create_outlet.restype = c_void_p
lib.lsl_create_inlet.restype = c_void_p
lib.lsl_get_fullinfo.restype = c_void_p
lib.lsl_open_stream.restype = c_void_p
lib.lsl_time_correction.restype = c_double
lib.lsl_pull_sample_f.restype = c_double
lib.lsl_pull_sample_d.restype = c_double
lib.lsl_pull_sample_l.restype = c_double
lib.lsl_pull_sample_i.restype = c_double
lib.lsl_pull_sample_s.restype = c_double
lib.lsl_pull_sample_c.restype = c_double
lib.lsl_pull_sample_str.restype = c_double
lib.lsl_pull_sample_buf.restype = c_double
lib.lsl_first_child.restype = c_void_p
lib.lsl_first_child.argtypes = [c_void_p, ]
lib.lsl_last_child.restype = c_void_p
lib.lsl_last_child.argtypes = [c_void_p, ]
lib.lsl_next_sibling.restype = c_void_p
lib.lsl_next_sibling.argtypes = [c_void_p, ]
lib.lsl_previous_sibling.restype = c_void_p
lib.lsl_previous_sibling.argtypes = [c_void_p, ]
lib.lsl_parent.restype = c_void_p
lib.lsl_parent.argtypes = [c_void_p, ]
lib.lsl_child.restype = c_void_p
lib.lsl_child.argtypes = [c_void_p, c_char_p]
lib.lsl_next_sibling_n.restype = c_void_p
lib.lsl_next_sibling_n.argtypes = [c_void_p, c_char_p]
lib.lsl_previous_sibling_n.restype = c_void_p
lib.lsl_previous_sibling_n.argtypes = [c_void_p, c_char_p]
lib.lsl_name.restype = c_char_p
lib.lsl_name.argtypes = [c_void_p, ]
lib.lsl_value.restype = c_char_p
lib.lsl_value.argtypes = [c_void_p, ]
lib.lsl_child_value.restype = c_char_p
lib.lsl_child_value.argtypes = [c_void_p, ]
lib.lsl_child_value_n.restype = c_char_p
lib.lsl_child_value_n.argtypes = [c_void_p, c_char_p]
lib.lsl_append_child_value.restype = c_void_p
lib.lsl_append_child_value.argtypes = [c_void_p, c_char_p, c_char_p]
lib.lsl_prepend_child_value.restype = c_void_p
lib.lsl_prepend_child_value.argtypes = [c_void_p, c_char_p, c_char_p]
# Return type for lsl_set_child_value, lsl_set_name, lsl_set_value is int
lib.lsl_set_child_value.argtypes = [c_void_p, c_char_p, c_char_p]
lib.lsl_set_name.argtypes = [c_void_p, c_char_p]
lib.lsl_set_value.argtypes = [c_void_p, c_char_p]
lib.lsl_append_child.restype = c_void_p
lib.lsl_append_child.argtypes = [c_void_p, c_char_p]
lib.lsl_prepend_child.restype = c_void_p
lib.lsl_prepend_child.argtypes = [c_void_p, c_char_p]
lib.lsl_append_copy.restype = c_void_p
lib.lsl_append_copy.argtypes = [c_void_p, c_void_p]
lib.lsl_prepend_copy.restype = c_void_p
lib.lsl_prepend_copy.argtypes = [c_void_p, c_void_p]
lib.lsl_remove_child_n.argtypes = [c_void_p, c_char_p]
lib.lsl_remove_child.argtypes = [c_void_p, c_void_p]
lib.lsl_destroy_string.argtypes = [c_void_p]
# noinspection PyBroadException
try:
    lib.lsl_pull_chunk_f.restype = c_long
    lib.lsl_pull_chunk_d.restype = c_long
    lib.lsl_pull_chunk_l.restype = c_long
    lib.lsl_pull_chunk_i.restype = c_long
    lib.lsl_pull_chunk_s.restype = c_long
    lib.lsl_pull_chunk_c.restype = c_long
    lib.lsl_pull_chunk_str.restype = c_long
    lib.lsl_pull_chunk_buf.restype = c_long
except:
    print("pylsl: chunk transfer functions not available in your liblsl "
          "version.")
# noinspection PyBroadException
try:
    lib.lsl_create_continuous_resolver.restype = c_void_p
    lib.lsl_create_continuous_resolver_bypred.restype = c_void_p
    lib.lsl_create_continuous_resolver_byprop.restype = c_void_p
except:
    print("pylsl: ContinuousResolver not (fully) available in your liblsl "
          "version.")

# set up some type maps
string2fmt = {'float32': cf_float32, 'double64': cf_double64,
              'string': cf_string, 'int32': cf_int32, 'int16': cf_int16,
              'int8': cf_int8, 'int64': cf_int64}
fmt2string = ['undefined', 'float32', 'double64', 'string', 'int32', 'int16',
              'int8', 'int64']
fmt2type = [[], c_float, c_double, c_char_p, c_int, c_short, c_byte, c_longlong]
fmt2push_sample = [[], lib.lsl_push_sample_ftp, lib.lsl_push_sample_dtp,
                   lib.lsl_push_sample_strtp, lib.lsl_push_sample_itp,
                   lib.lsl_push_sample_stp, lib.lsl_push_sample_ctp, []]
fmt2pull_sample = [[], lib.lsl_pull_sample_f, lib.lsl_pull_sample_d,
                   lib.lsl_pull_sample_str, lib.lsl_pull_sample_i,
                   lib.lsl_pull_sample_s, lib.lsl_pull_sample_c, []]
# noinspection PyBroadException
try:
    fmt2push_chunk = [[], lib.lsl_push_chunk_ftp, lib.lsl_push_chunk_dtp,
                      lib.lsl_push_chunk_strtp, lib.lsl_push_chunk_itp,
                      lib.lsl_push_chunk_stp, lib.lsl_push_chunk_ctp, []]
    fmt2pull_chunk = [[], lib.lsl_pull_chunk_f, lib.lsl_pull_chunk_d,
                      lib.lsl_pull_chunk_str, lib.lsl_pull_chunk_i,
                      lib.lsl_pull_chunk_s, lib.lsl_pull_chunk_c, []]
except:
    # if not available
    fmt2push_chunk = [None, None, None, None, None, None, None, None]
    fmt2pull_chunk = [None, None, None, None, None, None, None, None]


from addok.db import DB
from addok.helpers import keys, magenta, white
from addok.helpers.search import preprocess_query

def pair_key(s):
    return 'p|{}'.format(s)

class PairsIndexer:

    @staticmethod
    def index(pipe, key, doc, tokens, **kwargs):
        for token in list(set(tokens.keys())):  # Unique values.
            pairs = set(t for t in tokens if t != token)
            if pairs:
                pipe.sadd(pair_key(token), *pairs)

    @staticmethod
    def deindex(db, key, doc, tokens, **kwargs):
        tokens = list(set(tokens))  # Unique values.
        for i, token in enumerate(tokens):
            for token2 in tokens[i:]:
                if token != token2:
                    tmp_key = '|'.join(['didx', token, token2])
                    # Do we have other documents that share token and token2?
                    commons = db.zinterstore(tmp_key, [keys.token_key(token),
                                                       keys.token_key(token2)])
                    db.delete(tmp_key)
                    if not commons:
                        db.srem(pair_key(token), token2)
                        db.srem(pair_key(token2), token)

def pair(cmd, word):
    """See all token associated with a given token.
    PAIR lilas"""
    word = list(preprocess_query(word))[0]
    key = pair_key(word)
    tokens = [t.decode() for t in DB.smembers(key)]
    tokens.sort()
    print(white(tokens))
    print(magenta('(Total: {})'.format(len(tokens))))

def register_shell_command(cmd):
    cmd.register_command(pair)


import redis

from addok.config import config
from addok.db import DB
from addok.helpers import keys as dbkeys
from addok.helpers import magenta, parallelize, white
from addok.helpers.index import token_key_frequency
from addok.helpers.search import preprocess_query
from addok.helpers.text import compute_edge_ngrams
from addok.pairs import pair_key

def edge_ngram_key(s):
    return 'n|{}'.format(s)

def index_edge_ngrams(pipe, token):
    for ngram in compute_edge_ngrams(token):
        pipe.sadd(edge_ngram_key(ngram), token)

def deindex_edge_ngrams(token):
    for ngram in compute_edge_ngrams(token):
        DB.srem(edge_ngram_key(ngram), token)

class EdgeNgramIndexer:

    @staticmethod
    def index(pipe, key, doc, tokens, **kwargs):
        if config.INDEX_EDGE_NGRAMS:  # Allow to disable for mass indexing.
            for token in tokens.keys():
                index_edge_ngrams(pipe, token)

    @staticmethod
    def deindex(db, key, doc, tokens, **kwargs):
        if config.INDEX_EDGE_NGRAMS:
            for token in tokens:
                tkey = dbkeys.token_key(token)
                if not DB.exists(tkey):
                    deindex_edge_ngrams(token)

def only_commons_but_geohash_try_autocomplete_collector(helper):
    if helper.geohash_key and helper.only_commons:
        autocomplete(helper, helper.tokens, use_geohash=True)

def only_commons_try_autocomplete_collector(helper):
    if helper.only_commons:
        autocomplete(helper, helper.tokens, skip_commons=True)
        if not helper.bucket_empty:
            helper.debug('Only common terms. Return.')
            return True

def no_meaningful_but_common_try_autocomplete_collector(helper):
    if not helper.meaningful and helper.common:
        # Only commons terms, try to reduce with autocomplete.
        helper.debug('Only commons, trying autocomplete')
        autocomplete(helper, helper.common)
        helper.meaningful = helper.common[:1]
        if not helper.pass_should_match_threshold:
            return False
        if helper.bucket_full or helper.bucket_overflow or helper.has_cream():
            return True

def autocomplete_meaningful_collector(helper):
    if helper.bucket_overflow:
        return
    if not helper.autocomplete:
        helper.debug('Autocomplete not active. Abort.')
        return
    if helper.geohash_key:
        autocomplete(helper, helper.meaningful, use_geohash=True)
    autocomplete(helper, helper.meaningful)

def autocomplete(helper, tokens, skip_commons=False, use_geohash=False):
    helper.debug('Autocompleting %s', helper.last_token)
    keys = [t.db_key for t in tokens if not t.is_last]
    pair_keys = [pair_key(t) for t in tokens if not t.is_last]
    key = edge_ngram_key(helper.last_token)
    autocomplete_tokens = DB.sinter(pair_keys + [key])
    helper.debug('Found tokens to autocomplete %s', autocomplete_tokens)
    for token in autocomplete_tokens:
        key = dbkeys.token_key(token.decode())
        if skip_commons\
           and token_key_frequency(key) > config.COMMON_THRESHOLD:
            helper.debug('Skip common token to autocomplete %s', key)
            continue
        if not helper.bucket_overflow or helper.last_token in helper.not_found:
            helper.debug('Trying to extend bucket. Autocomplete %s', key)
            extra_keys = [key]
            if use_geohash and helper.geohash_key:
                extra_keys.append(helper.geohash_key)
            helper.add_to_bucket(keys + extra_keys)

def index_ngram_keys(*keys):
    pipe = DB.pipeline(transaction=False)
    for key in keys:
        key = key.decode()
        _, token = key.split('|')
        if token.isdigit():
            continue
        index_edge_ngrams(pipe, token)
    try:
        pipe.execute()
    except redis.RedisError as e:
        msg = 'Error while generating ngrams:\n{}'.format(str(e))
        raise ValueError(msg)
    return keys

def create_edge_ngrams(*args):
    pattern = '{}*'.format(dbkeys.TOKEN_PREFIX)
    parallelize(index_ngram_keys, DB.scan_iter(match=pattern),
                chunk_size=10000, throttle=1000)

def register_command(subparsers):
    parser = subparsers.add_parser('ngrams', help='Create edge ngrams.')
    parser.set_defaults(func=create_edge_ngrams)

def do_AUTOCOMPLETE(cmd, s):
    """Shows autocomplete results for a given token."""
    s = list(preprocess_query(s))[0]
    keys = [k.decode() for k in DB.smembers(edge_ngram_key(s))]
    print(white(keys))
    print(magenta('({} elements)'.format(len(keys))))

def register_shell_command(cmd):
    cmd.register_command(do_AUTOCOMPLETE)


import re
from functools import lru_cache
from pathlib import Path

from unidecode import unidecode

from addok.config import config
from addok.db import DB
from addok.helpers import keys, yielder
from addok.helpers.index import token_frequency

import editdistance

PATTERN = re.compile(r"[\w]+", re.U | re.X)

class Token(str):

    __slots__ = ('_positions', 'is_last', 'db_key', 'raw', '_frequency',
                 '_key', 'kind', 'is_first')

    def __new__(cls, value, position=None, is_last=False, raw=None, kind=None):
        obj = str.__new__(cls, value)
        obj._positions = []
        position = position or 0
        obj.position = position
        obj.is_first = obj.position[0] == 0
        obj.is_last = is_last
        obj.db_key = None
        obj.raw = raw or value  # Allow to keep raw on update.
        obj.kind = kind
        return obj

    def __repr__(self):
        return '<Token {}>'.format(self)

    def update(self, value, **kwargs):
        default = dict(is_last=self.is_last, raw=self.raw, kind=self.kind,
                       position=self.position[:])
        # Never replace position through `update`.
        position = kwargs.pop('position', None)
        default.update(kwargs)
        token = Token(value=value, **default)
        if position is not None:
            token.position = position
        return token

    def search(self):
        if DB.exists(self.key):
            self.db_key = self.key

    @property
    def is_common(self):
        return self.frequency > config.COMMON_THRESHOLD

    @property
    def frequency(self):
        if not hasattr(self, '_frequency'):
            self._frequency = token_frequency(self)
        return self._frequency

    @property
    def key(self):
        if not hasattr(self, '_key'):
            self._key = keys.token_key(self)
        return self._key

    @property
    def position(self):
        # Allow to store multiple positions when a token is subdivided.
        return self._positions

    @position.setter
    def position(self, position):
        if isinstance(position, list):
            self._positions = position
        else:
            self._positions.append(position)

def _tokenize(text):
    """Split text into a list of tokens."""
    return PATTERN.findall(text)

def tokenize(pipe):
    for text in pipe:
        for position, token in enumerate(_tokenize(text)):
            yield Token(token, position=position)

def _normalize(s):
    return s.update(unidecode(s.lower()))
normalize = yielder(_normalize)

@config.on_load
def load_synonyms():
    config.SYNONYMS = {}
    path = config.SYNONYMS_PATH
    if not path:
        return  # pragma: no cover
    with Path(path).open() as f:
        for line in f:
            if line.startswith('#'):
                continue
            synonyms, wanted = line.split('=>')
            wanted = wanted.strip()
            synonyms = synonyms.split(',')
            for synonym in synonyms:
                synonym = synonym.strip()
                if not synonym:
                    continue
                config.SYNONYMS[synonym] = wanted

def synonymize(tokens):
    for token in tokens:
        for position, subtoken in enumerate(config.SYNONYMS.get(token, token).split()):
            yield token.update(subtoken, position=position)

class ascii(str):
    """Just like a str, but ascii folded and cached."""

    __slots__ = ['_cache', '_raw']

    def __new__(cls, value):
        try:
            cache = value._cache
        except AttributeError:
            cache = alphanumerize(unidecode(value.lower()))
            obj = str.__new__(cls, cache)
            obj._cache = cache
            obj._raw = getattr(value, '_raw', value)
            return obj
        else:
            return value

    def __str__(self):
        return self._raw

@lru_cache(maxsize=512)
def ngrams(text, n=3):
    # Make sure strings are at least 3 chars long, and a given token will have
    # same ngrams whether at the start, the middle or the end of the string.
    text = ' ' + text + ' '
    return set([text[i:i+n] for i in range(0, len(text)-(n-1))])

def compare_str(left, right):
    left = ascii(left)
    right = ascii(right)
    left_n = ngrams(left)
    right_n = ngrams(right)
    # Evaluate editdistance limited to common text portion.
    distance = ((editdistance.eval(left, right) - abs(len(left)-len(right)))
                / max(len(left), len(right)))
    return (len(left_n & right_n) / len(right_n) * 0.85
            + len(left_n & right_n) / len(left_n) * 0.05
            + (1 - distance) * 0.1)

def contains(candidate, target):
    candidate = ascii(candidate)
    target = ascii(target)
    return candidate in target

def startswith(candidate, target):
    candidate = ascii(candidate)
    target = ascii(target)
    return target.startswith(candidate)

def equals(candidate, target):
    candidate = ascii(candidate)
    target = ascii(target)
    return target == candidate

def alphanumerize(text):
    return re.sub(' {2,}', ' ', re.sub('[^\w]', ' ', text))

def compute_edge_ngrams(token, min=None):
    """Compute edge ngram of token from min. Does not include token itself."""
    if min is None:
        min = config.MIN_EDGE_NGRAMS
    token = token[:config.MAX_EDGE_NGRAMS + 1]
    return [token[:i] for i in range(min, len(token))]

class EntityTooLarge(ValueError):
    pass

@yielder
def check_query_length(q):
    if len(q) > config.QUERY_MAX_LENGTH:
        raise EntityTooLarge('Query too long, {} chars, limit is {}'.format(
            len(q), config.QUERY_MAX_LENGTH))
    return q

@yielder
def flag_housenumber(token):
    """Very basic housenumber flagging. Make your own for your specific needs.
    Eg. in addok-france:
    https://github.com/addok/addok-france/blob/master/addok_france/utils.py#L106
    """
    if token.is_first and token.isdigit():
        token.kind = 'housenumber'
    return token


import csv
from functools import wraps
from importlib import import_module
from math import asin, cos, exp, radians, sin, sqrt
from multiprocessing.pool import RUN, IMapUnorderedIterator, Pool

from progressist import ProgressBar

from addok.config import config

def load_file(filepath):
    with open(filepath) as f:
        for line in f:
            yield line

def load_msgpack_file(filepath):
    import msgpack  # We don't want to make it a required dependency.
    with open(filepath, mode='rb') as f:
        for line in msgpack.Unpacker(f, encoding='utf-8'):
            yield line

def load_csv_file(filepath):
    with open(filepath, newline='') as csvfile:
        dialect = csv.Sniffer().sniff(csvfile.read(1024))
        csvfile.seek(0)
        reader = csv.DictReader(csvfile, dialect=dialect)
        yield from reader

def iter_pipe(pipe, processors):
    """Allow for iterators to return either an item or an iterator of items."""
    if isinstance(pipe, str):
        pipe = [pipe]
    for it in processors:
        pipe = it(pipe)
    yield from pipe

def import_by_path(path):
    """
    Import functions or class by their path. Should be of the form:
    path.to.module.func
    """
    if not isinstance(path, str):
        return path
    module_path, *name = path.rsplit('.', 1)
    func = import_module(module_path)
    if name:
        func = getattr(func, name[0])
    return func

def yielder(func):
    @wraps(func)
    def wrapper(pipe):
        for item in pipe:
            yield func(item)
    return wrapper

def haversine_distance(point1, point2):
    """
    Calculate the great circle distance between two points
    on the earth (specified in decimal degrees).
    """
    lat1, lon1 = point1
    lat2, lon2 = point2

    # Convert decimal degrees to radians.
    lon1, lat1, lon2, lat2 = map(radians, [lon1, lat1, lon2, lat2])

    # Haversine formula.
    dlon = lon2 - lon1
    dlat = lat2 - lat1
    a = sin(dlat/2)**2 + cos(lat1) * cos(lat2) * sin(dlon/2)**2
    c = 2 * asin(sqrt(a))

    # 6367 km is the radius of the Earth.
    km = 6367 * c
    return km

def km_to_score(km):
    # Score between 0 and 0.1 (close to 0 km will be close to 0.1, and 100 and
    # above will be 0).
    return 0.0 if km > 100 else exp(-(km / 50.0) ** 2)

COLORS = {
    'red': '31',
    'green': '32',
    'yellow': '33',
    'blue': '34',
    'magenta': '35',
    'cyan': '36',
    'white': '37',
    'reset': '39'
}

def colorText(s, color):
    # color should be a string from COLORS
    return '\033[%sm%s\033[%sm' % (COLORS[color], s, COLORS['reset'])

def red(s):
    return colorText(s, 'red')

def green(s):
    return colorText(s, 'green')

def yellow(s):
    return colorText(s, 'yellow')

def blue(s):
    return colorText(s, 'blue')

def magenta(s):
    return colorText(s, 'magenta')

def cyan(s):
    return colorText(s, 'cyan')

def white(s):
    return colorText(s, 'white')

class Bar(ProgressBar):
    animation = '{spinner}'
    template = '{prefix} {animation} Done: {done} | Elapsed: {elapsed}'

class ChunkedPool(Pool):

    def imap_unordered(self, func, iterable, chunksize):
        """Customized version of imap_unordered.

        Directly send chunks to func, instead of iterating in each process and
        sending one by one.

        Original:
        https://hg.python.org/cpython/file/tip/Lib/multiprocessing/pool.py#l271

        Other tried options:
        - map_async: makes a list(iterable), so it loads all the data for each
          process into RAM
        - apply_async: needs manual chunking
        """
        assert self._state == RUN
        task_batches = Pool._get_tasks(func, iterable, chunksize)
        result = IMapUnorderedIterator(self._cache)
        tasks = ((result._job, i, func, chunk, {})
                 for i, (_, chunk) in enumerate(task_batches))
        self._taskqueue.put((tasks, result._set_length))
        return result

def parallelize(func, iterable, chunk_size, **bar_kwargs):
    bar = Bar(prefix='Processing…', **bar_kwargs)
    with ChunkedPool(processes=config.BATCH_WORKERS) as pool:
        for chunk in pool.imap_unordered(func, iterable, chunk_size):
            bar(step=len(chunk))
        bar.finish()


import string

from addok.db import DB
from addok.helpers import keys as dbkeys
from addok.helpers import blue, white
from addok.helpers.search import preprocess_query
from addok.helpers.text import Token
from addok.pairs import pair_key

def make_fuzzy(word, max=1):
    """Naive neighborhoods algo."""
    # inversions
    neighbors = []
    for i in range(0, len(word) - 1):
        neighbor = list(word)
        neighbor[i], neighbor[i+1] = neighbor[i+1], neighbor[i]
        neighbors.append(''.join(neighbor))
    # substitutions
    for letter in string.ascii_lowercase:
        for i in range(0, len(word)):
            neighbor = list(word)
            if letter != neighbor[i]:
                neighbor[i] = letter
                neighbors.append(''.join(neighbor))
    # insertions
    for letter in string.ascii_lowercase:
        for i in range(0, len(word) + 1):
            neighbor = list(word)
            neighbor.insert(i, letter)
            neighbors.append(''.join(neighbor))
    if len(word) > 3:
        # removal
        for i in range(0, len(word)):
            neighbor = list(word)
            del neighbor[i]
            neighbors.append(''.join(neighbor))
    return neighbors

def fuzzy_collector(helper):
    if helper.fuzzy and not helper.has_cream():
        if helper.not_found:
            try_fuzzy(helper, helper.not_found)
        if helper.bucket_dry and not helper.has_cream():
            try_fuzzy(helper, helper.meaningful)
        if helper.bucket_dry and not helper.has_cream() and helper.common:
            try_fuzzy(helper, helper.meaningful, include_common=False)

def try_fuzzy(helper, tokens, include_common=True):
    if not helper.bucket_dry or not tokens:
        return
    helper.debug('Fuzzy on. Trying with %s.', tokens)
    tokens.sort(key=lambda t: len(t), reverse=True)
    allkeys = helper.keys[:]
    if include_common:
        # As we are in fuzzy, try to narrow as much as possible by adding
        # unused common tokens.
        allkeys.extend([t.db_key for t in helper.common
                        if t.db_key not in helper.keys])
    for try_one in tokens:
        if helper.bucket_full:
            break
        keys = allkeys[:]
        if try_one.db_key in keys:
            keys.remove(try_one.db_key)
        if try_one.isdigit():
            continue
        helper.debug('Going fuzzy with %s and %s', try_one, keys)
        neighbors = make_fuzzy(try_one, max=helper.fuzzy)
        if len(keys):
            # Only retain tokens that have been seen in the index at least
            # once with the other tokens.
            DB.sadd(helper.pid, *neighbors)
            interkeys = [pair_key(k[2:]) for k in keys]
            interkeys.append(helper.pid)
            fuzzy_words = DB.sinter(interkeys)
            DB.delete(helper.pid)
            # Keep the priority we gave in building fuzzy terms (inversion
            # first, then substitution, etc.).
            fuzzy_words = [w.decode() for w in fuzzy_words]
            fuzzy_words.sort(key=lambda x: neighbors.index(x))
        else:
            # The token we are considering is alone.
            fuzzy_words = []
            for neighbor in neighbors:
                key = dbkeys.token_key(neighbor)
                count = DB.zcard(key)
                if count:
                    fuzzy_words.append(neighbor)
        if fuzzy_words:
            helper.debug('Found fuzzy candidates %s', fuzzy_words)
            fuzzy_keys = [dbkeys.token_key(w) for w in fuzzy_words]
            for key in fuzzy_keys:
                if helper.bucket_dry:
                    helper.add_to_bucket(keys + [key])

def do_fuzzy(self, word):
    """Compute fuzzy extensions of word.
    FUZZY lilas"""
    word = list(preprocess_query(word))[0]
    print(white(make_fuzzy(word)))

def do_fuzzyindex(self, word):
    """Compute fuzzy extensions of word that exist in index.
    FUZZYINDEX lilas"""
    word = list(preprocess_query(word))[0]
    token = Token(word)
    neighbors = make_fuzzy(token)
    neighbors = [(n, DB.zcard(dbkeys.token_key(n))) for n in neighbors]
    neighbors.sort(key=lambda n: n[1], reverse=True)
    for token, freq in neighbors:
        if freq == 0:
            break
        print(white(token), blue(freq))

def register_shell_command(cmd):
    cmd.register_commands(do_fuzzy, do_fuzzyindex)


from collections import defaultdict
from itertools import product

from addok.config import config
from addok.db import DB
from addok.helpers import scripts
from addok.pairs import pair_key

def no_tokens_but_housenumbers_and_geohash(helper):
    if not helper.tokens and helper.housenumbers and helper.geohash_key:
        helper.new_bucket([helper.geohash_key], config.BUCKET_MIN)

def no_available_tokens_abort(helper):
    if not helper.tokens:
        return True  # Stop processing.

def only_commons(helper):
    if helper.only_commons:
        # Only common terms, shortcut to search
        keys = [t.db_key for t in helper.tokens]
        if helper.geohash_key:
            keys.append(helper.geohash_key)
            helper.debug('Adding geohash %s', helper.geohash_key)
        if len(keys) == 1 or helper.geohash_key:
            helper.add_to_bucket(keys)
        if helper.bucket_dry and len(keys) > 1:
            # Scan the less frequent token.
            helper.tokens.sort(key=lambda t: t.frequency)
            keys = [t.db_key for t in helper.tokens]
            first = helper.tokens[0]
            if first.frequency < config.INTERSECT_LIMIT:
                helper.debug('Under INTERSECT_LIMIT, force intersect.')
                helper.add_to_bucket(keys)
            else:
                helper.debug('INTERSECT_LIMIT hit, manual scan')
                if helper.filters:
                    # Always consider filters when doing manual intersect.
                    keys = keys + helper.filters
                    # But, hey, can we brute force again?
                    if any(DB.scard(k) < config.INTERSECT_LIMIT
                           for k in helper.filters):
                        helper.debug('Filters under INTERSECT_LIMIT, force')
                        helper.add_to_bucket(keys)
                        return
                helper.debug('manual scan on "%s"', first)
                ids = scripts.manual_scan(keys=keys, args=[helper.wanted])
                helper.bucket.update(ids)
                helper.debug('%s results after scan', len(helper.bucket))

def bucket_with_meaningful(helper):
    if not helper.meaningful:
        return
    if len(helper.meaningful) == 1 and helper.common and not helper.filters:
        # Avoid running with too less tokens while having commons terms.
        for token in helper.common:
            if token not in helper.meaningful:
                helper.meaningful.append(token)
                break  # We want only one more.
    helper.keys = [t.db_key for t in helper.meaningful]
    if helper.bucket_empty:
        helper.new_bucket(helper.keys, config.BUCKET_MIN)
        if len(helper.bucket) == config.BUCKET_MIN:
            # Do not rerun if bucket with limit 10 has returned less
            # than 10 results.
            helper.new_bucket(helper.keys)
        if (not helper.autocomplete and helper.has_cream() and
                helper.cream < config.BUCKET_MIN):
            # Do not check cream before computing autocomplete when
            # autocomplete is on.
            # If we have too much cream, do not consider our bucket is good.
            helper.debug('Cream found. Returning.')
            return True
    else:
        helper.add_to_bucket(helper.keys)

def reduce_with_other_commons(helper):
    if helper.only_commons:
        return
    for token in helper.common:  # Already ordered by frequency asc.
        if token not in helper.meaningful and helper.bucket_overflow:
            helper.debug('Now considering also common token %s', token)
            helper.meaningful.append(token)
            helper.keys = [t.db_key for t in helper.meaningful]
            helper.new_bucket(helper.keys)

def ensure_geohash_results_are_included_if_center_is_given(helper):
    if helper.bucket_overflow and helper.geohash_key:
        helper.debug('Bucket overflow and center, force nearby look up')
        helper.add_to_bucket(helper.keys + [helper.geohash_key],
                             max(helper.wanted, 10))

def extend_results_reducing_tokens(helper):
    if helper.bucket_full or helper.has_cream():
        return True
    if not helper.bucket_dry:
        return  # No need.
    # Only if bucket is empty or we have margin on should_match_threshold.
    if (helper.bucket_empty
            or len(helper.meaningful) - 1 > helper.should_match_threshold):
        helper.debug('Bucket dry. Trying to remove 1 meaningful token.')

        def sorter(t):
            # First numbers, then by frequency
            return (2 if t.isdigit() else 1, t.frequency)

        helper.meaningful.sort(key=sorter, reverse=True)
        for token in helper.meaningful:
            keys = helper.keys[:]
            keys.remove(token.db_key)
            helper.add_to_bucket(keys)
            if helper.bucket_overflow:
                break

        if helper.bucket_empty and len(helper.meaningful) > 3:
            helper.debug("Bucket still empty, remove 2 meaningful tokens.")
            for token, token2 in product(helper.meaningful, helper.meaningful):
                if token != token2:
                    keys = helper.keys[:]
                    keys.remove(token.db_key)
                    keys.remove(token2.db_key)
                    helper.add_to_bucket(keys)
                    if helper.bucket_overflow:
                        break

def extend_results_extrapoling_relations(helper):
    """Try to extract the bigger group of interlinked tokens.

    Should generally be used at last in the collectors chain.
    """
    if not helper.bucket_dry:
        return  # No need.
    tokens = set(helper.meaningful + helper.common)
    for relation in _extract_manytomany_relations(tokens):
        helper.add_to_bucket([t.db_key for t in relation])
        if helper.bucket_overflow:
            break
    else:
        helper.debug('No relation extrapolated.')

def _extract_manytomany_relations(tokens):
    o2m_relations = _compute_onetomany_relations(tokens)
    m2m_relations = _extrapolate_manytomany_relations(o2m_relations)
    return _deduplicate_sets(m2m_relations)

def _compute_onetomany_relations(tokens):
    relations = defaultdict(list)
    for token in tokens:
        for other in tokens:
            if other == token:
                continue
            if (token in relations[other]
                    or DB.sismember(pair_key(token), other)):
                relations[token].append(other)
    return relations

def _extrapolate_manytomany_relations(o2m_relations):
    m2m_relations = []
    for origin, others in o2m_relations.items():
        relation = [origin]
        for token in others:
            if all(token in o2m_relations[o] for o in relation):
                relation.append(token)
        # Commons tokens, given that they are "common", are more luckily to
        # create false positives (but we kept them until now because if some
        # other token is not related to some common token, that is a clear true
        # negative).
        relation = set([token for token in relation if not token.is_common])
        if len(relation) > 1:
            m2m_relations.append(relation)
    return m2m_relations

def _deduplicate_sets(sets):
    unique = []
    for set_ in sorted(sets, key=len, reverse=True):
        if not any(set_.issubset(g) for g in unique):
            unique.append(set_)
    return unique