__author__ = 'breddels'
import numpy as np
import logging
import threading
import uuid
import time
import ast
from .dataset import Dataset
from .utils import _issequence
from .tasks import Task
from .legacy import Subspace
import vaex.promise
import vaex.settings
import vaex.utils
from tornado.httpclient import AsyncHTTPClient, HTTPClient
import tornado.httputil
import tornado.websocket
from tornado.concurrent import Future
from tornado import gen
from .dataset import default_shape
import tornado.ioloop
import json
import astropy.units
from vaex.utils import _ensure_strings_from_expressions, _ensure_string_from_expression

try:
    import __builtin__
except ImportError:
    import builtins as __builtin__

try:
    import Cookie  # py2
except ImportError:
    import http.cookies as Cookie  # py3

logger = logging.getLogger("vaex.remote")

DEFAULT_REQUEST_TIMEOUT = 60 * 5  # 5 minutes

def wrap_future_with_promise(future):
    if isinstance(future, vaex.promise.Promise):  # TODO: not so nice, sometimes we pass a promise
        return future
    promise = vaex.promise.Promise()

    def callback(future):
        # print(("callback", future, future.result()))
        e = future.exception()
        if e:
            # print(("reject", e))
            promise.reject(e)
        else:
            promise.fulfill(future.result())
    future.add_done_callback(callback)
    return promise

def listify(value):
    if isinstance(value, list):
        value = list([listify(item) for item in value])
    if hasattr(value, "tolist"):
        value = value.tolist()
    return value

try:
    from urllib.request import urlopen
    from urllib.parse import urlparse, urlencode
except ImportError:
    from urlparse import urlparse
    from urllib import urlopen, urlencode

def _check_error(object):
    if "error" in object:
        raise RuntimeError("Server responded with error: %r" % object["error"])

