#!/usr/bin/env python
# -*- coding: utf-8 -*-
# File: remote.py
# Author: Yuxin Wu <ppwwyyxxc@gmail.com>

import time
from collections import deque
from .base import DataFlow
from ..utils import logger, get_tqdm
from ..utils.serialize import dumps, loads, dumps_for_tfop
try:
    import zmq
except ImportError:
    logger.warn("Error in 'import zmq'. remote feature won't be available")
    __all__ = []
else:
    __all__ = ['send_dataflow_zmq', 'RemoteDataZMQ']


def send_dataflow_zmq(df, addr, hwm=50, print_interval=100, format='msgpack'):
    """
    Run DataFlow and send data to a ZMQ socket addr.
    It will dump and send each datapoint to this addr with a PUSH socket.

    Args:
        df (DataFlow): Will infinitely loop over the DataFlow.
        addr: a ZMQ socket addr.
        hwm (int): high water mark
        format (str): The serialization format.
            'msgpack' is the default format corresponding to RemoteDataZMQ.
            Otherwise will use the format corresponding to the ZMQRecv TensorFlow Op.
    """
    dump_fn = dumps if format == 'msgpack' else dumps_for_tfop
    ctx = zmq.Context()
    socket = ctx.socket(zmq.PUSH)
    socket.set_hwm(hwm)
    socket.connect(addr)
    try:
        df.reset_state()
        logger.info("Serving data to {} ...".format(addr))
        q = deque(maxlen=print_interval)
        with get_tqdm(total=0) as pbar:
            while True:
                for dp in df.get_data():
                    start = time.time()
                    socket.send(dump_fn(dp), copy=False)
                    q.append(time.time() - start)
                    pbar.update(1)
                    if pbar.n % print_interval == 0:
                        pbar.write("Avg send time @{}: {}".format(pbar.n, sum(q) / len(q)))
    finally:
        socket.setsockopt(zmq.LINGER, 0)
        socket.close()
        if not ctx.closed:
            ctx.destroy(0)


class RemoteDataZMQ(DataFlow):
    """
    Produce data from ZMQ PULL socket(s).

    Attributes:
        cnt1, cnt2 (int): number of data points received from addr1 and addr2
    """
    def __init__(self, addr1, addr2=None):
        """
        Args:
            addr1,addr2 (str): addr of the socket to connect to.
                Use both if you need two protocols (e.g. both IPC and TCP).
                I don't think you'll ever need 3.
        """
        assert addr1
        self._addr1 = addr1
        self._addr2 = addr2

    def reset_state(self):
        self.cnt1 = 0
        self.cnt2 = 0

    def get_data(self):
        try:
            ctx = zmq.Context()
            if self._addr2 is None:
                socket = ctx.socket(zmq.PULL)
                socket.set_hwm(50)
                socket.bind(self._addr1)

                while True:
                    dp = loads(socket.recv(copy=False).bytes)
                    yield dp
                    self.cnt1 += 1
            else:
                socket1 = ctx.socket(zmq.PULL)
                socket1.set_hwm(50)
                socket1.bind(self._addr1)

                socket2 = ctx.socket(zmq.PULL)
                socket2.set_hwm(50)
                socket2.bind(self._addr2)

                poller = zmq.Poller()
                poller.register(socket1, zmq.POLLIN)
                poller.register(socket2, zmq.POLLIN)

                while True:
                    evts = poller.poll()
                    for sock, evt in evts:
                        dp = loads(sock.recv(copy=False).bytes)
                        yield dp
                        if sock == socket1:
                            self.cnt1 += 1
                        else:
                            self.cnt2 += 1
        finally:
            ctx.destroy(linger=0)


if __name__ == '__main__':
    from argparse import ArgumentParser
    from .raw import FakeData
    from .common import TestDataSpeed

    """
    Test the multi-producer single-consumer model
    """
    parser = ArgumentParser()
    parser.add_argument('-t', '--task', choices=['send', 'recv'], required=True)
    parser.add_argument('-a', '--addr1', required=True)
    parser.add_argument('-b', '--addr2', default=None)
    args = parser.parse_args()

    # tcp addr like "tcp://127.0.0.1:8877"
    # ipc addr like "ipc:///tmp/ipc-test"
    if args.task == 'send':
        # use random=True to make it slow and cpu-consuming
        ds = FakeData([(128, 244, 244, 3)], 1000, random=True)
        send_dataflow_zmq(ds, args.addr1)
    else:
        ds = RemoteDataZMQ(args.addr1, args.addr2)
        logger.info("Each DP is 73.5MB")
        TestDataSpeed(ds).start_test()
