# Copyright 2015 The Brotli Authors. All rights reserved.
#
# Distributed under MIT license.
# See file LICENSE for detail or copy at https://opensource.org/licenses/MIT

import os
import platform
import re
import unittest

try:
    from setuptools import Extension
    from setuptools import setup
except:
    from distutils.core import Extension
    from distutils.core import setup
from distutils.command.build_ext import build_ext
from distutils import errors
from distutils import dep_util
from distutils import log

CURR_DIR = os.path.abspath(os.path.dirname(os.path.realpath(__file__)))

def get_version():
    """ Return BROTLI_VERSION string as defined in 'common/version.h' file. """
    version_file_path = os.path.join(CURR_DIR, 'c', 'common', 'version.h')
    version = 0
    with open(version_file_path, 'r') as f:
        for line in f:
            m = re.match(r'#define\sBROTLI_VERSION\s+0x([0-9a-fA-F]+)', line)
            if m:
                version = int(m.group(1), 16)
    if version == 0:
        return ''
    # Semantic version is calculated as (MAJOR << 24) | (MINOR << 12) | PATCH.
    major = version >> 24
    minor = (version >> 12) & 0xFFF
    patch = version & 0xFFF
    return '{0}.{1}.{2}'.format(major, minor, patch)

def get_test_suite():
    test_loader = unittest.TestLoader()
    test_suite = test_loader.discover('python', pattern='*_test.py')
    return test_suite

class BuildExt(build_ext):

    def get_source_files(self):
        filenames = build_ext.get_source_files(self)
        for ext in self.extensions:
            filenames.extend(ext.depends)
        return filenames

    def build_extension(self, ext):
        if ext.sources is None or not isinstance(ext.sources, (list, tuple)):
            raise errors.DistutilsSetupError(
                "in 'ext_modules' option (extension '%s'), "
                "'sources' must be present and must be "
                "a list of source filenames" % ext.name)

        ext_path = self.get_ext_fullpath(ext.name)
        depends = ext.sources + ext.depends
        if not (self.force or dep_util.newer_group(depends, ext_path, 'newer')):
            log.debug("skipping '%s' extension (up-to-date)", ext.name)
            return
        else:
            log.info("building '%s' extension", ext.name)

        c_sources = []
        cxx_sources = []
        for source in ext.sources:
            if source.endswith('.c'):
                c_sources.append(source)
            else:
                cxx_sources.append(source)
        extra_args = ext.extra_compile_args or []

        objects = []
        for lang, sources in (('c', c_sources), ('c++', cxx_sources)):
            if lang == 'c++':
                if self.compiler.compiler_type == 'msvc':
                    extra_args.append('/EHsc')

            macros = ext.define_macros[:]
            if platform.system() == 'Darwin':
                macros.append(('OS_MACOSX', '1'))
            elif self.compiler.compiler_type == 'mingw32':
                # On Windows Python 2.7, pyconfig.h defines "hypot" as "_hypot",
                # This clashes with GCC's cmath, and causes compilation errors when
                # building under MinGW: http://bugs.python.org/issue11566
                macros.append(('_hypot', 'hypot'))
            for undef in ext.undef_macros:
                macros.append((undef,))

            objs = self.compiler.compile(
                sources,
                output_dir=self.build_temp,
                macros=macros,
                include_dirs=ext.include_dirs,
                debug=self.debug,
                extra_postargs=extra_args,
                depends=ext.depends)
            objects.extend(objs)

        self._built_objects = objects[:]
        if ext.extra_objects:
            objects.extend(ext.extra_objects)
        extra_args = ext.extra_link_args or []
        # when using GCC on Windows, we statically link libgcc and libstdc++,
        # so that we don't need to package extra DLLs
        if self.compiler.compiler_type == 'mingw32':
            extra_args.extend(['-static-libgcc', '-static-libstdc++'])

        ext_path = self.get_ext_fullpath(ext.name)
        # Detect target language, if not provided
        language = ext.language or self.compiler.detect_language(sources)

        self.compiler.link_shared_object(
            objects,
            ext_path,
            libraries=self.get_libraries(ext),
            library_dirs=ext.library_dirs,
            runtime_library_dirs=ext.runtime_library_dirs,
            extra_postargs=extra_args,
            export_symbols=self.get_export_symbols(ext),
            debug=self.debug,
            build_temp=self.build_temp,
            target_lang=language)

NAME = 'Brotli'

VERSION = get_version()

URL = 'https://github.com/google/brotli'

DESCRIPTION = 'Python bindings for the Brotli compression library'

AUTHOR = 'The Brotli Authors'

LICENSE = 'Apache 2.0'

PLATFORMS = ['Posix', 'MacOS X', 'Windows']

CLASSIFIERS = [
    'Development Status :: 4 - Beta',
    'Environment :: Console',
    'Intended Audience :: Developers',
    'License :: OSI Approved :: Apache Software License',
    'Operating System :: MacOS :: MacOS X',
    'Operating System :: Microsoft :: Windows',
    'Operating System :: POSIX :: Linux',
    'Programming Language :: C',
    'Programming Language :: C++',
    'Programming Language :: Python',
    'Programming Language :: Python :: 2',
    'Programming Language :: Python :: 2.7',
    'Programming Language :: Python :: 3',
    'Programming Language :: Python :: 3.3',
    'Programming Language :: Python :: 3.4',
    'Programming Language :: Python :: 3.5',
    'Programming Language :: Unix Shell',
    'Topic :: Software Development :: Libraries',
    'Topic :: Software Development :: Libraries :: Python Modules',
    'Topic :: System :: Archiving',
    'Topic :: System :: Archiving :: Compression',
    'Topic :: Text Processing :: Fonts',
    'Topic :: Utilities',
]

PACKAGE_DIR = {'': 'python'}

PY_MODULES = ['brotli']

EXT_MODULES = [
    Extension(
        '_brotli',
        sources=[
            'python/_brotli.cc',
            'c/common/dictionary.c',
            'c/common/transform.c',
            'c/dec/bit_reader.c',
            'c/dec/decode.c',
            'c/dec/huffman.c',
            'c/dec/state.c',
            'c/enc/backward_references.c',
            'c/enc/backward_references_hq.c',
            'c/enc/bit_cost.c',
            'c/enc/block_splitter.c',
            'c/enc/brotli_bit_stream.c',
            'c/enc/cluster.c',
            'c/enc/compress_fragment.c',
            'c/enc/compress_fragment_two_pass.c',
            'c/enc/dictionary_hash.c',
            'c/enc/encode.c',
            'c/enc/encoder_dict.c',
            'c/enc/entropy_encode.c',
            'c/enc/histogram.c',
            'c/enc/literal_cost.c',
            'c/enc/memory.c',
            'c/enc/metablock.c',
            'c/enc/static_dict.c',
            'c/enc/utf8_util.c',
        ],
        depends=[
            'c/common/constants.h',
            'c/common/context.h',
            'c/common/dictionary.h',
            'c/common/platform.h',
            'c/common/transform.h',
            'c/common/version.h',
            'c/dec/bit_reader.h',
            'c/dec/huffman.h',
            'c/dec/prefix.h',
            'c/dec/state.h',
            'c/enc/backward_references.h',
            'c/enc/backward_references_hq.h',
            'c/enc/backward_references_inc.h',
            'c/enc/bit_cost.h',
            'c/enc/bit_cost_inc.h',
            'c/enc/block_encoder_inc.h',
            'c/enc/block_splitter.h',
            'c/enc/block_splitter_inc.h',
            'c/enc/brotli_bit_stream.h',
            'c/enc/cluster.h',
            'c/enc/cluster_inc.h',
            'c/enc/command.h',
            'c/enc/compress_fragment.h',
            'c/enc/compress_fragment_two_pass.h',
            'c/enc/dictionary_hash.h',
            'c/enc/encoder_dict.h',
            'c/enc/entropy_encode.h',
            'c/enc/entropy_encode_static.h',
            'c/enc/fast_log.h',
            'c/enc/find_match_length.h',
            'c/enc/hash.h',
            'c/enc/hash_composite_inc.h',
            'c/enc/hash_forgetful_chain_inc.h',
            'c/enc/hash_longest_match64_inc.h',
            'c/enc/hash_longest_match_inc.h',
            'c/enc/hash_longest_match_quickly_inc.h',
            'c/enc/hash_rolling_inc.h',
            'c/enc/hash_to_binary_tree_inc.h',
            'c/enc/histogram.h',
            'c/enc/histogram_inc.h',
            'c/enc/literal_cost.h',
            'c/enc/memory.h',
            'c/enc/metablock.h',
            'c/enc/metablock_inc.h',
            'c/enc/params.h',
            'c/enc/prefix.h',
            'c/enc/quality.h',
            'c/enc/ringbuffer.h',
            'c/enc/static_dict.h',
            'c/enc/static_dict_lut.h',
            'c/enc/utf8_util.h',
            'c/enc/write_bits.h',
        ],
        include_dirs=[
            'c/include',
        ],
        language='c++'),
]

TEST_SUITE = 'setup.get_test_suite'

CMD_CLASS = {
    'build_ext': BuildExt,
}