class ServerRest(object):
    def __init__(self, hostname, port=5000, base_path="/", background=False, thread_mover=None, websocket=True):
        self.hostname = hostname
        self.port = port
        self.base_path = base_path if base_path.endswith("/") else (base_path + "/")
        # if delay:
        event = threading.Event()
        self.thread_mover = thread_mover or (lambda fn, *args, **kwargs: fn(*args, **kwargs))
        logger.debug("thread mover: %r", self.thread_mover)

        # jobs maps from uid to tasks
        self.jobs = {}

        def ioloop_threaded():
            logger.debug("creating tornado io_loop")
            self.io_loop = tornado.ioloop.IOLoop().current()
            event.set()
            logger.debug("started tornado io_loop...")

            self.io_loop.start()
            self.io_loop.close()
            logger.debug("stopped tornado io_loop")

        io_loop = tornado.ioloop.IOLoop.current(instance=False)
        if True:  # io_loop:# is None:
            logger.debug("no current io loop, starting it in thread")
            thread = threading.Thread(target=ioloop_threaded)
            thread.setDaemon(True)
            thread.start()
            event.wait()
        else:
            logger.debug("using current io loop")
            self.io_loop = io_loop

        self.io_loop.make_current()
        # if async:
        self.http_client_async = AsyncHTTPClient()
        self.http_client = HTTPClient()
        self.user_id = vaex.settings.webclient.get("cookie.user_id")
        self.use_websocket = websocket
        self.websocket = None

        self.submit = self.submit_http

        if self.use_websocket:
            logger.debug("connect via websocket")
            self.submit = self.submit_websocket
            self._websocket_connect()
            logger.debug("websocket connected")

    def close(self):
        # self.http_client_async.
        self.http_client.close()
        if self.use_websocket:
            if self.websocket:
                self.websocket.close()
        self.io_loop.stop()

    def _websocket_connect(self):
        def connected(websocket):
            logger.debug("connected to websocket: %s" % self._build_url(""))

        def failed(reason):
            logger.error("failed to connect to %s" % self._build_url(""))
        self.websocket_connected = vaex.promise.Promise()
        self.websocket_connected.then(connected, failed)
        if 0:
            connected = wrap_future_with_promise(tornado.websocket.websocket_connect(self._build_url("websocket"), on_message_callback=self._on_websocket_message))
            connected.get()
            self.websocket_connected.fulfill(connected)

        def do():
            try:
                logger.debug("wrapping promise")
                connected = wrap_future_with_promise(tornado.websocket.websocket_connect(self._build_url("websocket"), on_message_callback=self._on_websocket_message))
                logger.debug("continue")
                self.websocket_connected.fulfill(connected)
            except:
                logger.exception("error connecting")
                # raise
        logger.debug("add callback")
        self.io_loop.add_callback(do)
        logger.debug("added callback: ")
        # self.io_loop.start()
        # if self.port == 29345:
        # import pdb
        # pdb.set_trace()
        logger.debug("waiting for connection")
        result = self.websocket_connected.get()
        logger.debug("websocket connected")
        if self.websocket_connected.isRejected:
            raise self.websocket.reason

    def _on_websocket_message(self, msg):
        json_data, data = msg.split(b"\n", 1)
        response = json.loads(json_data.decode("utf8"))
        if data:
            import zlib
            data = zlib.decompress(data)
            numpy_array = np.frombuffer(data, dtype=np.dtype(response["dtype"])).reshape(ast.literal_eval(response["shape"]))
            response["result"] = numpy_array
        import sys
        if sys.getsizeof(msg) > 1024 * 4:
            logger.debug("socket read message: <large amount of data>",)
        else:
            logger.debug("socket read message: %s", msg)
            logger.debug("json response: %r", response)
        # for the moment, job == task, in the future a job can be multiple tasks
        job_id = response.get("job_id")
        if job_id:
            try:
                phase = response["job_phase"]
                logger.debug("job update %r, phase=%r", job_id, phase)
                if phase == "COMPLETED":
                    result = response["result"]  # [0]
                    # logger.debug("completed job %r, result=%r", job_id, result)
                    task = self.jobs[job_id]
                    logger.debug("completed job %r (delay=%r, thread_mover=%r)", job_id, task.delay, self.thread_mover)
                    processed_result = task.post_process(result)
                    if task.delay:
                        self.thread_mover(task.fulfill, processed_result)
                    else:
                        task.fulfill(processed_result)
                elif phase == "EXCEPTION":
                    logger.error("exception happened at server side: %r", response)
                    class_name = response["exception"]["class"]
                    msg = response["exception"]["msg"]
                    exception = getattr(__builtin__, class_name)(msg)
                    logger.debug("error in job %r, exception=%r", job_id, exception)
                    task = self.jobs[job_id]
                    if task.delay:
                        self.thread_mover(task.reject, exception)
                    else:
                        task.reject(exception)
                elif phase == "ERROR":
                    logger.error("error happened at server side: %r", response)
                    msg = response["error"]
                    exception = RuntimeError("error at server: %r" % msg)
                    task = self.jobs[job_id]
                    if task.delay:
                        self.thread_mover(task.reject, exception)
                    else:
                        task.reject(exception)
                elif phase == "PENDING":
                    fraction = response["progress"]
                    logger.debug("pending?: %r", phase)
                    task = self.jobs[job_id]
                    if task.delay:
                        self.thread_mover(task.signal_progress.emit, fraction)
                    else:
                        task.signal_progress.emit(fraction)
            except Exception as e:
                logger.exception("error in handling job", e)
                task = self.jobs[job_id]
                if task.delay:
                    self.thread_mover(task.reject, e)
                else:
                    task.reject(e)

    def wait(self):
        io_loop = tornado.ioloop.IOLoop.instance()
        io_loop.start()

    def submit_websocket(self, path, arguments, delay=False, progress=None, post_process=lambda x: x):
        assert self.use_websocket

        task = TaskServer(post_process=post_process, delay=delay)
        progressbars = vaex.utils.progressbars(progress)
        progressbars.add_task(task)
        logger.debug("created task: %r, %r (delay=%r)" % (path, arguments, delay))
        job_id = str(uuid.uuid4())
        self.jobs[job_id] = task
        arguments["job_id"] = job_id
        arguments["path"] = path
        arguments["user_id"] = self.user_id
        # arguments = dict({key: (value.tolist() if hasattr(value, "tolist") else value) for key, value in arguments.items()})
        arguments = dict({key: listify(value) for key, value in arguments.items()})

        def do():
            def write(socket):
                try:
                    logger.debug("write to websocket: %r", arguments)
                    socket.write_message(json.dumps(arguments))
                    return
                except:
                    import traceback
                    traceback.print_exc()
            # return
            logger.debug("will schedule a write to the websocket")
            self.websocket_connected.then(write).end()  # .then(task.fulfill)

        self.io_loop.add_callback(do)
        logger.debug("we can continue (main thread is %r)", threading.currentThread())
        if delay:
            return task
        else:
            return task.get()

    def submit_http(self, path, arguments, post_process, delay, progress=None, **kwargs):
        def pre_post_process(response):
            cookie = Cookie.SimpleCookie()
            for cookieset in response.headers.get_list("Set-Cookie"):
                cookie.load(cookieset)
                logger.debug("cookie load: %r", cookieset)
            logger.debug("cookie: %r", cookie)
            if "user_id" in cookie:
                user_id = cookie["user_id"].value
                logger.debug("user_id: %s", user_id)
                if self.user_id != user_id:
                    self.user_id = user_id
                    vaex.settings.webclient.store("cookie.user_id", self.user_id)
            data = response.body
            is_json = False
            logger.info("response is: %r", response.body)
            logger.info("content_type is: %r", response.headers["Content-Type"])
            if response.headers["Content-Type"] == "application/numpy-array":
                shape, dtype, data = response.body.split(b"\n", 2)
                shape = shape.decode("ascii")
                dtype = dtype.decode("ascii")
                import ast
                numpy_array = np.fromstring(data, dtype=np.dtype(dtype)).reshape(ast.literal_eval(shape))
                return post_process(numpy_array)
            else:
                try:
                    data = json.loads(response.body.decode("ascii"))
                    is_json = True
                except Exception as e:
                    logger.info("couldn't convert to json (error is %s, assume it's raw data): %s", e, data)
                    # logger.info("couldn't convert to json (error is %s, assume it's raw data)", e)
                if is_json:
                    self._check_exception(data)
                    return post_process(data["result"])
                else:
                    return post_process(data)

        arguments = {key: listify(value) for key, value in arguments.items()}
        import pdb
        arguments_json = {key: json.dumps(value) for key, value in arguments.items()}
        headers = tornado.httputil.HTTPHeaders()

        url = self._build_url(path + "?" + urlencode(arguments_json))
        logger.debug("fetch %s, delay=%r", url, delay)

        if self.user_id is not None:
            headers.add("Cookie", "user_id=%s" % self.user_id)
            logger.debug("adding user_id %s to request", self.user_id)
        if delay:
            task = TaskServer(pre_post_process, delay=delay)
            # tornado doesn't like that we call fetch while ioloop is running in another thread, we should use ioloop.add_callbacl

            def do():
                self.thread_mover(task.signal_progress.emit, 0.5)
                future = self.http_client_async.fetch(url, headers=headers, request_timeout=DEFAULT_REQUEST_TIMEOUT, **kwargs)

                def thread_save_succes(value):
                    self.thread_mover(task.signal_progress.emit, 1.0)
                    self.thread_mover(task.fulfill, value)

                def thread_save_failure(value):
                    self.thread_mover(task.reject, value)
                wrap_future_with_promise(future).then(pre_post_process).then(thread_save_succes, thread_save_failure)
            self.io_loop.add_callback(do)
            return task  # promise.then(self._move_to_thread)
        else:
            return pre_post_process(self.http_client.fetch(url, headers=headers, request_timeout=DEFAULT_REQUEST_TIMEOUT, **kwargs))

    def datasets(self, as_dict=False, delay=False):
        def post(result):
            logger.debug("datasets result: %r", result)

            def create(server, state):
                dataset = DatasetRest(self, name=state['name'],
                                      length_original=state['length_original'],
                                      column_names=state['column_names'],
                                      dtypes=state['dtypes'])
                dataset.state_set(state['state'])
                return dataset
            datasets = [create(self, kwargs) for kwargs in result]
            logger.debug("datasets: %r", datasets)
            return datasets if not as_dict else dict([(ds.name, ds) for ds in datasets])
        return self.submit(path="datasets", arguments={}, post_process=post, delay=delay)

    def _build_url(self, method):
        protocol = "ws" if self.use_websocket else "http"
        return "%s://%s:%d%s%s" % (protocol, self.hostname, self.port, self.base_path, method)

    def _call_subspace(self, method_name, subspace, **kwargs):
        def post_process(result):
            if method_name == "histogram":  # histogram is the exception..
                # TODO: don't do binary transfer, just json, now we cannot handle exception
                import base64
                # logger.debug("result: %r", result)
                # result = base64.b64decode(result) #.decode("base64")
                # result = base64.
                data = np.fromstring(result, dtype=np.float64)
                shape = (kwargs["size"],) * subspace.dimension
                data = data.reshape(shape)
                return data
            else:
                try:
                    return np.array(result)
                except ValueError:
                    return result
        dataset_name = subspace.dataset.name
        expressions = subspace.expressions
        delay = subspace.delay
        path = "datasets/%s/%s" % (dataset_name, method_name)
        url = self._build_url(path)
        arguments = dict(kwargs)
        dataset_remote = subspace.dataset
        arguments["selection"] = subspace.is_masked
        arguments['state'] = dataset_remote.state_get()
        arguments['auto_fraction'] = dataset_remote.get_auto_fraction()
        arguments.update(dict(expressions=expressions))
        return self.submit(path, arguments, post_process=post_process, delay=delay)

    def _call_dataset(self, method_name, dataset_remote, delay, numpy=False, progress=None, **kwargs):
        def post_process(result):
            # result = self._check_exception(json.loads(result.body))["result"]
            if numpy:
                return np.fromstring(result, dtype=np.float64)
            else:
                try:
                    return np.array(result)
                except ValueError:
                    return result
        path = "datasets/%s/%s" % (dataset_remote.name, method_name)
        arguments = dict(kwargs)
        arguments['state'] = dataset_remote.state_get()
        arguments['auto_fraction'] = dataset_remote.get_auto_fraction()
        body = urlencode(arguments)

        return self.submit(path, arguments, post_process=post_process, progress=progress, delay=delay)

    def _schedule_call(self, method_name, dataset_remote, delay, **kwargs):
        def post_process(result):
            # result = self._check_exception(json.loads(result.body))["result"]
            try:
                return np.array(result)
            except ValueError:
                return result
        method = "%s/%s" % (dataset_remote.name, method_name)
        return self.schedule(path, arguments, post_process=post_process, delay=delay)

    def _check_exception(self, reply_json):
        if "exception" in reply_json:
            logger.error("exception happened at server side: %r", reply_json)
            class_name = reply_json["exception"]["class"]
            msg = reply_json["exception"]["msg"]
            raise getattr(__builtin__, class_name)(msg)
        if "error" in reply_json:
            raise ValueError("unknown error occured at server")
        else:
            return reply_json

class SubspaceRemote(Subspace):
    def toarray(self, list):
        return np.array(list)

    @property
    def dimension(self):
        return len(self.expressions)

    def _task(self, promise):
        """Helper function for returning tasks results, result when immediate is True, otherwise the task itself, which is a promise"""
        if self.delay:
            return promise
        else:
            return promise

    def sleep(self, seconds, delay=False):
        return self.dataset.server.call("sleep", seconds, delay=delay)

    def minmax(self):
        return self._task(self.dataset.server._call_subspace("minmax", self))
        # return self._task(task)

    def histogram(self, limits, size=256, weight=None):
        return self._task(self.dataset.server._call_subspace("histogram", self, size=size, limits=limits, weight=weight))

    def nearest(self, point, metric=None):
        point = vaex.utils.make_list(point)
        result = self.dataset.server._call_subspace("nearest", self, point=point, metric=metric)
        return self._task(result)

    def mean(self):
        return self.dataset.server._call_subspace("mean", self)

    def correlation(self, means=None, vars=None):
        return self.dataset.server._call_subspace("correlation", self, means=means, vars=vars)

    def var(self, means=None):
        return self.dataset.server._call_subspace("var", self, means=means)

    def sum(self):
        return self.dataset.server._call_subspace("sum", self)

    def limits_sigma(self, sigmas=3, square=False):
        return self.dataset.server._call_subspace("limits_sigma", self, sigmas=sigmas, square=square)

    def mutual_information(self, limits=None, size=256):
        return self.dataset.server._call_subspace("mutual_information", self, limits=limits, size=size)