setup(
    name=NAME,
    description=DESCRIPTION,
    version=VERSION,
    url=URL,
    author=AUTHOR,
    license=LICENSE,
    platforms=PLATFORMS,
    classifiers=CLASSIFIERS,
    package_dir=PACKAGE_DIR,
    py_modules=PY_MODULES,
    ext_modules=EXT_MODULES,
    test_suite=TEST_SUITE,
    cmdclass=CMD_CLASS)


__author__ = 'maartenbreddels'
import PIL.Image
import PIL.ImageDraw
try:
    from StringIO import StringIO
    py3 = False
except ImportError:
    from io import StringIO, BytesIO
    py3 = True
import collections
import numpy as np
import matplotlib.colors
import warnings
from base64 import b64encode

def rgba_2_pil(rgba):
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        im = PIL.Image.frombuffer("RGBA", rgba.shape[:2], rgba, 'raw')  # , "RGBA", 0, -1)
    return im

def pil_2_data(im, format="png"):
    if py3:  # py3 case
        f = BytesIO()
    else:
        f = StringIO()
    im.save(f, format)
    return f.getvalue()

def rgba_to_url(rgba):
    bit8 = rgba.dtype == np.uint8
    if not bit8:
        rgba = (rgba * 255.).astype(np.uint8)
    im = rgba_2_pil(rgba)
    data = pil_2_data(im)
    data = b64encode(data)
    if py3:
        data = data.decode("ascii")
    imgurl = "data:image/png;base64," + data + ""
    return imgurl

pdf_modes = collections.OrderedDict()
pdf_modes["multiply"] = np.multiply
pdf_modes["screen"] = lambda a, b: a + b - a * b
pdf_modes["darken"] = np.minimum
pdf_modes["lighten"] = np.maximum

cairo_modes = collections.OrderedDict()
cairo_modes["saturate"] = (lambda aA, aB: np.minimum(1, aA + aB),
                           lambda aA, xA, aB, xB, aR: (np.minimum(aA, 1 - aB) * xA + aB * xB) / aR)
cairo_modes["add"] = (lambda aA, aB: np.minimum(1, aA + aB),
                      lambda aA, xA, aB, xB, aR: (aA * xA + aB * xB) / aR)

modes = list(pdf_modes.keys()) + list(cairo_modes.keys())

def background(shape, color="white", alpha=1, bit8=True):
    rgba = np.zeros(shape + (4,))
    rgba[:] = np.array(matplotlib.colors.colorConverter.to_rgba(color))
    rgba[..., 3] = alpha
    if bit8:
        return (rgba * 255).astype(np.uint8)
    else:
        return rgba

def fade(image_list, opacity=0.5, blend_mode="multiply"):
    result = image_list[0] * 1.
    for i in range(1, len(image_list)):
        result[result[..., 3] > 0, 3] = 1
        layer = image_list[i] * 1.0
        layer[layer[..., 3] > 0, 3] = opacity
        result = blend([result, layer], blend_mode=blend_mode)
    return result

def blend(image_list, blend_mode="multiply"):
    bit8 = image_list[0].dtype == np.uint8
    image_list = image_list[::-1]
    rgba_dest = image_list[0] * 1
    if bit8:
        rgba_dest = (rgba_dest / 255.).astype(np.float)
    for i in range(1, len(image_list)):
        rgba_source = image_list[i]
        if bit8:
            rgba_source = (rgba_source / 255.).astype(np.float)
        # assert rgba_source.dtype == image_list[0].dtype, "images have different types: first has %r, %d has %r" % (image_list[0].dtype, i, rgba_source.dtype)

        aA = rgba_source[:, :, 3]
        aB = rgba_dest[:, :, 3]

        if blend_mode in pdf_modes:
            aR = aA + aB * (1 - aA)
        else:
            aR = cairo_modes[blend_mode][0](aA, aB)
        mask = aR > 0
        for c in range(3):  # for r, g and b
            xA = rgba_source[..., c]
            xB = rgba_dest[..., c]
            if blend_mode in pdf_modes:
                f = pdf_modes[blend_mode](xB, xA)
                with np.errstate(divide='ignore', invalid='ignore'):  # these are fine, we are ok with nan's in vaex
                    result = ((1. - aB) * aA * xA + (1. - aA) * aB * xB + aA * aB * f) / aR
            else:
                result = cairo_modes[blend_mode][1](aA, xA, aB, xB, aR)
            with np.errstate(divide='ignore', invalid='ignore'):  # these are fine, we are ok with nan's in vaex
                result = (np.minimum(aA, 1 - aB) * xA + aB * xB) / aR
            # print(result)
            rgba_dest[:, :, c][[mask]] = np.clip(result[[mask]], 0, 1)
        rgba_dest[:, :, 3] = np.clip(aR, 0., 1)
    rgba = rgba_dest
    if bit8:
        rgba = (rgba * 255).astype(np.uint8)
    return rgba

def monochrome(I, color, vmin=None, vmax=None):
    """Turns a intensity array to a monochrome 'image' by replacing each intensity by a scaled 'color'

    Values in I between vmin  and vmax get scaled between 0 and 1, and values outside this range are clipped to this.

    Example
    >>> I = np.arange(16.).reshape(4,4)
    >>> color = (0, 0, 1) # red
    >>> rgb = vx.image.monochrome(I, color) # shape is (4,4,3)

    :param I: ndarray of any shape (2d for image)
    :param color: sequence of a (r, g and b) value
    :param vmin: normalization minimum for I, or np.nanmin(I) when None
    :param vmax: normalization maximum for I, or np.nanmax(I) when None
    :return:
    """
    if vmin is None:
        vmin = np.nanmin(I)
    if vmax is None:
        vmax = np.nanmax(I)
    normalized = (I - vmin) / (vmax - vmin)
    return np.clip(normalized[..., np.newaxis], 0, 1) * np.array(color)

def polychrome(I, colors, vmin=None, vmax=None, axis=-1):
    """Similar to monochrome, but now do it for multiple colors

    Example
    >>> I = np.arange(32.).reshape(4,4,2)
    >>> colors = [(0, 0, 1), (0, 1, 0)] # red and green
    >>> rgb = vx.image.polychrome(I, colors) # shape is (4,4,3)

    :param I: ndarray of any shape (3d will result in a 2d image)
    :param colors: sequence of [(r,g,b), ...] values
    :param vmin: normalization minimum for I, or np.nanmin(I) when None
    :param vmax: normalization maximum for I, or np.nanmax(I) when None
    :param axis: axis which to sum over, by default the last
    :return:
    """
    axes_length = len(I.shape)
    allaxes = list(range(axes_length))
    otheraxes = list(allaxes)
    otheraxes.remove((axis + axes_length) % axes_length)
    otheraxes = tuple(otheraxes)

    if vmin is None:
        vmin = np.nanmin(I, axis=otheraxes)
    if vmax is None:
        vmax = np.nanmax(I, axis=otheraxes)
    normalized = (I - vmin) / (vmax - vmin)
    return np.clip(normalized, 0, 1).dot(colors)

# c_r + c_g + c_b


from __future__ import print_function
# -*- coding: utf-8 -*-
import threading
import queue
import sys
import os
import pickle
from io import StringIO
import traceback
import time
import signal
import multiprocessing

pickle_encoding = "hex"
pickle_protocol = 2

def log(*args):
	#print "LOG: " + " ".join([str(k) for k in args])
	pass

def serialize(obj):
	data = pickle.dumps(obj, pickle_protocol)
	rawdata = data.encode(pickle_encoding)
	return rawdata

def deserialize(rawdata):
	data = rawdata.decode(pickle_encoding)
	obj = pickle.loads(data)
	return obj

class InfoThread(threading.Thread):
	def __init__(self, fullsize, executions):
		threading.Thread.__init__(self)
		self.fullsize = fullsize
		self.executions = executions
		self.setDaemon(True)
		#self.

	def n_done(self):
		n = 0
		for execution in self.executions:
			n += len(execution.results)
		return n

	def run(self):
		while 1:
			count = self.n_done() #self.fullsize - self.taskQueue.qsize()
			print("%d out of %d tasks completed (%5.2f%%)" % (count, self.fullsize, float(count)/self.fullsize*100))
			time.sleep(0.1)

class Watchdog(threading.Thread):
	def __init__(self, executions):
		threading.Thread.__init__(self)
		#self.fullsize = fullsize
		self.executions = executions
		#self.setDaemon(True):
		self.setDaemon(True)

	def run(self):
		while 1:
			time.sleep(0.5)
			for execution in self.executions:
				print(execution.isAlive(), execution.error)

import gavi.progressbar
class InfoThreadProgressBar(InfoThread):
	def __init__(self, fullsize, executions):
		InfoThread.__init__(self, fullsize, executions)
		self.bar = gavi.progressbar.ProgressBar(0, fullsize)

	def run(self):
		done = False
		error = False
		while (not done) and (not error):
			count = self.n_done() #self.fullsize - self.taskQueue.qsize()
			self.bar.update(count)
			time.sleep(0.1)
			done = count == self.fullsize
			for execution in self.executions:
				if execution.error:
					error = True

class LocalExecutor(object):
	def __init__(self, thread, function):
		self.thread = thread
		self.function = function

	def stop(self):
		pass

	def init(self):
		pass

	def __call__(self, *args, **kwargs):
		info = "ok"
		exc_info = None
		result = None
		try:
			result = self.function(*args, **kwargs)
		except:
			info = "exception"
			exc_info = traceback.format_exc()
		return info, exc_info, result