class DatasetRemote(Dataset):
    def __init__(self, name, server, column_names):
        super(DatasetRemote, self).__init__(name, column_names)
        self.server = server

class DatasetRest(DatasetRemote):
    def __init__(self, server, name, column_names, dtypes, length_original):
        DatasetRemote.__init__(self, name, server.hostname, column_names)
        self.server = server
        self.name = name
        self.column_names = column_names
        self._dtypes = {name: np.dtype(dtype) for name, dtype in dtypes.items()}
        for column_name in self.get_column_names(virtual=True, strings=True):
            self._save_assign_expression(column_name)
        self._length_original = length_original
        self._length_unfiltered = length_original
        self._index_end = length_original
        self.path = self.filename = self.server._build_url("%s" % name)
        self.fraction = 1
        self.executor = ServerExecutor()

    def copy(self, column_names=None, virtual=True):
        dtypes = {name: self.dtype(name) for name in self.get_column_names(strings=True, virtual=False)}
        ds = DatasetRest(self.server, self.name, self.column_names, dtypes=dtypes, length_original=self._length_original)
        state = self.state_get()
        if not virtual:
            state['virtual_columns'] = {}
        ds.state_set(state)
        return ds

    def count(self, expression=None, binby=[], limits=None, shape=default_shape, selection=False, delay=False, edges=False, progress=None):
        return self._delay(delay, self.server._call_dataset("count", self, delay=True, progress=progress, expression=expression, binby=binby, limits=limits, shape=shape, selection=selection, edges=edges))

    def mean(self, expression, binby=[], limits=None, shape=default_shape, selection=False, delay=False, progress=None):
        return self._delay(delay, self.server._call_dataset("mean", self, delay=True, progress=progress, expression=expression, binby=binby, limits=limits, shape=shape, selection=selection))

    def sum(self, expression, binby=[], limits=None, shape=default_shape, selection=False, delay=False, progress=None):
        return self._delay(delay, self.server._call_dataset("sum", self, delay=True, progress=progress, expression=expression, binby=binby, limits=limits, shape=shape, selection=selection))

    def var(self, expression, binby=[], limits=None, shape=default_shape, selection=False, delay=False, progress=None):
        return self._delay(delay, self.server._call_dataset("var", self, delay=True, progress=progress, expression=expression, binby=binby, limits=limits, shape=shape, selection=selection))

    def minmax(self, expression, binby=[], limits=None, shape=default_shape, selection=False, delay=False, progress=None):
        return self._delay(delay, self.server._call_dataset("minmax", self, delay=True, progress=progress, expression=expression, binby=binby, limits=limits, shape=shape, selection=selection))

    # def count(self, expression=None, binby=[], limits=None, shape=default_shape, selection=False, delay=False):
    def cov(self, x, y=None, binby=[], limits=None, shape=default_shape, selection=False, delay=False, progress=None):
        return self._delay(delay, self.server._call_dataset("cov", self, delay=True, progress=progress, x=x, y=y, binby=binby, limits=limits, shape=shape, selection=selection))

    def correlation(self, x, y=None, binby=[], limits=None, shape=default_shape, sort=False, sort_key=np.abs, selection=False, delay=False, progress=None):
        # TODO: sort and sort_key should be done locally
        return self._delay(delay, self.server._call_dataset("correlation", self, delay=True, progress=progress, x=x, y=y, binby=binby, limits=limits, shape=shape, selection=selection))

    def _delay(self, delay, task, progressbar=False):
        if delay:
            return task
        else:
            result = task.get()
            logger.debug("result = %r", result)
            return result

    def dtype(self, expression):
        if expression in self._dtypes:
            return self._dtypes[expression]
        else:
            return np.zeros(1, dtype=np.float64).dtype

    def is_local(self): return False

    def __repr__(self):
        name = self.__class__.__module__ + "." + self.__class__.__name__
        return "<%s(server=%r, name=%r, column_names=%r, __len__=%r)> instance at 0x%x" % (name, self.server, self.name, self.column_names, len(self), id(self))

    def __call__(self, *expressions, **kwargs):
        return SubspaceRemote(self, expressions, kwargs.get("executor") or self.executor, delay=kwargs.get("delay", False))

    def evaluate(self, expression, i1=None, i2=None, out=None, selection=None, delay=False):
        expression = _ensure_strings_from_expressions(expression)
        """basic support for evaluate at server, at least to run some unittest, do not expect this to work from strings"""
        result = self.server._call_dataset("evaluate", self, expression=expression, i1=i1, i2=i2, selection=selection, delay=delay)
        # TODO: we ignore out
        return result

# we may get rid of this when we group together tasks
class ServerExecutor(object):
    def __init__(self):
        self.signal_begin = vaex.events.Signal("begin")
        self.signal_progress = vaex.events.Signal("progress")
        self.signal_end = vaex.events.Signal("end")
        self.signal_cancel = vaex.events.Signal("cancel")

    def execute(self):
        logger.debug("dummy execute")

class TaskServer(Task):
    def __init__(self, post_process, delay):
        Task.__init__(self, None, [])
        self.post_process = post_process
        self.delay = delay
        self.task_queue = []

    def schedule(self, task):
        self.task_queue.append(task)
        logger.info("task added, queue: %r", self.task_queue)
        return task

    def execute(self):
        logger.debug("starting with execute")
        if self._is_executing:
            logger.debug("nested execute call")
            # this situation may happen since in this methods, via a callback (to update a progressbar) we enter
            # Qt's eventloop, which may execute code that will call execute again
            # as long as that code is using delay tasks (i.e. promises) we can simple return here, since after
            # the execute is almost finished, any new tasks added to the task_queue will get executing
            return
        # u 'column' is uniquely identified by a tuple of (dataset, expression)
        self._is_executing = True
        try:
            t0 = time.time()
            task_queue_all = list(self.task_queue)
            if not task_queue_all:
                logger.info("only had cancelled tasks")
            logger.info("clearing queue")
            # self.task_queue = [] # Ok, this was stupid.. in the meantime there may have been new tasks, instead, remove the ones we copied
            for task in task_queue_all:
                logger.info("remove from queue: %r", task)
                self.task_queue.remove(task)
            logger.info("left in queue: %r", self.task_queue)
            task_queue_all = [task for task in task_queue_all if not task.cancelled]
            logger.debug("executing queue: %r" % (task_queue_all))

            # for task in self.task_queue:
            # $	print task, task.expressions_all
            datasets = set(task.dataset for task in task_queue_all)
            cancelled = [False]

            def cancel():
                logger.debug("cancelling")
                self.signal_cancel.emit()
                cancelled[0] = True
            try:
                # process tasks per dataset
                self.signal_begin.emit()
                for dataset in datasets:
                    task_queue = [task for task in task_queue_all if task.dataset == dataset]
                    expressions = list(set(expression for task in task_queue for expression in task.expressions_all))

                    for task in task_queue:
                        task._results = []
                        task.signal_progress.emit(0)
                    self.server.execute_queue(task_queue)
                    self._is_executing = False
            except:
                # on any error we flush the task queue
                self.signal_cancel.emit()
                logger.exception("error in task, flush task queue")
                raise
            logger.debug("executing took %r seconds" % (time.time() - t0))
            # while processing the self.task_queue, new elements will be added to it, so copy it
            logger.debug("cancelled: %r", cancelled)
            if cancelled[0]:
                logger.debug("execution aborted")
                task_queue = task_queue_all
                for task in task_queue:
                    # task._result = task.reduce(task._results)
                    # task.reject(UserAbort("cancelled"))
                    # remove references
                    task._result = None
                    task._results = None
            else:
                task_queue = task_queue_all
                for task in task_queue:
                    logger.debug("fulfill task: %r", task)
                    if not task.cancelled:
                        task._result = task.reduce(task._results)
                        task.fulfill(task._result)
                        # remove references
                    task._result = None
                    task._results = None
                self.signal_end.emit()
                # if new tasks were added as a result of this, execute them immediately
                # TODO: we may want to include infinite recursion protection
                self._is_executing = False
                if len(self.task_queue) > 0:
                    logger.debug("task queue not empty.. start over!")
                    self.execute()
        finally:
            self._is_executing = False