class IOExecutor(object):
	def __init__(self, thread, function, init, id_number):
		self.thread = thread
		self.function = function
		self.initf = init
		self.id_number = id_number
		self.createIO()
		self.error = False

	def __call__(self, *args, **kwargs):
		self.log("p: sending execute command")
		#os.close(w) # we are not gonna write
		self.output.write("execute\n")
		self.output.write(serialize((args, kwargs)))
		self.output.write("\n")
		self.output.flush()
		self.log("p: reading response")
		response = self.input.readline().strip()
		#self.log("p: response:", `response`)
		error, exc_info, result = deserialize(response)
		return error, exc_info, result

	def init(self):
		return True

	def stop(self, error=False):
		self.error = error
		self.log("p: sending stop command")
		self.output.write("stop\n")
		self.output.flush()
		self.log("waitpid")
		self.join()
		self.log("waitpid done")
		self.input.close()
		self.output.close()

	def join(self):
		os.waitpid(self.pid, 0)

	def log(self, *args):
		log(self.thread.getName()+":", " ".join([str(k) for k in args]))

class ForkExecutor(IOExecutor):
	#def __init__(self, thread, function):
	#	IOExecutor.__init__(self, thread, function)

	def createIO(self):
		self.log("forking")
		#self.r, self.w = os.pipe()
		#self.output = os.fdopen(self.w, 'w')
		#self.input = os.fdopen(self.r, 'r')
		self.parent_input_p, self.child_output_p = os.pipe()
		self.child_input_p, self.parent_output_p = os.pipe()
		#self.r_c, self.w_c = os.pipe()
		#self.output_p = os.fdopen(self.w_p, 'w')
		#self.input_p = os.fdopen(self.r_p, 'r')
		#self.output_c = os.fdopen(self.w_c, 'w')
		#self.input_c = os.fdopen(self.r_c, 'r')
		self.pid = os.fork()
		if not self.pid: # we are the child
			#if self.id_number != 0:
			#	signal.signal(signal.SIGINT, signal.SIG_IGN)
			self.child_output = os.fdopen(self.child_output_p, 'w')
			self.child_input = os.fdopen(self.child_input_p, 'r')
			self.output = self.child_output
			self.input = self.child_input
			os.close(self.parent_output_p)
			os.close(self.parent_input_p)
			try:
				self.childPart()
			except BaseException:
				print("oops")
			self.input.close()
			self.output.close()
			os._exit(0)
		else:
			self.parent_output = os.fdopen(self.parent_output_p, 'w')
			self.parent_input = os.fdopen(self.parent_input_p, 'r')
			self.output = self.parent_output
			self.input = self.parent_input
			os.close(self.child_output_p)
			os.close(self.child_input_p)
		# else just continue

	def childPart(self):
		done = False
		self.log("c: child")
		if self.initf is not None:
			self.initf(self.id_number)

		while not done:
			self.log("c: waiting for command...")
			command = self.input.readline().strip()
			self.log("command:", command)
			if command == "execute":
				# execute...
				response = self.input.readline().strip()
				self.log("c: args:", repr(response))
				args, kwargs = deserialize(response)
				info = "ok"
				exc_info = None
				result = None
				try:
					#self.log("c: calling")
					result = self.function(*args, **kwargs)
					#self.log("c: result" +result)
				except BaseException:
					info = "exception"
					exc_info = traceback.format_exc()
					#raise "bla"
					#done = True
				except KeyboardInterrupt:
					info = "exception"
					exc_info = traceback.format_exc()
					#raise "bla"
					#done = True
				# encode result
				self.log("c: pickling")
				self.output.write(serialize((info, exc_info, result)))
				self.output.write("\n")
				self.output.flush()
				self.log("c: closing")
			elif command == "stop":
				self.log("c: stopping...")
				done = True
				self.log("c: closed, exiting")
			else:
				done = True
				self.log("c: unknown command", repr(command))
				# and exit
				#os._exit(0) # os._exit is used for children

	def join(self):
		os.waitpid(self.pid, 0)

class Execution(threading.Thread):
	def __init__(self, taskQueue, fork, function, init, id_number, args, kwargs):
		self.taskQueue = taskQueue
		self.fork = fork
		self.function = function
		self.init = init
		self.id_number = id_number
		self.args = args
		self.kwargs = kwargs
		self.results = []
		self.done = False
		self.error = False
		threading.Thread.__init__(self)
		if self.fork:
			self.executor = ForkExecutor(self, self.function, self.init, self.id_number)
		else:
			self.executor = LocalExecutor(self, self.function)

	def log(self, *args):
		log(self.getName()+":", " ".join([str(k) for k in args]))

	def run(self):
		task = None
		if self.executor.init():
			task = None
			self.log("starting")
			try:
				task = self.taskQueue.get(False)
			except queue.Empty:
				self.log("empty at first try")
			while task is not None and self.error is False:
				tasknr, args = task
				args = list(args)
				common_args = list(self.args)
				args.extend(common_args)
				kwargs = dict(self.kwargs)

				info, exc_info, result = self.executor(*args, **kwargs)
				if info == "exception":
					print(exc_info)
					self.error = True
				self.log("r: got result")
				self.results.append((tasknr, result))
				try:
					task = self.taskQueue.get(False)
				except queue.Empty:
					self.log("empty queue")
					task = None
			#if not self.error:
			self.executor.stop(self.error)
			self.log("done")
		else:
			self.log("failure starting slave")
			self.error = True
		self.done = True

def countcores():
	stdin, stdout = os.popen2("cat /proc/cpuinfo | grep processor")
	lines = stdout.readlines()
	return len(lines)

def timed(f):
	def execute(*args, **kwargs):
		t0 = time.time()
		utime0, stime0, child_utime0, child_stime0, walltime0 = os.times()
		result = f(*args, **kwargs)
		dt = time.time() - t0
		utime, stime, child_utime, child_stime, walltime = os.times()
		#name = f.func_name
		print()
		print("user time:            % 9.3f sec." % (utime-utime0))
		print("system time:          % 9.3f sec." % (stime-stime0))
		print("user time(children):  % 9.3f sec." % (child_utime-child_utime0))
		print("system time(children):% 9.3f sec." % (child_stime-child_stime0))
		print()
		dt_total = child_utime-child_utime0 + child_stime-child_stime0+utime-utime0+stime-stime0
		print("total cpu time:       % 9.3f sec. (time it would take on a single cpu/core)" % (dt_total))
		print("elapsed time:         % 9.3f sec. (normal wallclock time it took)" % (walltime-walltime0))
		dt = walltime-walltime0
		if dt == 0:
			eff = 0.
		else:
			eff = dt_total/(dt)
		print("efficiency factor     % 9.3f      (ratio of the two above ~= # cores)" % eff)
		return result
	return execute

def parallelize(cores=None, fork=True, flatten=False, info=False, infoclass=InfoThreadProgressBar, init=None, *args, **kwargs):
	"""Function decorater that executes the function in parallel

	Usage::

		@parallelize(cores=10, info=True)
		def f(x):
			return x**2

		x = numpy.arange(0, 100, 0.1)
		y = f(x) # this gets executed parallel

	:param cores: number of cpus/cores to use (if None, it counts the cores using /proc/cpuinfo)
	:param fork: fork a process (should always be true since of the GIT, but can be used with c modules that release the GIT)
	:param flatten: if False and each return value is a list, final result will be a list of lists, if True, all lists are combined to one big list
	:param info: show progress bar (see infoclass)
	:param infoclass: class to instantiate that shows the progress (default shows progressbar)
	:param init: function to be called in each forked process before executing, can be used to set the seed, takes a integer as parameter (number that identifies the process)
	:param args: extra arguments passed to function
	:param kwargs: extra keyword arguments passed to function

	Example::

		@parallelize(cores=10, info=True, n=2)
		def f(x, n):
			return x**n

		x = numpy.arange(0, 100, 0.1)
		y = f(x) # this gets executed parallel

	"""
	if cores == None:
		cores = multiprocessing.cpu_count()
	def wrapper(f):
		def execute(*multiargs):
			results = []
			len(list(zip(*multiargs)))
			N = len(multiargs[0])
			if info:
				print("running %i jobs on %i cores" % (N, cores))
			taskQueue = queue.Queue(len(multiargs[0]))
			#for timenr in range(times):
			#	taskQueue.put(timenr)
			for tasknr, _args in enumerate(zip(*multiargs)):
				taskQueue.put((tasknr, list(_args)))
			#for timenr in range(times):
			#	result = f(*args, **kwargs)
			#	results.append(result)
			executions = [Execution(taskQueue, fork, f, init, corenr, args, kwargs) for corenr in range(cores)]
			if info:
				infoobj = infoclass(len(multiargs[0]), executions)
				infoobj.start()
			for i, execution in enumerate(executions):
				execution.setName("T-%d" % i)
				execution.start()
			#if 1:
			#	watchdog = Watchdog(executions)
			#	watchdog.start()
			error = False
			for execution in executions:
				log("joining:",execution.getName())
				try:
					execution.join()
				except BaseException:
					error = True
				results.extend(execution.results)
				if execution.error:
					error = True
			if info:
				infoobj.join()
			if error:
				print("error", file=sys.stderr)
				results = None
				raise Exception("error in one or more of the executors")
			else:
				results.sort(cmp=lambda a, b: cmp(a[0], b[0]))
				results = [k[1] for k in results]
				#print "bla", results
				if flatten:
					flatresults = []
					for result in results:
						flatresults.extend(result)
					results = flatresults
			return results
		return execute
	return wrapper