if __name__ == "__main__":
    import vaex
    import sys
    vaex.set_log_level_debug()
    server = vaex.server(sys.argv[1], port=int(sys.argv[2]))
    datasets = server.datasets()
    print(datasets)
    dataset = datasets[0]
    dataset = vaex.example()
    print(dataset("x").minmax())
    dataset.select("x < 0")
    print(dataset.selected_length(), len(dataset))
    print(dataset("x").selected().is_masked)
    print(dataset("x").selected().minmax())


__author__ = 'maartenbreddels'
import aplus

"""
Mini library to make working with promises/futures a bit more Python like.

Example:

@delayed
def f(grid):
    return grid**2

f(ds.count(delay=True))

"""

def promisify(value):
    # TODO, support futures etc..
    if isinstance(value, aplus.Promise):
        return value
    if isinstance(value, (list, tuple)):
        return aplus.listPromise(*list([promisify(k) for k in value]))
    else:
        return aplus.Promise.fulfilled(value)

def delayed(f):
    '''Decorator to transparantly accept delayed computation.

    Example:

    >>> delayed_sum = ds.sum(ds.E, binby=ds.x, limits=limits,
    >>>                   shape=4, delay=True)
    >>> @vaex.delayed
    >>> def total_sum(sums):
    >>>     return sums.sum()
    >>> sum_of_sums = total_sum(delayed_sum)
    >>> ds.execute()
    >>> sum_of_sums.get()
    See the tutorial for a more complete example https://docs.vaex.io/en/latest/tutorial.html#Parallel-computations
    '''

    def wrapped(*args, **kwargs):
        # print "calling", f, "with", kwargs
        # key_values = kwargs.items()
        key_promise = list([(key, promisify(value)) for key, value in kwargs.items()])
        # key_promise = [(key, promisify(value)) for key, value in key_values]
        arg_promises = list([promisify(value) for value in args])
        kwarg_promises = list([promise for key, promise in key_promise])
        promises = arg_promises + kwarg_promises
        for promise in promises:
            def echo_error(exc, promise=promise):
                print("error with ", promise, "exception is", exc)
                # raise exc

            def echo(value, promise=promise):
                print("done with ", repr(promise), "value is", value)
            # promise.then(echo, echo_error)

        # print promises
        allarguments = aplus.listPromise(*promises)

        def call(_):
            kwargs_real = {key: promise.get() for key, promise in key_promise}
            args_real = list([promise.get() for promise in arg_promises])
            return f(*args_real, **kwargs_real)

        def error(exc):
            print("error", exc)
            raise exc
        return allarguments.then(call, error)
    return wrapped

@delayed
def delayed_args(*args):
    return args

@delayed
def delayed_list(l):
    return delayed_args(*l)


import logging
import numpy as np

import vaex.expression
from .utils import _split_and_combine_mask, as_flat_float

logger = logging.getLogger('vaex.selections')

def _select_replace(maskold, masknew):
    return masknew

def _select_and(maskold, masknew):
    return masknew if maskold is None else maskold & masknew

def _select_or(maskold, masknew):
    return masknew if maskold is None else maskold | masknew

def _select_xor(maskold, masknew):
    return masknew if maskold is None else maskold ^ masknew

def _select_subtract(maskold, masknew):
    return ~masknew if maskold is None else (maskold) & ~masknew

_select_functions = {"replace": _select_replace,
                     "and": _select_and,
                     "or": _select_or,
                     "xor": _select_xor,
                     "subtract": _select_subtract
                     }

class Selection(object):
    def __init__(self, previous_selection, mode):
        # we don't care about the previous selection if we simply replace the current selection
        self.previous_selection = previous_selection if mode != "replace" else None
        self.mode = mode

    def execute(self, datexecutor, execute_fully=False):
        if execute_fully and self.previous_selection:
            self.previous_selection.execute(executor=executor, execute_fully=execute_fully)

    def _depending_columns(self, ds):
        '''Find all columns that this selection depends on for df ds'''
        depending = set()
        for expression in self.expressions:
            expression = ds._expr(expression)  # make sure it is an expression
            depending |= expression.variables()
        if self.previous_selection:
            depending |= self.previous_selection._depending_columns(ds)
        return depending

class SelectionDropNa(Selection):
    def __init__(self, drop_nan, drop_masked, column_names, previous_selection, mode):
        super(SelectionDropNa, self).__init__(previous_selection, mode)
        self.drop_nan = drop_nan
        self.drop_masked = drop_masked
        self.column_names = column_names
        self.expressions = self.column_names

    def to_dict(self):
        previous = None
        if self.previous_selection:
            previous = self.previous_selection.to_dict()
        return dict(type="dropna", drop_nan=self.drop_nan, drop_masked=self.drop_masked, column_names=self.column_names,
                    mode=self.mode, previous_selection=previous)

    def _rename(self, df, old, new):
        pass  # TODO: do we need to rename the column_names?

    def evaluate(self, df, name, i1, i2):
        if self.previous_selection:
            previous_mask = df.evaluate_selection_mask(name, i1, i2, selection=self.previous_selection)
        else:
            previous_mask = None
        mask = np.ones(i2 - i1, dtype=np.bool)
        for name in self.column_names:
            data = df._evaluate(name, i1, i2)
            if self.drop_nan and data.dtype.kind in 'O':
                if np.ma.isMaskedArray(data):
                    strings = data.data.astype(str)
                else:
                    strings = data.astype(str)
                mask = mask & (strings != 'nan')
            elif self.drop_nan and data.dtype.kind == "f":
                if np.ma.isMaskedArray(data):
                    mask = mask & ~np.isnan(data.data)
                else:
                    mask = mask & ~np.isnan(data)
            if self.drop_masked and np.ma.isMaskedArray(data):
                mask = mask & ~data.mask  # ~np.ma.getmaskarray(data)
        if previous_mask is None:
            logger.debug("setting mask")
        else:
            logger.debug("combining previous mask with current mask using op %r", self.mode)
            mode_function = _select_functions[self.mode]
            mask = mode_function(previous_mask, mask)
        return mask

def _rename_expression_string(df, e, old, new):
    return vaex.expression.Expression(self.df, self.boolean_expression)._rename(old, new).expression