if __name__ == "__main__":
	@timed
	@parallelize(cores=6, fork=True, flatten=True, text="common argument")
	def testprime(from_nr, to_nr, text=""):
		#print text, from_nr, to_nr
		primes = []
		from_nr = max(from_nr, 2)
		for p in range(from_nr, to_nr):
			isprime = True
			time.sleep(1)

			for i in range(2, p):
				if p % i == 0:
					isprime = False
					break
			if isprime:
				primes.append(p)
		return primes

	testnumbers = list(range(0, 10001, 100))
	from_nrs = testnumbers[:-1]
	to_nrs = testnumbers[1:]
	results = testprime(from_nrs, to_nrs)
	#print results
	#print countcores()


from __future__ import division, print_function
__author__ = 'maartenbreddels'
import os
import sys
import collections
import numpy as np
import logging
import vaex
import vaex.utils
import vaex.execution
import vaex.export
import vaex.hdf5.dataset
from vaex.column import ColumnStringArrow, str_type

max_length = int(1e5)

on_rtd = os.environ.get('READTHEDOCS', None) == 'True'
try:
    import h5py
except:
    if not on_rtd:
        raise

logger = logging.getLogger("vaex.hdf5.export")

max_int32 = 2**31-1

def export_hdf5_v1(dataset, path, column_names=None, byteorder="=", shuffle=False, selection=False, progress=None, virtual=True):
    """
    :param DatasetLocal dataset: dataset to export
    :param str path: path for file
    :param lis[str] column_names: list of column names to export or None for all columns
    :param str byteorder: = for native, < for little endian and > for big endian
    :param bool shuffle: export rows in random order
    :param bool selection: export selection or not
    :param progress: progress callback that gets a progress fraction as argument and should return True to continue,
            or a default progress bar when progress=True
    :param: bool virtual: When True, export virtual columns
    :return:
    """

    if selection:
        if selection == True:  # easier to work with the name
            selection = "default"
    # first open file using h5py api
    with h5py.File(path, "w") as h5file_output:

        h5data_output = h5file_output.require_group("data")
        # i1, i2 = dataset.current_slice
        N = len(dataset) if not selection else dataset.selected_length(selection)
        if N == 0:
            raise ValueError("Cannot export empty table")
        logger.debug("virtual=%r", virtual)
        logger.debug("exporting %d rows to file %s" % (N, path))
        # column_names = column_names or (dataset.get_column_names() + (list(dataset.virtual_columns.keys()) if virtual else []))
        column_names = column_names or dataset.get_column_names(virtual=virtual, strings=True)

        logger.debug("exporting columns(hdf5): %r" % column_names)
        for column_name in column_names:
            dtype = dataset.dtype(column_name)
            if column_name in dataset.get_column_names(strings=True):
                column = dataset.columns[column_name]
                shape = (N,) + column.shape[1:]
            else:
                shape = (N,)
            if dtype.type == np.datetime64:
                array = h5file_output.require_dataset("/data/%s" % column_name, shape=shape, dtype=np.int64)
                array.attrs["dtype"] = dtype.name
            else:
                try:
                    array = h5file_output.require_dataset("/data/%s" % column_name, shape=shape, dtype=dtype.newbyteorder(byteorder))
                except:
                    logging.exception("error creating dataset for %r, with type %r " % (column_name, dtype))
            array[0] = array[0]  # make sure the array really exists
        random_index_name = None
        column_order = list(column_names)  # copy
        if shuffle:
            random_index_name = "random_index"
            while random_index_name in dataset.get_column_names():
                random_index_name += "_new"
            shuffle_array = h5file_output.require_dataset("/data/" + random_index_name, shape=(N,), dtype=byteorder + "i8")
            shuffle_array[0] = shuffle_array[0]
            column_order.append(random_index_name)  # last item
        h5data_output.attrs["column_order"] = ",".join(column_order)  # keep track or the ordering of columns

    # after this the file is closed,, and reopen it using out class
    dataset_output = vaex.hdf5.dataset.Hdf5MemoryMapped(path, write=True)

    column_names = vaex.export._export(dataset_input=dataset, dataset_output=dataset_output, path=path, random_index_column=random_index_name,
                                       column_names=column_names, selection=selection, shuffle=shuffle, byteorder=byteorder,
                                       progress=progress)
    import getpass
    import datetime
    user = getpass.getuser()
    date = str(datetime.datetime.now())
    source = dataset.path
    description = "file exported by vaex, by user %s, on date %s, from source %s" % (user, date, source)
    if dataset.description:
        description += "previous description:\n" + dataset.description
    dataset_output.copy_metadata(dataset)
    dataset_output.description = description
    logger.debug("writing meta information")
    dataset_output.write_meta()
    dataset_output.close_files()
    return

def export_hdf5(dataset, path, column_names=None, byteorder="=", shuffle=False, selection=False, progress=None, virtual=True, sort=None, ascending=True):
    """
    :param DatasetLocal dataset: dataset to export
    :param str path: path for file
    :param lis[str] column_names: list of column names to export or None for all columns
    :param str byteorder: = for native, < for little endian and > for big endian
    :param bool shuffle: export rows in random order
    :param bool selection: export selection or not
    :param progress: progress callback that gets a progress fraction as argument and should return True to continue,
            or a default progress bar when progress=True
    :param: bool virtual: When True, export virtual columns
    :return:
    """

    if selection:
        if selection == True:  # easier to work with the name
            selection = "default"
    # first open file using h5py api
    with h5py.File(path, "w") as h5file_output:

        h5table_output = h5file_output.require_group("/table")
        h5table_output.attrs["type"] = "table"
        h5columns_output = h5file_output.require_group("/table/columns")
        # i1, i2 = dataset.current_slice
        N = len(dataset) if not selection else dataset.selected_length(selection)
        if N == 0:
            raise ValueError("Cannot export empty table")
        logger.debug("virtual=%r", virtual)
        logger.debug("exporting %d rows to file %s" % (N, path))
        # column_names = column_names or (dataset.get_column_names() + (list(dataset.virtual_columns.keys()) if virtual else []))
        column_names = column_names or dataset.get_column_names(virtual=virtual, strings=True)

        logger.debug("exporting columns(hdf5): %r" % column_names)
        sparse_groups = collections.defaultdict(list)
        sparse_matrices = {}  # alternative to a set of matrices, since they are not hashable
        for column_name in list(column_names):
            sparse_matrix = dataset._sparse_matrix(column_name)
            if sparse_matrix is not None:
                # sparse columns are stored differently
                sparse_groups[id(sparse_matrix)].append(column_name)
                sparse_matrices[id(sparse_matrix)] = sparse_matrix
                continue
            dtype = dataset.dtype(column_name)
            if column_name in dataset.get_column_names(virtual=False):
                column = dataset.columns[column_name]
                shape = (N,) + column.shape[1:]
            else:
                shape = (N,)
            h5column_output = h5columns_output.require_group(column_name)
            if dtype == str_type:
                # TODO: if no selection or filter, we could do this
                # if isinstance(column, ColumnStringArrow):
                #     data_shape = column.bytes.shape
                #     indices_shape = column.indices.shape
                # else:

                byte_length = dataset[column_name].str.byte_length().sum(selection=selection)
                if byte_length > max_int32:
                    dtype_indices = 'i8'
                else:
                    dtype_indices = 'i4'

                data_shape = (byte_length, )
                indices_shape = (N+1, )

                array = h5column_output.require_dataset('data', shape=data_shape, dtype='S1')
                array[0] = array[0]  # make sure the array really exists

                index_array = h5column_output.require_dataset('indices', shape=indices_shape, dtype=dtype_indices)
                index_array[0] = index_array[0]  # make sure the array really exists

                null_value_count = N - dataset.count(column_name, selection=selection)
                if null_value_count > 0:
                    null_shape = ((N + 7) // 8, )  # TODO: arrow requires padding right?
                    null_bitmap_array = h5column_output.require_dataset('null_bitmap', shape=null_shape, dtype='u1')
                    null_bitmap_array[0] = null_bitmap_array[0]  # make sure the array really exists

                array.attrs["dtype"] = 'str'
                # TODO: masked support ala arrow?
            else:
                if dtype.kind in 'mM':
                    array = h5column_output.require_dataset('data', shape=shape, dtype=np.int64)
                    array.attrs["dtype"] = dtype.name
                elif dtype.kind == 'U':
                    # numpy uses utf32 for unicode
                    char_length = dtype.itemsize // 4
                    shape = (N, char_length)
                    array = h5column_output.require_dataset('data', shape=shape, dtype=np.uint8)
                    array.attrs["dtype"] = 'utf32'
                    array.attrs["dlength"] = char_length
                else:
                    try:
                        array = h5column_output.require_dataset('data', shape=shape, dtype=dtype.newbyteorder(byteorder))
                    except:
                        logging.exception("error creating dataset for %r, with type %r " % (column_name, dtype))
                        del h5columns_output[column_name]
                        column_names.remove(column_name)
                array[0] = array[0]  # make sure the array really exists

                data = dataset.evaluate(column_name, 0, 1)
                if np.ma.isMaskedArray(data):
                    mask = h5column_output.require_dataset('mask', shape=shape, dtype=np.bool)
                    mask[0] = mask[0]  # make sure the array really exists
        random_index_name = None
        column_order = list(column_names)  # copy
        if shuffle:
            random_index_name = "random_index"
            while random_index_name in dataset.get_column_names():
                random_index_name += "_new"
            shuffle_array = h5columns_output.require_dataset(random_index_name + "/data", shape=(N,), dtype=byteorder + "i8")
            shuffle_array[0] = shuffle_array[0]
            column_order.append(random_index_name)  # last item
        h5columns_output.attrs["column_order"] = ",".join(column_order)  # keep track or the ordering of columns

        sparse_index = 0
        for sparse_matrix in sparse_matrices.values():
            columns = sorted(sparse_groups[id(sparse_matrix)], key=lambda col: dataset.columns[col].column_index)
            name = "sparse" + str(sparse_index)
            sparse_index += 1
            # TODO: slice columns
            # sparse_matrix = sparse_matrix[:,]
            sparse_group = h5columns_output.require_group(name)
            sparse_group.attrs['type'] = 'csr_matrix'
            ar = sparse_group.require_dataset('data', shape=(len(sparse_matrix.data), ), dtype=sparse_matrix.dtype)
            ar[0] = ar[0]
            ar = sparse_group.require_dataset('indptr', shape=(len(sparse_matrix.indptr), ), dtype=sparse_matrix.indptr.dtype)
            ar[0] = ar[0]
            ar = sparse_group.require_dataset('indices', shape=(len(sparse_matrix.indices), ), dtype=sparse_matrix.indices.dtype)
            ar[0] = ar[0]
            for i, column_name in enumerate(columns):
                h5column = sparse_group.require_group(column_name)
                h5column.attrs['column_index'] = i

    # after this the file is closed,, and reopen it using out class
    dataset_output = vaex.hdf5.dataset.Hdf5MemoryMapped(path, write=True)

    column_names = vaex.export._export(dataset_input=dataset, dataset_output=dataset_output, path=path, random_index_column=random_index_name,
                                       column_names=column_names, selection=selection, shuffle=shuffle, byteorder=byteorder,
                                       progress=progress, sort=sort, ascending=ascending)
    import getpass
    import datetime
    user = getpass.getuser()
    date = str(datetime.datetime.now())
    source = dataset.path
    description = "file exported by vaex, by user %s, on date %s, from source %s" % (user, date, source)
    if dataset.description:
        description += "previous description:\n" + dataset.description
    dataset_output.copy_metadata(dataset)
    dataset_output.description = description
    logger.debug("writing meta information")
    dataset_output.write_meta()
    dataset_output.close_files()
    return


"""Convert between arrow and vaex/numpy columns/arrays without doing memory copies."""
import pyarrow
import numpy as np
from vaex.column import ColumnStringArrow

def arrow_array_from_numpy_array(array):
    dtype = array.dtype
    mask = None
    if np.ma.isMaskedArray(array):
        mask = array.mask
    if dtype.kind == 'S':
        type = pyarrow.binary(dtype.itemsize)
        arrow_array = pyarrow.array(array, type, mask=mask)
    else:
        if dtype.isnative:
            arrow_array = pyarrow.array(array, mask=mask)
        else:
            # TODO: we copy here, but I guess we should not... or give some warning
            arrow_array = pyarrow.array(array.astype(dtype.newbyteorder('=')), mask=mask)
    return arrow_array

from vaex.dataframe import Column

def column_from_arrow_array(arrow_array):
    arrow_type = arrow_array.type
    buffers = arrow_array.buffers()
    if len(buffers) == 2:
        return numpy_array_from_arrow_array(arrow_array)
    elif len(buffers) == 3 and  isinstance(arrow_array.type, type(pyarrow.string())):
        bitmap_buffer, offsets, string_bytes = arrow_array.buffers()
        if arrow_array.null_count == 0:
            null_bitmap = None  # we drop any null_bitmap when there are no null counts
        else:
            null_bitmap = np.frombuffer(bitmap_buffer, 'uint8', len(bitmap_buffer))
        offsets = np.frombuffer(offsets, np.int32, len(offsets)//4)
        string_bytes = np.frombuffer(string_bytes, 'S1', len(string_bytes))
        column = ColumnStringArrow(offsets, string_bytes, len(arrow_array), null_bitmap=null_bitmap)
        return column
    else:
        raise TypeError('type unsupported: %r' % arrow_type)

def numpy_array_from_arrow_array(arrow_array):
    arrow_type = arrow_array.type
    buffers = arrow_array.buffers()
    assert len(buffers) == 2
    bitmap_buffer, data_buffer = buffers
    if isinstance(arrow_type, type(pyarrow.binary(1))):  # todo, is there a better way to typecheck?
        # mimics python/pyarrow/array.pxi::Array::to_numpy
        assert len(buffers) == 2
        dtype = "S" + str(arrow_type.byte_width)
        # arrow seems to do padding, check if it is all ok
        expected_length = arrow_type.byte_width * len(arrow_array)
        actual_length = len(buffers[-1])
        if actual_length < expected_length:
            raise ValueError('buffer is smaller (%d) than expected (%d)' % (actual_length, expected_length))
        array = np.frombuffer(buffers[-1], dtype, len(arrow_array))# TODO: deal with offset ? [arrow_array.offset:arrow_array.offset + len(arrow_array)]
    else:
        dtype = arrow_array.type.to_pandas_dtype()
    if np.bool_ == dtype:
        # TODO: this will also be a copy, we probably want to support bitmasks as well
        bitmap = np.frombuffer(data_buffer, np.uint8, len(data_buffer))
        array = numpy_mask_from_arrow_mask(bitmap, len(arrow_array))
    else:
        array = np.frombuffer(data_buffer, dtype, len(arrow_array))

    if bitmap_buffer is not None:
        bitmap = np.frombuffer(bitmap_buffer, np.uint8, len(bitmap_buffer))
        mask = numpy_mask_from_arrow_mask(bitmap, len(arrow_array))
        array = np.ma.MaskedArray(array, mask=mask)
    return array

def numpy_mask_from_arrow_mask(bitmap, length):
    # arrow uses a bitmap https://github.com/apache/arrow/blob/master/format/Layout.md
    # we do have to change the ordering of the bits
    return 1-np.unpackbits(bitmap).reshape((len(bitmap),8))[:,::-1].reshape(-1)[:length]

def arrow_table_from_vaex_df(ds, column_names=None, selection=None, strings=True, virtual=False):
    """Implementation of Dataset.to_arrow_table"""
    names = []
    arrays = []
    for name, array in ds.to_items(column_names=column_names, selection=selection, strings=strings, virtual=virtual):
        names.append(name)
        arrays.append(arrow_array_from_numpy_array(array))
    return pyarrow.Table.from_arrays(arrays, names)

def vaex_df_from_arrow_table(table):
    from .dataset import DatasetArrow
    return DatasetArrow(table=table)


from vaex.dataset import Dataset
import numpy as np
import math

def patch(f):
    '''Adds method f to the Dataset class'''
    name = f.__name__
    Dataset.__hidden__[name] = f
    return f

# All assume equinox J2000
comat = {'eq2ecl': [[0.9999999999999928, 1.1102233723050031e-07, 4.411803426976324e-08],
                     [-1.1941015020086788e-07, 0.9174821814419274, 0.39777688059582816],
                     [3.684608657254395e-09, -0.39777688059583055, 0.9174821814419342]],
         'eq2gal': [[-0.05487553939574265, -0.8734371047275962, -0.48383499177002515],
                     [0.49410945362774394, -0.4448295942975751, 0.7469822486998918],
                     [-0.8676661356833737, -0.19807638961301982, 0.45598379452141985]],
         'gal2eq': [[-0.0548756577126198, 0.4941094371971076, -0.8676661375571625],
                     [-0.873437051955779, -0.44482972122205366, -0.19807633727507046],
                     [-0.48383507361641837, 0.7469821839845096, 0.45598381369115243]]}

@patch
def add_virtual_columns_eq2ecl(self, long_in="ra", lat_in="dec", long_out="lambda_", lat_out="beta", name_prefix="__celestial_eq2ecl", radians=False):
    """Add ecliptic coordates (long_out, lat_out) from equatorial coordinates.

    :param long_in: Name/expression for right ascension
    :param lat_in: Name/expression for declination
    :param long_out:  Output name for lambda coordinate
    :param lat_out: Output name for beta coordinate
    :param name_prefix:
    :param radians: input and output in radians (True), or degrees (False)
    :return:
    """

    self.add_virtual_columns_celestial(long_in, lat_in, long_out, lat_out, name_prefix=name_prefix, radians=radians, _matrix='eq2ecl')

@patch
def add_virtual_columns_eq2gal(self, long_in="ra", lat_in="dec", long_out="l", lat_out="b", name_prefix="__celestial_eq2gal", radians=False):
    """Add galactic coordates (long_out, lat_out) from equatorial coordinates.

    :param long_in: Name/expression for right ascension
    :param lat_in: Name/expression for declination
    :param long_out:  Output name for galactic longitude
    :param lat_out: Output name for galactic latitude
    :param name_prefix:
    :param radians: input and output in radians (True), or degrees (False)
    :return:
    """

    self.add_virtual_columns_celestial(long_in, lat_in, long_out, lat_out, name_prefix=name_prefix, radians=radians, _matrix='eq2gal')

@patch
def add_virtual_columns_gal2eq(self, long_in='l', lat_in='b', long_out='ra', lat_out='dec', name_prefix="__celestial_gal2eq", radians=False):
    """
    Convert from galactic (l,b) to equatorial (ra,dec) spherical coordinate system.
    :param long_in: longitudinal angle l
    :param lat_in: latitudinal angle b
    :param long_out: right ascension
    :param lat_out: declination
    :param name_prefix:
    :param radians: input and output in radians (True), or degrees (False)
    """

    self.add_virtual_columns_celestial(long_in, lat_in, long_out, lat_out, name_prefix=name_prefix, radians=radians, _matrix='gal2eq')

@patch
def add_virtual_columns_distance_from_parallax(self, parallax="parallax", distance_name="distance", parallax_uncertainty=None, uncertainty_postfix="_uncertainty"):
    """Convert parallax to distance (i.e. 1/parallax)

    :param parallax: expression for the parallax, e.g. "parallax"
    :param distance_name: name for the virtual column of the distance, e.g. "distance"
    :param parallax_uncertainty: expression for the uncertainty on the parallax, e.g. "parallax_error"
    :param uncertainty_postfix: distance_name + uncertainty_postfix is the name for the virtual column, e.g. "distance_uncertainty" by default
    :return:
    """
    """

    """
    import astropy.units
    unit = self.unit(parallax)
    # if unit:
    #   convert = unit.to(astropy.units.mas)
    #   distance_expression = "%f/(%s)" % (convert, parallax)
    # else:
    distance_expression = "1/%s" % (parallax)
    self.ucds[distance_name] = "pos.distance"
    self.descriptions[distance_name] = "Derived from parallax (%s)" % parallax
    if unit:
        if unit == astropy.units.milliarcsecond:
            self.units[distance_name] = astropy.units.kpc
        if unit == astropy.units.arcsecond:
            self.units[distance_name] = astropy.units.parsec
    self.add_virtual_column(distance_name, distance_expression)
    if parallax_uncertainty:
        """
        y = 1/x
        sigma_y**2 = (1/x**2)**2 sigma_x**2
        sigma_y = (1/x**2) sigma_x
        sigma_y = y**2 sigma_x
        sigma_y/y = (1/x) sigma_x
        """
        name = distance_name + uncertainty_postfix
        distance_uncertainty_expression = "{parallax_uncertainty}/({parallax})**2".format(**locals())
        self.add_virtual_column(name, distance_uncertainty_expression)
        self.descriptions[name] = "Uncertainty on parallax (%s)" % parallax
        self.ucds[name] = "stat.error;pos.distance"

@patch
def add_virtual_columns_cartesian_velocities_to_pmvr(self, x="x", y="y", z="z", vx="vx", vy="vy", vz="vz", vr="vr", pm_long="pm_long", pm_lat="pm_lat", distance=None):
    """Concert velocities from a cartesian system to proper motions and radial velocities

    TODO: errors

    :param x: name of x column (input)
    :param y:         y
    :param z:         z
    :param vx:       vx
    :param vy:       vy
    :param vz:       vz
    :param vr: name of the column for the radial velocity in the r direction (output)
    :param pm_long: name of the column for the proper motion component in the longitude direction  (output)
    :param pm_lat: name of the column for the proper motion component in the latitude direction, positive points to the north pole (output)
    :param distance: Expression for distance, if not given defaults to sqrt(x**2+y**2+z**2), but if this column already exists, passing this expression may lead to a better performance
    :return:
    """
    if distance is None:
        distance = "sqrt({x}**2+{y}**2+{z}**2)".format(**locals())
    k = 4.74057
    self.add_variable("k", k, overwrite=False)
    self.add_virtual_column(vr, "({x}*{vx}+{y}*{vy}+{z}*{vz})/{distance}".format(**locals()))
    self.add_virtual_column(pm_long, "-({vx}*{y}-{x}*{vy})/sqrt({x}**2+{y}**2)/{distance}/k".format(**locals()))
    self.add_virtual_column(pm_lat, "-({z}*({x}*{vx}+{y}*{vy}) - ({x}**2+{y}**2)*{vz})/( ({x}**2+{y}**2+{z}**2) * sqrt({x}**2+{y}**2) )/k".format(**locals()))

@patch
def add_virtual_columns_proper_motion_eq2gal(self, long_in="ra", lat_in="dec", pm_long="pm_ra", pm_lat="pm_dec", pm_long_out="pm_l", pm_lat_out="pm_b",
                                            name_prefix="__proper_motion_eq2gal",
                                            right_ascension_galactic_pole=192.85,
                                            declination_galactic_pole=27.12,
                                            propagate_uncertainties=False,
                                            radians=False, inverse=False):
    """Transform/rotate proper motions from equatorial to galactic coordinates

    Taken from http://arxiv.org/abs/1306.2945

    :param long_in: Name/expression for right ascension
    :param lat_in: Name/expression for declination
    :param pm_long: Proper motion for ra
    :param pm_lat: Proper motion for dec
    :param pm_long_out:  Output name for output proper motion on l direction
    :param pm_lat_out: Output name for output proper motion on b direction
    :param name_prefix:
    :param radians: input and output in radians (True), or degrees (False)
    :parap inverse: (For internal use) convert from galactic to equatorial instead
    :return:
    """
    """mu_gb =  mu_dec*(cdec*sdp-sdec*cdp*COS(ras))/cgb $
      - mu_ra*cdp*SIN(ras)/cgb"""
    long_in_original = long_in = self._expr(long_in)
    lat_in_original = lat_in = self._expr(lat_in)
    pm_long = self._expr(pm_long)
    pm_lat = self._expr(pm_lat)
    if not radians:
        long_in = long_in * np.pi/180
        lat_in = lat_in * np.pi/180
    c1_name = name_prefix + "_C1"
    c2_name = name_prefix + "_C2"
    right_ascension_galactic_pole = math.radians(right_ascension_galactic_pole)
    declination_galactic_pole = math.radians(declination_galactic_pole)
    self[c1_name] = c1 = np.sin(declination_galactic_pole) * np.cos(lat_in) - np.cos(declination_galactic_pole)*np.sin(lat_in)*np.cos(long_in-right_ascension_galactic_pole)
    self[c2_name] = c2 = np.cos(declination_galactic_pole) * np.sin(long_in - right_ascension_galactic_pole)
    c1 = self[c1_name]
    c2 = self[c2_name]
    if inverse:
        self[pm_long_out] = ( c1 * pm_long + -c2 * pm_lat)/np.sqrt(c1**2+c2**2)
        self[pm_lat_out] =  ( c2 * pm_long +  c1 * pm_lat)/np.sqrt(c1**2+c2**2)
    else:
        self[pm_long_out] = ( c1 * pm_long + c2 * pm_lat)/np.sqrt(c1**2+c2**2)
        self[pm_lat_out] =  (-c2 * pm_long + c1 * pm_lat)/np.sqrt(c1**2+c2**2)
    if propagate_uncertainties:
        self.propagate_uncertainties([self[pm_long_out], self[pm_lat_out]])

@patch
def add_virtual_columns_proper_motion_gal2eq(self, long_in="ra", lat_in="dec", pm_long="pm_l", pm_lat="pm_b", pm_long_out="pm_ra", pm_lat_out="pm_dec",
                                            name_prefix="__proper_motion_gal2eq",
                                            right_ascension_galactic_pole=192.85,
                                            declination_galactic_pole=27.12,
                                            propagate_uncertainties=False,
                                            radians=False):
    """Transform/rotate proper motions from galactic to equatorial coordinates.

    Inverse of :py:`add_virtual_columns_proper_motion_eq2gal`
    """
    kwargs = dict(**locals())
    kwargs.pop('self')
    kwargs['inverse'] = True
    self.add_virtual_columns_proper_motion_eq2gal(**kwargs)

@patch
def add_virtual_columns_lbrvr_proper_motion2vcartesian(self, long_in="l", lat_in="b", distance="distance", pm_long="pm_l", pm_lat="pm_b",
                                                       vr="vr", vx="vx", vy="vy", vz="vz",
                                                       center_v=(0, 0, 0),
                                                       propagate_uncertainties=False, radians=False):
    """Convert radial velocity and galactic proper motions (and positions) to cartesian velocities wrt the center_v

    Based on http://adsabs.harvard.edu/abs/1987AJ.....93..864J

    :param long_in: Name/expression for galactic longitude
    :param lat_in: Name/expression for galactic latitude
    :param distance: Name/expression for heliocentric distance
    :param pm_long: Name/expression for the galactic proper motion in latitude direction (pm_l*, so cosine(b) term should be included)
    :param pm_lat: Name/expression for the galactic proper motion in longitude direction
    :param vr: Name/expression for the radial velocity
    :param vx: Output name for the cartesian velocity x-component
    :param vy: Output name for the cartesian velocity y-component
    :param vz: Output name for the cartesian velocity z-component
    :param center_v: Extra motion that should be added, for instance lsr + motion of the sun wrt the galactic restframe
    :param radians: input and output in radians (True), or degrees (False)
    :return:
    """
    k = 4.74057
    a, d, distance = self._expr(long_in, lat_in, distance)
    pm_long, pm_lat, vr = self._expr(pm_long, pm_lat, vr)
    if not radians:
        a = a * np.pi/180
        d = d * np.pi/180
    A = [[np.cos(a)*np.cos(d), -np.sin(a), -np.cos(a)*np.sin(d)],
         [np.sin(a)*np.cos(d),  np.cos(a), -np.sin(a)*np.sin(d)],
         [np.sin(d), d*0, np.cos(d)]]
    self.add_virtual_columns_matrix3d(vr, k * pm_long * distance, k * pm_lat * distance, vx, vy, vz, A, translation=center_v)
    if propagate_uncertainties:
        self.propagate_uncertainties([self[vx], self[vy], self[vz]])

@patch
def add_virtual_columns_equatorial_to_galactic_cartesian(self, alpha, delta, distance, xname, yname, zname, radians=True, alpha_gp=np.radians(192.85948), delta_gp=np.radians(27.12825), l_omega=np.radians(32.93192)):
    """From http://arxiv.org/pdf/1306.2945v2.pdf"""
    if not radians:
        alpha = "pi/180.*%s" % alpha
        delta = "pi/180.*%s" % delta
    self.virtual_columns[zname] = "{distance} * (cos({delta}) * cos({delta_gp}) * cos({alpha} - {alpha_gp}) + sin({delta}) * sin({delta_gp}))".format(**locals())
    self.virtual_columns[xname] = "{distance} * (cos({delta}) * sin({alpha} - {alpha_gp}))".format(**locals())
    self.virtual_columns[yname] = "{distance} * (sin({delta}) * cos({delta_gp}) - cos({delta}) * sin({delta_gp}) * cos({alpha} - {alpha_gp}))".format(**locals())
    # self.write_virtual_meta()

@patch
def add_virtual_columns_celestial(self, long_in, lat_in, long_out, lat_out, name_prefix="__celestial", radians=False, _matrix=None):
    #TODO: return wrapped angle, currently returning angles between [-2pi:2pi]. Need modulo arhitmetics on vaex add_virtual_column expressions.
    #      for example see http://docs.astropy.org/en/stable/_modules/astropy/coordinates/angles.html#Angle.wrap_at

    matrix = comat[_matrix]
    if not radians:
        long_in = "pi/180.*%s" % long_in
        lat_in = "pi/180.*%s" % lat_in
    x_in = name_prefix + "_in_x"
    y_in = name_prefix + "_in_y"
    z_in = name_prefix + "_in_z"
    x_out = name_prefix + "_out_x"
    y_out = name_prefix + "_out_y"
    z_out = name_prefix + "_out_z"
    self.add_virtual_column(x_in, "cos({long_in})*cos({lat_in})".format(**locals()))
    self.add_virtual_column(y_in, "sin({long_in})*cos({lat_in})".format(**locals()))
    self.add_virtual_column(z_in, "sin({lat_in})".format(**locals()))
    self.add_virtual_columns_matrix3d(x_in, y_in, z_in, x_out, y_out, z_out,
                                      matrix, name_prefix + "_matrix")
    transform = "" if radians else "*180./pi"
    x = x_out
    y = y_out
    z = z_out
    self.add_virtual_column(long_out, "arctan2({y}, {x}){transform}".format(**locals()))
    self.add_virtual_column(lat_out, "(-arccos({z}/sqrt({x}**2+{y}**2+{z}**2))+pi/2){transform}".format(**locals()))

@patch
def add_virtual_columns_proper_motion2vperpendicular(self, distance="distance", pm_long="pm_l", pm_lat="pm_b",
                                                     vl="vl", vb="vb",
                                                     propagate_uncertainties=False,
                                                     radians=False):
    """Convert proper motion to perpendicular velocities.

    :param distance:
    :param pm_long:
    :param pm_lat:
    :param vl:
    :param vb:
    :param cov_matrix_distance_pm_long_pm_lat:
    :param uncertainty_postfix:
    :param covariance_postfix:
    :param radians:
    :return:
    """
    k = 4.74057
    self.add_variable("k", k, overwrite=False)
    self.add_virtual_column(vl, "k*{pm_long}*{distance}".format(**locals()))
    self.add_virtual_column(vb, "k* {pm_lat}*{distance}".format(**locals()))
    if propagate_uncertainties:
        self.propagate_uncertainties([self[vl], self[vb]])

@patch
def add_virtual_columns_cartesian_angular_momenta(self, x='x', y='y', z='z',
                                                  vx='vx', vy='vy', vz='vz',
                                                  Lx='Lx', Ly='Ly', Lz='Lz',
                                                  propagate_uncertainties=False):
    """
    Calculate the angular momentum components provided Cartesian positions and velocities.
    Be mindful of the point of origin: ex. if considering Galactic dynamics, and positions and
    velocities should be as seen from the Galactic centre.

    :param x: x-position Cartesian component
    :param y: y-position Cartesian component
    :param z: z-position Cartesian component
    :param vx: x-velocity Cartesian component
    :param vy: y-velocity Cartesian component
    :param vz: z-velocity Cartesian component
    :param Lx: name of virtual column
    :param Ly: name of virtual column
    :param Lz: name of virtial column
    :propagate_uncertainties: (bool) whether to propagate the uncertainties of
    the positions and velocities to the angular momentum components
    """

    x, y, z, vx, vy, vz = self._expr(x, y, z, vx, vy, vz)
    self.add_virtual_column(Lx, y * vz - z * vy)
    self.add_virtual_column(Ly, z * vx - x * vz)
    self.add_virtual_column(Lz, x * vy - y * vx)
    if propagate_uncertainties:
        self.propagate_uncertainties([self[Lx], self[Ly], self[Lz]])

def add_plugin():
    pass  # importing this module already does the job


__author__ = 'maartenbreddels'

import functools

import matplotlib.patches as patches
import numpy as np
import matplotlib.artist as artist

import vaex.ui.plugin
from vaex.ui.qt import *
import logging

logger = logging.getLogger("plugin.dispersions")

import matplotlib.transforms as transforms
from matplotlib.path import Path

class DispersionEllipse(patches.Patch):
	"""
	This ellipse has it's center in user coordinates, and the width and height in device coordinates
	such that is is not deformed
	"""
	def __str__(self):
		return "DispersionEllipse(%s,%s;%sx%s)" % (self.center[0], self.center[1],
										 self.width, self.height)

	#@docstring.dedent_interpd
	def __init__(self, xy, width, height, scale=1.0, angle=0.0, **kwargs):
		"""
		*xy*
		  center of ellipse

		*width*
		  total length (diameter) of horizontal axis

		*height*
		  total length (diameter) of vertical axis

		*angle*
		  rotation in degrees (anti-clockwise)

		Valid kwargs are:
		%(Patch)s
		"""
		patches.Patch.__init__(self, **kwargs)

		self.center = xy
		self.width, self.height = width, height
		self.scale = scale
		self.angle = angle
		self._path = Path.unit_circle()
		# Note: This cannot be calculated until this is added to an Axes
		self._patch_transform = transforms.IdentityTransform()

	def _recompute_transform(self):
		"""NOTE: This cannot be called until after this has been added
				 to an Axes, otherwise unit conversion will fail. This
				 maxes it very important to call the accessor method and
				 not directly access the transformation member variable.
		"""
		center = (self.convert_xunits(self.center[0]),
				  self.convert_yunits(self.center[1]))
		width = self.width #self.convert_xunits(self.width)
		height = self.height #self.convert_yunits(self.height)
		trans = artist.Artist.get_transform(self)
		self._patch_transform = transforms.Affine2D() \
			.scale(width * 0.5 * self.scale, height * 0.5* self.scale) \
			.rotate_deg(self.angle) \
			.translate(*trans.transform(center))

	def get_path(self):
		"""
		Return the vertices of the rectangle
		"""
		return self._path

	def get_transform(self):
		"""
		Return the :class:`~matplotlib.transforms.Transform` applied
		to the :class:`Patch`.
		"""
		return self.get_patch_transform()

	def get_patch_transform(self):
		self._recompute_transform()
		return self._patch_transform

	def contains(self, ev):
		if ev.x is None or ev.y is None:
			return False, {}
		x, y = self.get_transform().inverted().transform_point((ev.x, ev.y))
		return (x * x + y * y) <= 1.0, {}

class DispersionPlugin(vaex.ui.plugin.PluginLayer):
	name = "dispersion"
	def __init__(self, parent, layer):
		super(DispersionPlugin, self).__init__(parent, layer)
		dialog.plug_page(self.plug_page, "Dispersions", 2.25, 1.0)
		dialog.plug_grids(self.define_grids, self.draw_grids)

	def define_grids(self, grids):
		#grids.define_grid("counts_vector", self.dialog.gridsize_vector, "VZ*0+1")

		# covariance matrix terms
		# diagonals
		for dimension in range(self.dialog.dimensions):
			axis_name = self.dialog.axisnames[dimension].lower()
			expression = self.expressions[dimension].strip()
			if len(expression) > 0:
				grids.define_grid(axis_name + "_mom1", self.dialog.vector_grid_size, expression)
				grids.define_grid(axis_name + "_mom2", self.dialog.vector_grid_size, "(" + expression + ")**2")
			else:
				grids.define_grid(axis_name + "_mom1", self.dialog.vector_grid_size, None)
				grids.define_grid(axis_name + "_mom2", self.dialog.vector_grid_size, None)

		if 1:
			for dimension1 in range(self.dialog.dimensions):
				for dimension2 in range(dimension1+1, self.dialog.dimensions):
					axis_name1 = self.dialog.axisnames[dimension1].lower()
					axis_name2 = self.dialog.axisnames[dimension2].lower()
					expression1 = self.expressions[dimension1].strip()
					expression2 = self.expressions[dimension2].strip()
					if len(expression1) > 0 and  len(expression2) > 0:
						grids.define_grid("cov_" + axis_name1 +"_" +axis_name2, self.dialog.vector_grid_size, "(" + expression1 + ")*(" + expression2 +")")

	def draw_grids(self, axes, grid_map, grid_map_vector):
		if not self.dispersions_draw:
			return
		self.ellipses = []
		dispersions = []
		counts = grid_map_vector["counts"]
		#print "counts check", np.sum(counts), np.sum(grid_map["counts"])
		#print counts
		#print grid_map_vector.keys()
		if self.dialog.dimensions == 2:
			axis_name1 = self.dialog.axisnames[0].lower()
			axis_name2 = self.dialog.axisnames[1].lower()
			if len(self.expressions[0]) > 0 and len(self.expressions[1]) > 0:
				meanx = grid_map_vector[axis_name1 + "_mom1"]/counts
				meany = grid_map_vector[axis_name2 + "_mom1"]/counts
				varx  = grid_map_vector[axis_name1 + "_mom2"]/counts
				vary  = grid_map_vector[axis_name2 + "_mom2"]/counts
				covxy = grid_map_vector["cov_" +axis_name1 + "_" +axis_name2]/counts - meanx*meany
				sigmax = (varx-meanx**2)**0.5
				sigmay = (vary-meany**2)**0.5
				mask = counts > 0
				x = grid_map_vector["x"]
				y = grid_map_vector["y"]
				x, y = np.meshgrid(x, y)

				vmax = np.nanmax(np.sqrt(sigmax.reshape(-1)**2 + sigmay.reshape(-1)**2))

				width, height = self.dialog.canvas.get_width_height()
				#print "width,height", width, height
				max_size = min(width, height) / float(self.dialog.vector_grid_size)# * 0.9
				#print max_size
				#identity_transform = matplotlib.transforms.IdentityTransform()
				#deltax = self.dialog.ranges_show[0][1] - self.dialog.ranges_show[0][0]
				#deltay = self.dialog.ranges_show[1][1] - self.dialog.ranges_show[1][0]
				#aspect = deltay / float(height) / (deltax/float(width))
				#for grid in [x, y, sigmax, sigmay, covxy, counts, mask]:
				#	print grid.shape
				for x, y, sigmax, sigmay, covxy in zip(x[mask].reshape(-1), y[mask].reshape(-1), sigmax[mask].reshape(-1), sigmay[mask].reshape(-1), covxy[mask].reshape(-1)):
					try:
						covmatrix = [[sigmax**2, covxy], [covxy, sigmay**2]]
						eigen_values, eigen_vectors = np.linalg.eig(covmatrix)
					except:
						pass
					else:
						scaling = 1./vmax * max_size
						device_width =  (np.sqrt(np.max(eigen_values)) * scaling)
						device_height = (np.sqrt(np.min(eigen_values)) * scaling)
						if self.dispersions_unit_length:
							length = np.sqrt(device_width**2+device_height**2)
							device_width  /= float(length) / max_size
							device_height /= float(length) / max_size
						#ellipse_width = np.sqrt(np.max(eigen_values)) * scaling / width * deltax
						#ellipse_height = np.sqrt(np.min(eigen_values)) * scaling / height * deltay
						#ellipse_height /= aspect
						if sigmax < sigmay: # if x was smaller, the largest eigenvalue corresponds to the y value
							device_width, device_height = device_height, device_width
							#ellipse_width, ellipse_height = ellipse_height, ellipse_width
						#ellipse_height /= aspect
						angle = np.arctan(2*covxy / (sigmax**2-sigmay**2))/2.
						#angle2 = np.arctan(2*covxy / (sigmax**2-sigmay**2))/2.
						#angle = angle2 = 0
						#print aspect, sigmax, sigmay, sigmax/sigmay, covxy/(sigmax*sigmay), ellipse_width/ellipse_height
						#aspect = 0.1
						#m = [[np.cos(angle2), np.sin(angle2)*aspect], [-np.sin(angle2), np.cos(angle2)*aspect]]
						#ellipse_width, ellipse_height = np.dot(m, [ellipse_width, ellipse_height])
						#print covxy/(sigmax*sigmay), angle, sigmax, sigmay, covxy
						#device_x, device_y = axes.transData.transform((x, y))
						#print device_x, device_y, device_width, device_height
						#ellipse = patches.Ellipse(xy=(device_x, device_y), width=device_width, height=device_height, angle=angle, transform=identity_transform,
						#                         alpha=0.4, color="blue") #rand()*360
						#ellipse = patches.Ellipse(xy=(x, y), width=ellipse_width, height=ellipse_height, angle=np.degrees(angle),
						#                         alpha=0.4, color="blue") #rand()*360
						ellipse = DispersionEllipse(xy=(x, y), width=device_width, height=device_height, angle=np.degrees(angle), scale=self.scale_dispersion,
												  alpha=0.4, facecolor="green", edgecolor="black") #rand()*360

						axes.add_artist(ellipse)
						self.ellipses.append(ellipse)
						#axes.quiver()

						#[Ellipse(xy=rand(2)*10, width=rand(), height=rand(), angle=rand()*360)

	#@staticmethod
	#def useon(dialog_class):
	#	return issubclass(dialog_class, vaex.plot_windows.VolumeRenderingPlotDialog)

	def plug_page(self, page):
		layout = self.layout = QtGui.QGridLayout()
		page.setLayout(self.layout)
		layout.setSpacing(0)
		layout.setContentsMargins(0,0,0,0)
		layout.setAlignment(QtCore.Qt.AlignTop)

		row = 0

		self.dispersions_draw = bool(eval(self.dialog.options.get("disp_draw", "True")))
		def setter(value):
			self.dispersions_draw = value
			self.dialog.plot()
		self.dispersions_draw_checkbox = self.dialog.create_checkbox(page, "Draw dispersion tensors", lambda : self.dispersions_draw, setter)
		layout.addWidget(self.dispersions_draw_checkbox, row, 1)
		row += 1

		self.dispersions_unit_length = bool(eval(self.dialog.options.get("disp_unit", "False")))
		def setter(value):
			self.dispersions_unit_length = value
			self.dialog.plot()
		self.dispersions_unit_lengthcheckbox = self.dialog.create_checkbox(page, "Unit length", lambda : self.dispersions_unit_length, setter)
		layout.addWidget(self.dispersions_unit_lengthcheckbox, row, 1)
		row += 1

		self.expressions = []
		self.expression_boxes = []

		for dimension in range(self.dialog.dimensions):
			axis_name = self.dialog.axisnames[dimension]
			expression_box = QtGui.QComboBox(page)
			expression_box.setEditable(True)
			expression_box.setMinimumContentsLength(10)
			self.expression_boxes.append(expression_box)

			self.layout.addWidget(QtGui.QLabel(axis_name + '-axis:', page), row, 0)
			self.layout.addWidget(expression_box, row, 1, QtCore.Qt.AlignLeft)

			expression = self.dialog.options.get("disp"+axis_name.lower(), "")
			expression_box.lineEdit().setText(expression)
			self.expressions.append(expression)
			#self.onExpressionChangedPartials.append()
			#expression_box.lineEdit().editingFinished.connect(self.onExpressionChangedPartials[axisIndex])
			calllback = functools.partial(self.onExpressionChanged, axis_index=dimension)
			expression_box.lineEdit().editingFinished.connect(calllback)
			row += 1

		self.scale_dispersion = eval(self.dialog.options.get("disp_scale", "1"))
		def setter(value):
			self.scale_dispersion = value
			for ellipse in self.ellipses:
				ellipse.scale = self.scale_dispersion
			self.dialog.canvas.draw()
			#self.dialog.plot()
		self.scale_dispersion_label, self.scale_dispersion_slider, self.scale_dispersion_value_label =\
				self.dialog.create_slider(page, "scale: ", 1./100, 100., lambda : self.scale_dispersion, setter, format=" {0:>05.2f}", transform=lambda x: 10**x, inverse=lambda x: np.log10(x))
		layout.addWidget(self.scale_dispersion_label, row, 0)
		layout.addWidget(self.scale_dispersion_slider, row, 1)
		layout.addWidget(self.scale_dispersion_value_label, row, 2)
		row += 1

	def onExpressionChanged(self, _=None, axis_index=-1):
		text = str(self.expression_boxes[axis_index].lineEdit().text())
		logger.debug("text set for axis %i: %s" % (axis_index, text))
		if text != self.expressions[axis_index]:
			axis_name = self.dialog.axisnames[axis_index].lower()
			self.expressions[axis_index] = text
			if text == "": # check if we can replot without doing the whole calculation
				self.dialog.plot()
			else:
				non_empty = [k for k in self.expressions if len(k) > 0]
				if len(non_empty) == len(self.expressions):
					self.dialog.compute()
					self.dialog.jobsManager.execute()
		else:
			logger.debug("nothing changed")