class SelectionExpression(Selection):
    def __init__(self,  boolean_expression, previous_selection, mode):
        super(SelectionExpression, self).__init__(previous_selection, mode)
        self.boolean_expression = str(boolean_expression)
        self.expressions = [self.boolean_expression]

    def _rename(self, df, old, new):
        boolean_expression = vaex.expression.Expression(df, self.boolean_expression)._rename(old, new).expression
        previous_selection = None
        if self.previous_selection:
            previous_selection = self.previous_selection._rename(df, old, new)
        return SelectionExpression(boolean_expression, previous_selection, self.mode)

    def to_dict(self):
        previous = None
        if self.previous_selection:
            previous = self.previous_selection.to_dict()
        return dict(type="expression", boolean_expression=str(self.boolean_expression), mode=self.mode, previous_selection=previous)

    def evaluate(self, df, name, i1, i2):
        if self.previous_selection:
            previous_mask = df._evaluate_selection_mask(name, i1, i2, selection=self.previous_selection)
        else:
            previous_mask = None
        current_mask = df._evaluate_selection_mask(self.boolean_expression, i1, i2).astype(np.bool)
        if previous_mask is None:
            logger.debug("setting mask")
            mask = current_mask
        else:
            logger.debug("combining previous mask with current mask using op %r", self.mode)
            mode_function = _select_functions[self.mode]
            mask = mode_function(previous_mask, current_mask)
        return mask

class SelectionInvert(Selection):
    def __init__(self, previous_selection):
        super(SelectionInvert, self).__init__(previous_selection, "")
        self.expressions = []

    def to_dict(self):
        previous = None
        if self.previous_selection:
            previous = self.previous_selection.to_dict()
        return dict(type="invert", previous_selection=previous)

    def evaluate(self, df, name, i1, i2):
        previous_mask = df._evaluate_selection_mask(name, i1, i2, selection=self.previous_selection)
        return ~previous_mask

class SelectionLasso(Selection):
    def __init__(self, boolean_expression_x, boolean_expression_y, xseq, yseq, previous_selection, mode):
        super(SelectionLasso, self).__init__(previous_selection, mode)
        self.boolean_expression_x = boolean_expression_x
        self.boolean_expression_y = boolean_expression_y
        self.xseq = xseq
        self.yseq = yseq
        self.expressions = [boolean_expression_x, boolean_expression_y]

    def evaluate(self, df, name, i1, i2):
        if self.previous_selection:
            previous_mask = df._evaluate_selection_mask(name, i1, i2, selection=self.previous_selection)
        else:
            previous_mask = None
        current_mask = np.zeros(i2 - i1, dtype=np.bool)
        x, y = np.array(self.xseq, dtype=np.float64), np.array(self.yseq, dtype=np.float64)
        meanx = x.mean()
        meany = y.mean()
        radius = np.sqrt((meanx - x)**2 + (meany - y)**2).max()
        blockx = df._evaluate(self.boolean_expression_x, i1=i1, i2=i2)
        blocky = df._evaluate(self.boolean_expression_y, i1=i1, i2=i2)
        (blockx, blocky), excluding_mask = _split_and_combine_mask([blockx, blocky])
        blockx = as_flat_float(blockx)
        blocky = as_flat_float(blocky)
        vaex.vaexfast.pnpoly(x, y, blockx, blocky, current_mask, meanx, meany, radius)
        if previous_mask is None:
            logger.debug("setting mask")
            mask = current_mask
        else:
            logger.debug("combining previous mask with current mask using op %r", self.mode)
            mode_function = _select_functions[self.mode]
            mask = mode_function(previous_mask, current_mask)
        if excluding_mask is not None:
            mask = mask & (~excluding_mask)
        return mask

    def to_dict(self):
        previous = None
        if self.previous_selection:
            previous = self.previous_selection.to_dict()
        return dict(type="lasso",
                    boolean_expression_x=str(self.boolean_expression_x),
                    boolean_expression_y=str(self.boolean_expression_y),
                    xseq=vaex.utils.make_list(self.xseq),
                    yseq=vaex.utils.make_list(self.yseq),
                    mode=self.mode,
                    previous_selection=previous)

def selection_from_dict(values):
    kwargs = dict(values)
    del kwargs["type"]
    if values["type"] == "lasso":
        kwargs["previous_selection"] = selection_from_dict(values["previous_selection"]) if values["previous_selection"] else None
        return SelectionLasso(**kwargs)
    elif values["type"] == "expression":
        kwargs["previous_selection"] = selection_from_dict(values["previous_selection"]) if values["previous_selection"] else None
        return SelectionExpression(**kwargs)
    elif values["type"] == "invert":
        kwargs["previous_selection"] = selection_from_dict(values["previous_selection"]) if values["previous_selection"] else None
        return SelectionInvert(**kwargs)
    elif values["type"] == "dropna":
        kwargs["previous_selection"] = selection_from_dict(values["previous_selection"]) if values["previous_selection"] else None
        return SelectionDropNa(**kwargs)
    else:
        raise ValueError("unknown type: %r, in dict: %r" % (values["type"], values))