# -*- coding: utf-8 -*-

from __future__ import print_function

import operator
import pytest

from pandas.compat import (zip, range, lrange, StringIO)
from pandas import DataFrame, Series, Index, MultiIndex, date_range
import pandas as pd
import numpy as np

from numpy.random import randn

from pandas.util.testing import (assert_series_equal,
                                 assert_frame_equal,
                                 makeCustomDataframe as mkdf)

import pandas.util.testing as tm
from pandas.core.computation import _NUMEXPR_INSTALLED

from pandas.tests.frame.common import TestData

PARSERS = 'python', 'pandas'
ENGINES = 'python', 'numexpr'

@pytest.fixture(params=PARSERS, ids=lambda x: x)
def parser(request):
    return request.param

@pytest.fixture(params=ENGINES, ids=lambda x: x)
def engine(request):
    return request.param

def skip_if_no_pandas_parser(parser):
    if parser != 'pandas':
        pytest.skip("cannot evaluate with parser {0!r}".format(parser))

def skip_if_no_ne(engine='numexpr'):
    if engine == 'numexpr':
        if not _NUMEXPR_INSTALLED:
            pytest.skip("cannot query engine numexpr when numexpr not "
                        "installed")

class TestCompat(object):

    def setup_method(self, method):
        self.df = DataFrame({'A': [1, 2, 3]})
        self.expected1 = self.df[self.df.A > 0]
        self.expected2 = self.df.A + 1

    def test_query_default(self):

        # GH 12749
        # this should always work, whether _NUMEXPR_INSTALLED or not
        df = self.df
        result = df.query('A>0')
        assert_frame_equal(result, self.expected1)
        result = df.eval('A+1')
        assert_series_equal(result, self.expected2, check_names=False)

    def test_query_None(self):

        df = self.df
        result = df.query('A>0', engine=None)
        assert_frame_equal(result, self.expected1)
        result = df.eval('A+1', engine=None)
        assert_series_equal(result, self.expected2, check_names=False)

    def test_query_python(self):

        df = self.df
        result = df.query('A>0', engine='python')
        assert_frame_equal(result, self.expected1)
        result = df.eval('A+1', engine='python')
        assert_series_equal(result, self.expected2, check_names=False)

    def test_query_numexpr(self):

        df = self.df
        if _NUMEXPR_INSTALLED:
            result = df.query('A>0', engine='numexpr')
            assert_frame_equal(result, self.expected1)
            result = df.eval('A+1', engine='numexpr')
            assert_series_equal(result, self.expected2, check_names=False)
        else:
            pytest.raises(ImportError,
                          lambda: df.query('A>0', engine='numexpr'))
            pytest.raises(ImportError,
                          lambda: df.eval('A+1', engine='numexpr'))

class TestDataFrameEval(TestData):

    def test_ops(self):

        # tst ops and reversed ops in evaluation
        # GH7198

        # smaller hits python, larger hits numexpr
        for n in [4, 4000]:

            df = DataFrame(1, index=range(n), columns=list('abcd'))
            df.iloc[0] = 2
            m = df.mean()

            for op_str, op, rop in [('+', '__add__', '__radd__'),
                                    ('-', '__sub__', '__rsub__'),
                                    ('*', '__mul__', '__rmul__'),
                                    ('/', '__truediv__', '__rtruediv__')]:

                base = (DataFrame(np.tile(m.values, n)  # noqa
                                  .reshape(n, -1),
                                  columns=list('abcd')))

                expected = eval("base{op}df".format(op=op_str))

                # ops as strings
                result = eval("m{op}df".format(op=op_str))
                assert_frame_equal(result, expected)

                # these are commutative
                if op in ['+', '*']:
                    result = getattr(df, op)(m)
                    assert_frame_equal(result, expected)

                # these are not
                elif op in ['-', '/']:
                    result = getattr(df, rop)(m)
                    assert_frame_equal(result, expected)

        # GH7192
        df = DataFrame(dict(A=np.random.randn(25000)))
        df.iloc[0:5] = np.nan
        expected = (1 - np.isnan(df.iloc[0:25]))
        result = (1 - np.isnan(df)).iloc[0:25]
        assert_frame_equal(result, expected)

    def test_query_non_str(self):
        # GH 11485
        df = pd.DataFrame({'A': [1, 2, 3], 'B': ['a', 'b', 'b']})

        msg = "expr must be a string to be evaluated"
        with tm.assert_raises_regex(ValueError, msg):
            df.query(lambda x: x.B == "b")

        with tm.assert_raises_regex(ValueError, msg):
            df.query(111)

    def test_query_empty_string(self):
        # GH 13139
        df = pd.DataFrame({'A': [1, 2, 3]})

        msg = "expr cannot be an empty string"
        with tm.assert_raises_regex(ValueError, msg):
            df.query('')

    def test_eval_resolvers_as_list(self):
        # GH 14095
        df = DataFrame(randn(10, 2), columns=list('ab'))
        dict1 = {'a': 1}
        dict2 = {'b': 2}
        assert (df.eval('a + b', resolvers=[dict1, dict2]) ==
                dict1['a'] + dict2['b'])
        assert (pd.eval('a + b', resolvers=[dict1, dict2]) ==
                dict1['a'] + dict2['b'])

class TestDataFrameQueryWithMultiIndex(object):

    def test_query_with_named_multiindex(self, parser, engine):
        tm.skip_if_no_ne(engine)
        skip_if_no_pandas_parser(parser)
        a = np.random.choice(['red', 'green'], size=10)
        b = np.random.choice(['eggs', 'ham'], size=10)
        index = MultiIndex.from_arrays([a, b], names=['color', 'food'])
        df = DataFrame(randn(10, 2), index=index)
        ind = Series(df.index.get_level_values('color').values, index=index,
                     name='color')

        # equality
        res1 = df.query('color == "red"', parser=parser, engine=engine)
        res2 = df.query('"red" == color', parser=parser, engine=engine)
        exp = df[ind == 'red']
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

        # inequality
        res1 = df.query('color != "red"', parser=parser, engine=engine)
        res2 = df.query('"red" != color', parser=parser, engine=engine)
        exp = df[ind != 'red']
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

        # list equality (really just set membership)
        res1 = df.query('color == ["red"]', parser=parser, engine=engine)
        res2 = df.query('["red"] == color', parser=parser, engine=engine)
        exp = df[ind.isin(['red'])]
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

        res1 = df.query('color != ["red"]', parser=parser, engine=engine)
        res2 = df.query('["red"] != color', parser=parser, engine=engine)
        exp = df[~ind.isin(['red'])]
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

        # in/not in ops
        res1 = df.query('["red"] in color', parser=parser, engine=engine)
        res2 = df.query('"red" in color', parser=parser, engine=engine)
        exp = df[ind.isin(['red'])]
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

        res1 = df.query('["red"] not in color', parser=parser, engine=engine)
        res2 = df.query('"red" not in color', parser=parser, engine=engine)
        exp = df[~ind.isin(['red'])]
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

    def test_query_with_unnamed_multiindex(self, parser, engine):
        tm.skip_if_no_ne(engine)
        skip_if_no_pandas_parser(parser)
        a = np.random.choice(['red', 'green'], size=10)
        b = np.random.choice(['eggs', 'ham'], size=10)
        index = MultiIndex.from_arrays([a, b])
        df = DataFrame(randn(10, 2), index=index)
        ind = Series(df.index.get_level_values(0).values, index=index)

        res1 = df.query('ilevel_0 == "red"', parser=parser, engine=engine)
        res2 = df.query('"red" == ilevel_0', parser=parser, engine=engine)
        exp = df[ind == 'red']
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

        # inequality
        res1 = df.query('ilevel_0 != "red"', parser=parser, engine=engine)
        res2 = df.query('"red" != ilevel_0', parser=parser, engine=engine)
        exp = df[ind != 'red']
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

        # list equality (really just set membership)
        res1 = df.query('ilevel_0 == ["red"]', parser=parser, engine=engine)
        res2 = df.query('["red"] == ilevel_0', parser=parser, engine=engine)
        exp = df[ind.isin(['red'])]
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

        res1 = df.query('ilevel_0 != ["red"]', parser=parser, engine=engine)
        res2 = df.query('["red"] != ilevel_0', parser=parser, engine=engine)
        exp = df[~ind.isin(['red'])]
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

        # in/not in ops
        res1 = df.query('["red"] in ilevel_0', parser=parser, engine=engine)
        res2 = df.query('"red" in ilevel_0', parser=parser, engine=engine)
        exp = df[ind.isin(['red'])]
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

        res1 = df.query('["red"] not in ilevel_0', parser=parser,
                        engine=engine)
        res2 = df.query('"red" not in ilevel_0', parser=parser, engine=engine)
        exp = df[~ind.isin(['red'])]
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

        # ## LEVEL 1
        ind = Series(df.index.get_level_values(1).values, index=index)
        res1 = df.query('ilevel_1 == "eggs"', parser=parser, engine=engine)
        res2 = df.query('"eggs" == ilevel_1', parser=parser, engine=engine)
        exp = df[ind == 'eggs']
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

        # inequality
        res1 = df.query('ilevel_1 != "eggs"', parser=parser, engine=engine)
        res2 = df.query('"eggs" != ilevel_1', parser=parser, engine=engine)
        exp = df[ind != 'eggs']
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

        # list equality (really just set membership)
        res1 = df.query('ilevel_1 == ["eggs"]', parser=parser, engine=engine)
        res2 = df.query('["eggs"] == ilevel_1', parser=parser, engine=engine)
        exp = df[ind.isin(['eggs'])]
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

        res1 = df.query('ilevel_1 != ["eggs"]', parser=parser, engine=engine)
        res2 = df.query('["eggs"] != ilevel_1', parser=parser, engine=engine)
        exp = df[~ind.isin(['eggs'])]
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

        # in/not in ops
        res1 = df.query('["eggs"] in ilevel_1', parser=parser, engine=engine)
        res2 = df.query('"eggs" in ilevel_1', parser=parser, engine=engine)
        exp = df[ind.isin(['eggs'])]
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

        res1 = df.query('["eggs"] not in ilevel_1', parser=parser,
                        engine=engine)
        res2 = df.query('"eggs" not in ilevel_1', parser=parser, engine=engine)
        exp = df[~ind.isin(['eggs'])]
        assert_frame_equal(res1, exp)
        assert_frame_equal(res2, exp)

    def test_query_with_partially_named_multiindex(self, parser, engine):
        tm.skip_if_no_ne(engine)
        skip_if_no_pandas_parser(parser)
        a = np.random.choice(['red', 'green'], size=10)
        b = np.arange(10)
        index = MultiIndex.from_arrays([a, b])
        index.names = [None, 'rating']
        df = DataFrame(randn(10, 2), index=index)
        res = df.query('rating == 1', parser=parser, engine=engine)
        ind = Series(df.index.get_level_values('rating').values, index=index,
                     name='rating')
        exp = df[ind == 1]
        assert_frame_equal(res, exp)

        res = df.query('rating != 1', parser=parser, engine=engine)
        ind = Series(df.index.get_level_values('rating').values, index=index,
                     name='rating')
        exp = df[ind != 1]
        assert_frame_equal(res, exp)

        res = df.query('ilevel_0 == "red"', parser=parser, engine=engine)
        ind = Series(df.index.get_level_values(0).values, index=index)
        exp = df[ind == "red"]
        assert_frame_equal(res, exp)

        res = df.query('ilevel_0 != "red"', parser=parser, engine=engine)
        ind = Series(df.index.get_level_values(0).values, index=index)
        exp = df[ind != "red"]
        assert_frame_equal(res, exp)

    def test_query_multiindex_get_index_resolvers(self):
        df = mkdf(10, 3, r_idx_nlevels=2, r_idx_names=['spam', 'eggs'])
        resolvers = df._get_index_resolvers()

        def to_series(mi, level):
            level_values = mi.get_level_values(level)
            s = level_values.to_series()
            s.index = mi
            return s

        col_series = df.columns.to_series()
        expected = {'index': df.index,
                    'columns': col_series,
                    'spam': to_series(df.index, 'spam'),
                    'eggs': to_series(df.index, 'eggs'),
                    'C0': col_series}
        for k, v in resolvers.items():
            if isinstance(v, Index):
                assert v.is_(expected[k])
            elif isinstance(v, Series):
                assert_series_equal(v, expected[k])
            else:
                raise AssertionError("object must be a Series or Index")

    def test_raise_on_panel_with_multiindex(self, parser, engine):
        tm.skip_if_no_ne()
        p = tm.makePanel(7)
        p.items = tm.makeCustomIndex(len(p.items), nlevels=2)
        with pytest.raises(NotImplementedError):
            pd.eval('p + 1', parser=parser, engine=engine)

    def test_raise_on_panel4d_with_multiindex(self, parser, engine):
        tm.skip_if_no_ne()
        p4d = tm.makePanel4D(7)
        p4d.items = tm.makeCustomIndex(len(p4d.items), nlevels=2)
        with pytest.raises(NotImplementedError):
            pd.eval('p4d + 1', parser=parser, engine=engine)

class TestDataFrameQueryNumExprPandas(object):

    @classmethod
    def setup_class(cls):
        cls.engine = 'numexpr'
        cls.parser = 'pandas'
        tm.skip_if_no_ne(cls.engine)

    @classmethod
    def teardown_class(cls):
        del cls.engine, cls.parser

    def test_date_query_with_attribute_access(self):
        engine, parser = self.engine, self.parser
        skip_if_no_pandas_parser(parser)
        df = DataFrame(randn(5, 3))
        df['dates1'] = date_range('1/1/2012', periods=5)
        df['dates2'] = date_range('1/1/2013', periods=5)
        df['dates3'] = date_range('1/1/2014', periods=5)
        res = df.query('@df.dates1 < 20130101 < @df.dates3', engine=engine,
                       parser=parser)
        expec = df[(df.dates1 < '20130101') & ('20130101' < df.dates3)]
        assert_frame_equal(res, expec)

    def test_date_query_no_attribute_access(self):
        engine, parser = self.engine, self.parser
        df = DataFrame(randn(5, 3))
        df['dates1'] = date_range('1/1/2012', periods=5)
        df['dates2'] = date_range('1/1/2013', periods=5)
        df['dates3'] = date_range('1/1/2014', periods=5)
        res = df.query('dates1 < 20130101 < dates3', engine=engine,
                       parser=parser)
        expec = df[(df.dates1 < '20130101') & ('20130101' < df.dates3)]
        assert_frame_equal(res, expec)

    def test_date_query_with_NaT(self):
        engine, parser = self.engine, self.parser
        n = 10
        df = DataFrame(randn(n, 3))
        df['dates1'] = date_range('1/1/2012', periods=n)
        df['dates2'] = date_range('1/1/2013', periods=n)
        df['dates3'] = date_range('1/1/2014', periods=n)
        df.loc[np.random.rand(n) > 0.5, 'dates1'] = pd.NaT
        df.loc[np.random.rand(n) > 0.5, 'dates3'] = pd.NaT
        res = df.query('dates1 < 20130101 < dates3', engine=engine,
                       parser=parser)
        expec = df[(df.dates1 < '20130101') & ('20130101' < df.dates3)]
        assert_frame_equal(res, expec)

    def test_date_index_query(self):
        engine, parser = self.engine, self.parser
        n = 10
        df = DataFrame(randn(n, 3))
        df['dates1'] = date_range('1/1/2012', periods=n)
        df['dates3'] = date_range('1/1/2014', periods=n)
        df.set_index('dates1', inplace=True, drop=True)
        res = df.query('index < 20130101 < dates3', engine=engine,
                       parser=parser)
        expec = df[(df.index < '20130101') & ('20130101' < df.dates3)]
        assert_frame_equal(res, expec)

    def test_date_index_query_with_NaT(self):
        engine, parser = self.engine, self.parser
        n = 10
        df = DataFrame(randn(n, 3))
        df['dates1'] = date_range('1/1/2012', periods=n)
        df['dates3'] = date_range('1/1/2014', periods=n)
        df.iloc[0, 0] = pd.NaT
        df.set_index('dates1', inplace=True, drop=True)
        res = df.query('index < 20130101 < dates3', engine=engine,
                       parser=parser)
        expec = df[(df.index < '20130101') & ('20130101' < df.dates3)]
        assert_frame_equal(res, expec)

    def test_date_index_query_with_NaT_duplicates(self):
        engine, parser = self.engine, self.parser
        n = 10
        d = {}
        d['dates1'] = date_range('1/1/2012', periods=n)
        d['dates3'] = date_range('1/1/2014', periods=n)
        df = DataFrame(d)
        df.loc[np.random.rand(n) > 0.5, 'dates1'] = pd.NaT
        df.set_index('dates1', inplace=True, drop=True)
        res = df.query('dates1 < 20130101 < dates3', engine=engine,
                       parser=parser)
        expec = df[(df.index.to_series() < '20130101') &
                   ('20130101' < df.dates3)]
        assert_frame_equal(res, expec)

    def test_date_query_with_non_date(self):
        engine, parser = self.engine, self.parser

        n = 10
        df = DataFrame({'dates': date_range('1/1/2012', periods=n),
                        'nondate': np.arange(n)})

        ops = '==', '!=', '<', '>', '<=', '>='

        for op in ops:
            with pytest.raises(TypeError):
                df.query('dates %s nondate' % op, parser=parser, engine=engine)

    def test_query_syntax_error(self):
        engine, parser = self.engine, self.parser
        df = DataFrame({"i": lrange(10), "+": lrange(3, 13),
                        "r": lrange(4, 14)})
        with pytest.raises(SyntaxError):
            df.query('i - +', engine=engine, parser=parser)

    def test_query_scope(self):
        from pandas.core.computation.ops import UndefinedVariableError
        engine, parser = self.engine, self.parser
        skip_if_no_pandas_parser(parser)

        df = DataFrame(np.random.randn(20, 2), columns=list('ab'))

        a, b = 1, 2  # noqa
        res = df.query('a > b', engine=engine, parser=parser)
        expected = df[df.a > df.b]
        assert_frame_equal(res, expected)

        res = df.query('@a > b', engine=engine, parser=parser)
        expected = df[a > df.b]
        assert_frame_equal(res, expected)

        # no local variable c
        with pytest.raises(UndefinedVariableError):
            df.query('@a > b > @c', engine=engine, parser=parser)

        # no column named 'c'
        with pytest.raises(UndefinedVariableError):
            df.query('@a > b > c', engine=engine, parser=parser)

    def test_query_doesnt_pickup_local(self):
        from pandas.core.computation.ops import UndefinedVariableError

        engine, parser = self.engine, self.parser
        n = m = 10
        df = DataFrame(np.random.randint(m, size=(n, 3)), columns=list('abc'))

        # we don't pick up the local 'sin'
        with pytest.raises(UndefinedVariableError):
            df.query('sin > 5', engine=engine, parser=parser)

    def test_query_builtin(self):
        from pandas.core.computation.engines import NumExprClobberingError
        engine, parser = self.engine, self.parser

        n = m = 10
        df = DataFrame(np.random.randint(m, size=(n, 3)), columns=list('abc'))

        df.index.name = 'sin'
        with tm.assert_raises_regex(NumExprClobberingError,
                                    'Variables in expression.+'):
            df.query('sin > 5', engine=engine, parser=parser)

    def test_query(self):
        engine, parser = self.engine, self.parser
        df = DataFrame(np.random.randn(10, 3), columns=['a', 'b', 'c'])

        assert_frame_equal(df.query('a < b', engine=engine, parser=parser),
                           df[df.a < df.b])
        assert_frame_equal(df.query('a + b > b * c', engine=engine,
                                    parser=parser),
                           df[df.a + df.b > df.b * df.c])

    def test_query_index_with_name(self):
        engine, parser = self.engine, self.parser
        df = DataFrame(np.random.randint(10, size=(10, 3)),
                       index=Index(range(10), name='blob'),
                       columns=['a', 'b', 'c'])
        res = df.query('(blob < 5) & (a < b)', engine=engine, parser=parser)
        expec = df[(df.index < 5) & (df.a < df.b)]
        assert_frame_equal(res, expec)

        res = df.query('blob < b', engine=engine, parser=parser)
        expec = df[df.index < df.b]

        assert_frame_equal(res, expec)

    def test_query_index_without_name(self):
        engine, parser = self.engine, self.parser
        df = DataFrame(np.random.randint(10, size=(10, 3)),
                       index=range(10), columns=['a', 'b', 'c'])

        # "index" should refer to the index
        res = df.query('index < b', engine=engine, parser=parser)
        expec = df[df.index < df.b]
        assert_frame_equal(res, expec)

        # test against a scalar
        res = df.query('index < 5', engine=engine, parser=parser)
        expec = df[df.index < 5]
        assert_frame_equal(res, expec)

    def test_nested_scope(self):
        engine = self.engine
        parser = self.parser

        skip_if_no_pandas_parser(parser)

        df = DataFrame(np.random.randn(5, 3))
        df2 = DataFrame(np.random.randn(5, 3))
        expected = df[(df > 0) & (df2 > 0)]

        result = df.query('(@df > 0) & (@df2 > 0)', engine=engine,
                          parser=parser)
        assert_frame_equal(result, expected)

        result = pd.eval('df[df > 0 and df2 > 0]', engine=engine,
                         parser=parser)
        assert_frame_equal(result, expected)

        result = pd.eval('df[df > 0 and df2 > 0 and df[df > 0] > 0]',
                         engine=engine, parser=parser)
        expected = df[(df > 0) & (df2 > 0) & (df[df > 0] > 0)]
        assert_frame_equal(result, expected)

        result = pd.eval('df[(df>0) & (df2>0)]', engine=engine, parser=parser)
        expected = df.query('(@df>0) & (@df2>0)', engine=engine, parser=parser)
        assert_frame_equal(result, expected)

    def test_nested_raises_on_local_self_reference(self):
        from pandas.core.computation.ops import UndefinedVariableError

        df = DataFrame(np.random.randn(5, 3))

        # can't reference ourself b/c we're a local so @ is necessary
        with pytest.raises(UndefinedVariableError):
            df.query('df > 0', engine=self.engine, parser=self.parser)

    def test_local_syntax(self):
        skip_if_no_pandas_parser(self.parser)

        engine, parser = self.engine, self.parser
        df = DataFrame(randn(100, 10), columns=list('abcdefghij'))
        b = 1
        expect = df[df.a < b]
        result = df.query('a < @b', engine=engine, parser=parser)
        assert_frame_equal(result, expect)

        expect = df[df.a < df.b]
        result = df.query('a < b', engine=engine, parser=parser)
        assert_frame_equal(result, expect)

    def test_chained_cmp_and_in(self):
        skip_if_no_pandas_parser(self.parser)
        engine, parser = self.engine, self.parser
        cols = list('abc')
        df = DataFrame(randn(100, len(cols)), columns=cols)
        res = df.query('a < b < c and a not in b not in c', engine=engine,
                       parser=parser)
        ind = (df.a < df.b) & (df.b < df.c) & ~df.b.isin(df.a) & ~df.c.isin(df.b)  # noqa
        expec = df[ind]
        assert_frame_equal(res, expec)

    def test_local_variable_with_in(self):
        engine, parser = self.engine, self.parser
        skip_if_no_pandas_parser(parser)
        a = Series(np.random.randint(3, size=15), name='a')
        b = Series(np.random.randint(10, size=15), name='b')
        df = DataFrame({'a': a, 'b': b})

        expected = df.loc[(df.b - 1).isin(a)]
        result = df.query('b - 1 in a', engine=engine, parser=parser)
        assert_frame_equal(expected, result)

        b = Series(np.random.randint(10, size=15), name='b')
        expected = df.loc[(b - 1).isin(a)]
        result = df.query('@b - 1 in a', engine=engine, parser=parser)
        assert_frame_equal(expected, result)

    def test_at_inside_string(self):
        engine, parser = self.engine, self.parser
        skip_if_no_pandas_parser(parser)
        c = 1  # noqa
        df = DataFrame({'a': ['a', 'a', 'b', 'b', '@c', '@c']})
        result = df.query('a == "@c"', engine=engine, parser=parser)
        expected = df[df.a == "@c"]
        assert_frame_equal(result, expected)

    def test_query_undefined_local(self):
        from pandas.core.computation.ops import UndefinedVariableError
        engine, parser = self.engine, self.parser
        skip_if_no_pandas_parser(parser)
        df = DataFrame(np.random.rand(10, 2), columns=list('ab'))
        with tm.assert_raises_regex(UndefinedVariableError,
                                    "local variable 'c' is not defined"):
            df.query('a == @c', engine=engine, parser=parser)

    def test_index_resolvers_come_after_columns_with_the_same_name(self):
        n = 1  # noqa
        a = np.r_[20:101:20]

        df = DataFrame({'index': a, 'b': np.random.randn(a.size)})
        df.index.name = 'index'
        result = df.query('index > 5', engine=self.engine, parser=self.parser)
        expected = df[df['index'] > 5]
        assert_frame_equal(result, expected)

        df = DataFrame({'index': a,
                        'b': np.random.randn(a.size)})
        result = df.query('ilevel_0 > 5', engine=self.engine,
                          parser=self.parser)
        expected = df.loc[df.index[df.index > 5]]
        assert_frame_equal(result, expected)

        df = DataFrame({'a': a, 'b': np.random.randn(a.size)})
        df.index.name = 'a'
        result = df.query('a > 5', engine=self.engine, parser=self.parser)
        expected = df[df.a > 5]
        assert_frame_equal(result, expected)

        result = df.query('index > 5', engine=self.engine, parser=self.parser)
        expected = df.loc[df.index[df.index > 5]]
        assert_frame_equal(result, expected)

    def test_inf(self):
        n = 10
        df = DataFrame({'a': np.random.rand(n), 'b': np.random.rand(n)})
        df.loc[::2, 0] = np.inf
        ops = '==', '!='
        d = dict(zip(ops, (operator.eq, operator.ne)))
        for op, f in d.items():
            q = 'a %s inf' % op
            expected = df[f(df.a, np.inf)]
            result = df.query(q, engine=self.engine, parser=self.parser)
            assert_frame_equal(result, expected)

class TestDataFrameQueryNumExprPython(TestDataFrameQueryNumExprPandas):

    @classmethod
    def setup_class(cls):
        super(TestDataFrameQueryNumExprPython, cls).setup_class()
        cls.engine = 'numexpr'
        cls.parser = 'python'
        tm.skip_if_no_ne(cls.engine)
        cls.frame = TestData().frame

    def test_date_query_no_attribute_access(self):
        engine, parser = self.engine, self.parser
        df = DataFrame(randn(5, 3))
        df['dates1'] = date_range('1/1/2012', periods=5)
        df['dates2'] = date_range('1/1/2013', periods=5)
        df['dates3'] = date_range('1/1/2014', periods=5)
        res = df.query('(dates1 < 20130101) & (20130101 < dates3)',
                       engine=engine, parser=parser)
        expec = df[(df.dates1 < '20130101') & ('20130101' < df.dates3)]
        assert_frame_equal(res, expec)

    def test_date_query_with_NaT(self):
        engine, parser = self.engine, self.parser
        n = 10
        df = DataFrame(randn(n, 3))
        df['dates1'] = date_range('1/1/2012', periods=n)
        df['dates2'] = date_range('1/1/2013', periods=n)
        df['dates3'] = date_range('1/1/2014', periods=n)
        df.loc[np.random.rand(n) > 0.5, 'dates1'] = pd.NaT
        df.loc[np.random.rand(n) > 0.5, 'dates3'] = pd.NaT
        res = df.query('(dates1 < 20130101) & (20130101 < dates3)',
                       engine=engine, parser=parser)
        expec = df[(df.dates1 < '20130101') & ('20130101' < df.dates3)]
        assert_frame_equal(res, expec)

    def test_date_index_query(self):
        engine, parser = self.engine, self.parser
        n = 10
        df = DataFrame(randn(n, 3))
        df['dates1'] = date_range('1/1/2012', periods=n)
        df['dates3'] = date_range('1/1/2014', periods=n)
        df.set_index('dates1', inplace=True, drop=True)
        res = df.query('(index < 20130101) & (20130101 < dates3)',
                       engine=engine, parser=parser)
        expec = df[(df.index < '20130101') & ('20130101' < df.dates3)]
        assert_frame_equal(res, expec)

    def test_date_index_query_with_NaT(self):
        engine, parser = self.engine, self.parser
        n = 10
        df = DataFrame(randn(n, 3))
        df['dates1'] = date_range('1/1/2012', periods=n)
        df['dates3'] = date_range('1/1/2014', periods=n)
        df.iloc[0, 0] = pd.NaT
        df.set_index('dates1', inplace=True, drop=True)
        res = df.query('(index < 20130101) & (20130101 < dates3)',
                       engine=engine, parser=parser)
        expec = df[(df.index < '20130101') & ('20130101' < df.dates3)]
        assert_frame_equal(res, expec)

    def test_date_index_query_with_NaT_duplicates(self):
        engine, parser = self.engine, self.parser
        n = 10
        df = DataFrame(randn(n, 3))
        df['dates1'] = date_range('1/1/2012', periods=n)
        df['dates3'] = date_range('1/1/2014', periods=n)
        df.loc[np.random.rand(n) > 0.5, 'dates1'] = pd.NaT
        df.set_index('dates1', inplace=True, drop=True)
        with pytest.raises(NotImplementedError):
            df.query('index < 20130101 < dates3', engine=engine, parser=parser)

    def test_nested_scope(self):
        from pandas.core.computation.ops import UndefinedVariableError
        engine = self.engine
        parser = self.parser
        # smoke test
        x = 1  # noqa
        result = pd.eval('x + 1', engine=engine, parser=parser)
        assert result == 2

        df = DataFrame(np.random.randn(5, 3))
        df2 = DataFrame(np.random.randn(5, 3))

        # don't have the pandas parser
        with pytest.raises(SyntaxError):
            df.query('(@df>0) & (@df2>0)', engine=engine, parser=parser)

        with pytest.raises(UndefinedVariableError):
            df.query('(df>0) & (df2>0)', engine=engine, parser=parser)

        expected = df[(df > 0) & (df2 > 0)]
        result = pd.eval('df[(df > 0) & (df2 > 0)]', engine=engine,
                         parser=parser)
        assert_frame_equal(expected, result)

        expected = df[(df > 0) & (df2 > 0) & (df[df > 0] > 0)]
        result = pd.eval('df[(df > 0) & (df2 > 0) & (df[df > 0] > 0)]',
                         engine=engine, parser=parser)
        assert_frame_equal(expected, result)

class TestDataFrameQueryPythonPandas(TestDataFrameQueryNumExprPandas):

    @classmethod
    def setup_class(cls):
        super(TestDataFrameQueryPythonPandas, cls).setup_class()
        cls.engine = 'python'
        cls.parser = 'pandas'
        cls.frame = TestData().frame

    def test_query_builtin(self):
        engine, parser = self.engine, self.parser

        n = m = 10
        df = DataFrame(np.random.randint(m, size=(n, 3)), columns=list('abc'))

        df.index.name = 'sin'
        expected = df[df.index > 5]
        result = df.query('sin > 5', engine=engine, parser=parser)
        assert_frame_equal(expected, result)

class TestDataFrameQueryPythonPython(TestDataFrameQueryNumExprPython):

    @classmethod
    def setup_class(cls):
        super(TestDataFrameQueryPythonPython, cls).setup_class()
        cls.engine = cls.parser = 'python'
        cls.frame = TestData().frame

    def test_query_builtin(self):
        engine, parser = self.engine, self.parser

        n = m = 10
        df = DataFrame(np.random.randint(m, size=(n, 3)), columns=list('abc'))

        df.index.name = 'sin'
        expected = df[df.index > 5]
        result = df.query('sin > 5', engine=engine, parser=parser)
        assert_frame_equal(expected, result)

class TestDataFrameQueryStrings(object):

    def test_str_query_method(self, parser, engine):
        tm.skip_if_no_ne(engine)
        df = DataFrame(randn(10, 1), columns=['b'])
        df['strings'] = Series(list('aabbccddee'))
        expect = df[df.strings == 'a']

        if parser != 'pandas':
            col = 'strings'
            lst = '"a"'

            lhs = [col] * 2 + [lst] * 2
            rhs = lhs[::-1]

            eq, ne = '==', '!='
            ops = 2 * ([eq] + [ne])

            for lhs, op, rhs in zip(lhs, ops, rhs):
                ex = '{lhs} {op} {rhs}'.format(lhs=lhs, op=op, rhs=rhs)
                pytest.raises(NotImplementedError, df.query, ex,
                              engine=engine, parser=parser,
                              local_dict={'strings': df.strings})
        else:
            res = df.query('"a" == strings', engine=engine, parser=parser)
            assert_frame_equal(res, expect)

            res = df.query('strings == "a"', engine=engine, parser=parser)
            assert_frame_equal(res, expect)
            assert_frame_equal(res, df[df.strings.isin(['a'])])

            expect = df[df.strings != 'a']
            res = df.query('strings != "a"', engine=engine, parser=parser)
            assert_frame_equal(res, expect)

            res = df.query('"a" != strings', engine=engine, parser=parser)
            assert_frame_equal(res, expect)
            assert_frame_equal(res, df[~df.strings.isin(['a'])])

    def test_str_list_query_method(self, parser, engine):
        tm.skip_if_no_ne(engine)
        df = DataFrame(randn(10, 1), columns=['b'])
        df['strings'] = Series(list('aabbccddee'))
        expect = df[df.strings.isin(['a', 'b'])]

        if parser != 'pandas':
            col = 'strings'
            lst = '["a", "b"]'

            lhs = [col] * 2 + [lst] * 2
            rhs = lhs[::-1]

            eq, ne = '==', '!='
            ops = 2 * ([eq] + [ne])

            for lhs, op, rhs in zip(lhs, ops, rhs):
                ex = '{lhs} {op} {rhs}'.format(lhs=lhs, op=op, rhs=rhs)
                with pytest.raises(NotImplementedError):
                    df.query(ex, engine=engine, parser=parser)
        else:
            res = df.query('strings == ["a", "b"]', engine=engine,
                           parser=parser)
            assert_frame_equal(res, expect)

            res = df.query('["a", "b"] == strings', engine=engine,
                           parser=parser)
            assert_frame_equal(res, expect)

            expect = df[~df.strings.isin(['a', 'b'])]

            res = df.query('strings != ["a", "b"]', engine=engine,
                           parser=parser)
            assert_frame_equal(res, expect)

            res = df.query('["a", "b"] != strings', engine=engine,
                           parser=parser)
            assert_frame_equal(res, expect)

    def test_query_with_string_columns(self, parser, engine):
        tm.skip_if_no_ne(engine)
        df = DataFrame({'a': list('aaaabbbbcccc'),
                        'b': list('aabbccddeeff'),
                        'c': np.random.randint(5, size=12),
                        'd': np.random.randint(9, size=12)})
        if parser == 'pandas':
            res = df.query('a in b', parser=parser, engine=engine)
            expec = df[df.a.isin(df.b)]
            assert_frame_equal(res, expec)

            res = df.query('a in b and c < d', parser=parser, engine=engine)
            expec = df[df.a.isin(df.b) & (df.c < df.d)]
            assert_frame_equal(res, expec)
        else:
            with pytest.raises(NotImplementedError):
                df.query('a in b', parser=parser, engine=engine)

            with pytest.raises(NotImplementedError):
                df.query('a in b and c < d', parser=parser, engine=engine)

    def test_object_array_eq_ne(self, parser, engine):
        tm.skip_if_no_ne(engine)
        df = DataFrame({'a': list('aaaabbbbcccc'),
                        'b': list('aabbccddeeff'),
                        'c': np.random.randint(5, size=12),
                        'd': np.random.randint(9, size=12)})
        res = df.query('a == b', parser=parser, engine=engine)
        exp = df[df.a == df.b]
        assert_frame_equal(res, exp)

        res = df.query('a != b', parser=parser, engine=engine)
        exp = df[df.a != df.b]
        assert_frame_equal(res, exp)

    def test_query_with_nested_strings(self, parser, engine):
        tm.skip_if_no_ne(engine)
        skip_if_no_pandas_parser(parser)
        raw = """id          event          timestamp
        1   "page 1 load"   1/1/2014 0:00:01
        1   "page 1 exit"   1/1/2014 0:00:31
        2   "page 2 load"   1/1/2014 0:01:01
        2   "page 2 exit"   1/1/2014 0:01:31
        3   "page 3 load"   1/1/2014 0:02:01
        3   "page 3 exit"   1/1/2014 0:02:31
        4   "page 1 load"   2/1/2014 1:00:01
        4   "page 1 exit"   2/1/2014 1:00:31
        5   "page 2 load"   2/1/2014 1:01:01
        5   "page 2 exit"   2/1/2014 1:01:31
        6   "page 3 load"   2/1/2014 1:02:01
        6   "page 3 exit"   2/1/2014 1:02:31
        """
        df = pd.read_csv(StringIO(raw), sep=r'\s{2,}', engine='python',
                         parse_dates=['timestamp'])
        expected = df[df.event == '"page 1 load"']
        res = df.query("""'"page 1 load"' in event""", parser=parser,
                       engine=engine)
        assert_frame_equal(expected, res)

    def test_query_with_nested_special_character(self, parser, engine):
        skip_if_no_pandas_parser(parser)
        tm.skip_if_no_ne(engine)
        df = DataFrame({'a': ['a', 'b', 'test & test'],
                        'b': [1, 2, 3]})
        res = df.query('a == "test & test"', parser=parser, engine=engine)
        expec = df[df.a == 'test & test']
        assert_frame_equal(res, expec)

    def test_query_lex_compare_strings(self, parser, engine):
        tm.skip_if_no_ne(engine=engine)
        import operator as opr

        a = Series(np.random.choice(list('abcde'), 20))
        b = Series(np.arange(a.size))
        df = DataFrame({'X': a, 'Y': b})

        ops = {'<': opr.lt, '>': opr.gt, '<=': opr.le, '>=': opr.ge}

        for op, func in ops.items():
            res = df.query('X %s "d"' % op, engine=engine, parser=parser)
            expected = df[func(df.X, 'd')]
            assert_frame_equal(res, expected)

    def test_query_single_element_booleans(self, parser, engine):
        tm.skip_if_no_ne(engine)
        columns = 'bid', 'bidsize', 'ask', 'asksize'
        data = np.random.randint(2, size=(1, len(columns))).astype(bool)
        df = DataFrame(data, columns=columns)
        res = df.query('bid & ask', engine=engine, parser=parser)
        expected = df[df.bid & df.ask]
        assert_frame_equal(res, expected)

    def test_query_string_scalar_variable(self, parser, engine):
        tm.skip_if_no_ne(engine)
        skip_if_no_pandas_parser(parser)
        df = pd.DataFrame({'Symbol': ['BUD US', 'BUD US', 'IBM US', 'IBM US'],
                           'Price': [109.70, 109.72, 183.30, 183.35]})
        e = df[df.Symbol == 'BUD US']
        symb = 'BUD US'  # noqa
        r = df.query('Symbol == @symb', parser=parser, engine=engine)
        assert_frame_equal(e, r)

class TestDataFrameEvalNumExprPandas(object):

    @classmethod
    def setup_class(cls):
        cls.engine = 'numexpr'
        cls.parser = 'pandas'
        tm.skip_if_no_ne()

    def setup_method(self, method):
        self.frame = DataFrame(randn(10, 3), columns=list('abc'))

    def teardown_method(self, method):
        del self.frame

    def test_simple_expr(self):
        res = self.frame.eval('a + b', engine=self.engine, parser=self.parser)
        expect = self.frame.a + self.frame.b
        assert_series_equal(res, expect)

    def test_bool_arith_expr(self):
        res = self.frame.eval('a[a < 1] + b', engine=self.engine,
                              parser=self.parser)
        expect = self.frame.a[self.frame.a < 1] + self.frame.b
        assert_series_equal(res, expect)

    def test_invalid_type_for_operator_raises(self):
        df = DataFrame({'a': [1, 2], 'b': ['c', 'd']})
        ops = '+', '-', '*', '/'
        for op in ops:
            with tm.assert_raises_regex(TypeError,
                                        "unsupported operand type\(s\) "
                                        "for .+: '.+' and '.+'"):
                df.eval('a {0} b'.format(op), engine=self.engine,
                        parser=self.parser)

class TestDataFrameEvalNumExprPython(TestDataFrameEvalNumExprPandas):

    @classmethod
    def setup_class(cls):
        super(TestDataFrameEvalNumExprPython, cls).setup_class()
        cls.engine = 'numexpr'
        cls.parser = 'python'
        tm.skip_if_no_ne(cls.engine)

class TestDataFrameEvalPythonPandas(TestDataFrameEvalNumExprPandas):

    @classmethod
    def setup_class(cls):
        super(TestDataFrameEvalPythonPandas, cls).setup_class()
        cls.engine = 'python'
        cls.parser = 'pandas'

class TestDataFrameEvalPythonPython(TestDataFrameEvalNumExprPython):

    @classmethod
    def setup_class(cls):
        cls.engine = cls.parser = 'python'


"""
========================================
Plotting topographic maps of evoked data
========================================

Load evoked data and plot topomaps for selected time points.

"""
# Authors: Christian Brodbeck <christianbrodbeck@nyu.edu>
#          Tal Linzen <linzen@nyu.edu>
#          Denis A. Engeman <denis.engemann@gmail.com>
#
# License: BSD (3-clause)

import numpy as np
import matplotlib.pyplot as plt
from mne.datasets import sample
from mne import read_evokeds

print(__doc__)

path = sample.data_path()
fname = path + '/MEG/sample/sample_audvis-ave.fif'

# load evoked and subtract baseline
condition = 'Left Auditory'
evoked = read_evokeds(fname, condition=condition, baseline=(None, 0))

# set time instants in seconds (from 50 to 150ms in a step of 10ms)
times = np.arange(0.05, 0.15, 0.01)
# If times is set to None only 10 regularly spaced topographies will be shown

# plot magnetometer data as topomaps
evoked.plot_topomap(times, ch_type='mag')

# compute a 50 ms bin to stabilize topographies
evoked.plot_topomap(times, ch_type='mag', average=0.05)

# plot gradiometer data (plots the RMS for each pair of gradiometers)
evoked.plot_topomap(times, ch_type='grad')

# plot magnetometer data as an animation
evoked.animate_topomap(ch_type='mag', times=times, frame_rate=10)

# plot magnetometer data as topomap at 1 time point : 100 ms
# and add channel labels and title
evoked.plot_topomap(0.1, ch_type='mag', show_names=True, colorbar=False,
                    size=6, res=128, title='Auditory response')
plt.subplots_adjust(left=0.01, right=0.99, bottom=0.01, top=0.88)


"""
python speedup_kmeans.py --profile
python speedup_kmeans.py

git worktree add workdir_master master
rob sedr "\<sklearn\>" sklearn_master True
git mv sklearn sklearn_master
python setup develop
python -c "import sklearn_master; print(sklearn_master.__file__)"
python -c "import sklearn; print(sklearn.__file__)"
"""
from __future__ import absolute_import, division, print_function, unicode_literals
import utool as ut
import sklearn  # NOQA
from sklearn.datasets.samples_generator import make_blobs
from sklearn.utils.extmath import row_norms, squared_norm  # NOQA
import sklearn.cluster
import numpy as np
from sklearn.metrics.pairwise import euclidean_distances  # NOQA

import sklearn_master.cluster
(print, rrr, profile) = ut.inject2(__name__, '[tester]')

def test_kmeans_plus_plus_speed(n_clusters=2000, n_features=128, per_cluster=10, asint=False, fix=True):
    """
    from speedup_kmeans import *
    from sklearn.cluster.k_means_ import *
    """
    rng = np.random.RandomState(42)
    # Make random cluster centers on a ball
    centers = rng.rand(n_clusters, n_features)
    centers /= np.linalg.norm(centers, axis=0)[None, :]
    centers = (centers * 512).astype(np.uint8) / 512
    centers /= np.linalg.norm(centers, axis=0)[None, :]

    n_samples = int(n_clusters * per_cluster)
    n_clusters, n_features = centers.shape
    X, true_labels = make_blobs(n_samples=n_samples, centers=centers,
                                cluster_std=1., random_state=42)

    if asint:
        X = (X * 512).astype(np.int32)

    x_squared_norms = row_norms(X, squared=True)

    if fix:
        _k_init = sklearn.cluster.k_means_._k_init
    else:
        _k_init = sklearn_master.cluster.k_means_._k_init
    random_state = np.random.RandomState(42)
    n_local_trials = None  # NOQA

    with ut.Timer('testing kmeans init') as t:
        centers = _k_init(X, n_clusters, random_state=random_state, x_squared_norms=x_squared_norms)
    return centers, t.ellapsed

def main():
    if True:
        import pandas as pd
        pd.options.display.max_rows = 1000
        pd.options.display.width = 1000

        basis = {
            #'n_clusters': [10, 100, 1000, 2000][::-1],
            #'n_features': [4, 32, 128, 512][::-1],
            #'per_cluster': [1, 10, 100, 200][::-1],
            'n_clusters': [10, 100, 500][::-1],
            'n_features': [32, 128][::-1],
            'per_cluster': [1, 10, 20][::-1],
            'asint': [True, False],
        }
        vals = []
        for kw in ut.ProgIter(ut.all_dict_combinations(basis), lbl='gridsearch',
                              bs=False, adjust=False, freq=1):
            print('kw = ' + ut.repr2(kw))
            exec(ut.execstr_dict(kw))
            centers1, new_speed = test_kmeans_plus_plus_speed(fix=True, **kw)
            centers2, old_speed = test_kmeans_plus_plus_speed(fix=False, **kw)
            import utool
            with utool.embed_on_exception_context:
                assert np.all(centers1 == centers2), 'new code disagrees'

            kw['new_speed'] = new_speed
            kw['old_speed'] = old_speed
            vals.append(kw)
            print('---------')

        df = pd.DataFrame.from_dict(vals)
        df['percent_change'] = 100 * (df['old_speed'] - df['new_speed']) / df['old_speed']
        df = df.reindex_axis(list(basis.keys()) + ['new_speed', 'old_speed', 'percent_change'], axis=1)
        df['absolute_change'] = (df['old_speed'] - df['new_speed'])
        print(df.sort('absolute_change', ascending=False))
        #print(df)

        print(df['percent_change'][df['absolute_change'] > .1].mean())

    #print(df.loc[df['percent_change'].argsort()[::-1]])
    else:
        new_speed = test_kmeans_plus_plus_speed()
        try:
            profile.dump_stats('out.lprof')
            profile.print_stats(stripzeros=True)
        except Exception:
            pass
        print('new_speed = %r' % (new_speed,))

if __name__ == '__main__':
    main()


# Copyright (c) 2016, Jeffrey Maggio and Joseph Bartelmo
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
# associated documentation files (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or substantial
# portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT
# LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import Tkinter as tk
from PIL import Image, ImageTk
import cv2
import sys
import numpy
import time
import matplotlib.path as matpath

class GlobalSurveyor(object):
    def __init__(self, root, ideal_image,  number_polys = 2, thresholds = (50, 200), polys = None):
        self.ideal_im = ideal_image
        self.parent = root
        self.done = False
        self.points = []
        self.warningColor = (128, 128, 0)
        self.errorColor = (255, 0, 0)

        self.number_polys = number_polys
        #self.number_sides = number_sides
        if polys == None:
            self.polys = []
        else:
            self.polys = polys

        self.thresh = thresholds
        self.ideal_counts = []
        self.calibrates_polygons()

    def calibrates_polygons(self):
        polys = self.define_polygons(self.ideal_im)
        ideal_kp = self.apply_fast(self.ideal_im)
        self.ideal_counts = self.keypoints_to_counts(ideal_kp)

    def _on_left_mouse(self, event):   # CITE THIS FROM STACKEXCHANGE
        # Mouse callback that gets called for every mouse event (i.e. moving, clicking, etc.)
        if self.done: # Nothing more to do
            return
        # Left click means adding a point at current position to the list of points
        print("Adding point #%d with position(%d,%d)" % (len(self.points), event.x, event.y))
        self.points.append((event.x, event.y))

    def _on_right_mouse(self, event):   # CITE THIS FROM STACKEXCHANGE
        if len(self.points) > 2:
            # Right click means we're done
            print("Completing polygon with %d points." % len(self.points))
            self.done = True

    def define_polygons(self, image):
        # use this to automate point picking and/or display point picker window in gui

        FINAL_LINE_COLOR = (255, 0, 0)
        WORKING_LINE_COLOR = (0, 0, 255)

        top = tk.Toplevel(self.parent)
        top.title("Select bounding points for the target area")
        print image.shape
        top.update_idletasks()
        width = image.shape[1]
        height = image.shape[0]
        x = (top.winfo_screenwidth() // 2) - (width // 2)
        y = (top.winfo_screenheight() // 2) - (height // 2)
        top.geometry('{}x{}+{}+{}'.format(width, height, x, y))

        #we are defaulting to the image size, i don't want to deal wiht resizing
        initial_im = tk.PhotoImage()
        image_label = tk.Label(top, image=initial_im)
        image_label.grid(row = 0, column = 0)
        image_label.bind("<Button-1>", self._on_left_mouse)
        image_label.bind("<Button-3>", self._on_right_mouse)
        top.grid()

        copy = image.copy()
        def display_image():
            imageFromArray = Image.fromarray(copy)
            try:
                tkImage = ImageTk.PhotoImage(image=imageFromArray)
                image_label.configure(image=tkImage)

                image_label._image_cache = tkImage  # avoid garbage collection

                top.update()
                return True
            except RuntimeError:
                print('Unable to update image frame. Assuming application has been killed unexpectidly.')
                return False

        for num in range(self.number_polys):
            if display_image() == False:
                top.destroy()
                return

            self.points = []
            self.done = False

            while (not self.done):
                if (len(self.points) > 0):
                    cv2.polylines(copy, numpy.asarray([self.points]), False, FINAL_LINE_COLOR, 1)

                imageFromArray = Image.fromarray(copy)
                if display_image() == False:
                    top.destroy()
                    return
            self.polys.append(self.points)
            print(self.polys)
            update = numpy.asarray(self.polys[num]).astype(numpy.int32)

            if (len(self.points) > 0):
                cv2.fillPoly(copy, [update], (0,0,255))
        display_image()

        time.sleep(0.5)
        top.destroy()

        return self.polys

    def apply_fast(self, image):
        fast = cv2.FastFeatureDetector()
        kp = fast.detect(image, None)

        return kp

    def keypoints_to_counts(self, im_keypoints): #polys = numpy.asarray[[[0,0],[0,ideal_im.shape[0]//2],[ideal_im.shape[0]//2,ideal_im.shape[1]],[ideal_im.shape[0]//2, 0]]]):
        #pass in im_keypoints and count how many are in each polygon
        #polygon points should also be an input

        im_keys = []
        for i in range(len(im_keypoints)):
            im_keys.append(im_keypoints[i].pt)

        # NOTE:: POLYS MUST BE IN X (COL), Y (ROW) FORMAT FOR FILLPOLY TO WORK

        # NOTE:: POLYS MUST ALSO BE OF DTYPE FLOAT

        data = self.polys
        polys = map(lambda data: map(lambda data: map(float, data), data), data)
        counts = []

        for poly in polys:
            # MATPLOTLIB PATH EXAMPLE GOES HERE
            bbPath = matpath.Path(poly)
            # sum boolean array to get number of kp within current poly
            try:
                truth_vector = bbPath.contains_points(im_keys)   # check input for contains_points. might only like tuples
            except:
                truth_vector = [1] * len(im_keys)
            counts.append(numpy.sum(truth_vector))

        return numpy.asarray(counts)

    def threshold(self, counts, thresh_tuple):
        # do the thresholding thing
        delta = abs(counts - self.ideal_counts)

        threshed_polys = numpy.where((delta >= thresh_tuple[0]) & (delta < thresh_tuple[1]))[0]   # ok to index like this because tuple represents vector

        # save which polys get marked true out so these can be used as overlay in fillpolys
        # this can be done by indexing into polys with truth vector saved here

        return threshed_polys.astype(int)   # a true/false vector of length polys

    def generate_output(self, image, threshed_polys, color, alpha = 0.5): #, colorAlert = (0,0,255)):
        # use polyfill to blend im_in and overlay (which is just the polygons we want colored in)
        colored_polys = []

        for threshed_poly in threshed_polys:
            colored_polys.append(numpy.array(self.polys[threshed_poly]).astype(numpy.int32))

        overlay = image.copy()
        im_out = image.copy()

        for cp in range(len(colored_polys)):
            cv2.fillPoly(overlay, [colored_polys[cp]], color)

        cv2.addWeighted(overlay, alpha, im_out, 1 - alpha, 0, im_out)

        return im_out

    def run_basic_fod(self, current_im):
        if current_im is None:
            return None
        # load ideal image
        im_in_kp = self.apply_fast(current_im)

        counts = self.keypoints_to_counts(im_keypoints = im_in_kp)

        warning_polys = self.threshold(counts, self.thresh)
        error_polys = self.threshold(counts, (self.thresh[1], sys.maxint))

        warning_overlay = self.generate_output(current_im, warning_polys, self.warningColor)
        return self.generate_output(warning_overlay, error_polys, self.errorColor)

def get_thresholds_widget(parent, values_list):
    top = tk.Toplevel(parent)
    top.title("Select your desired thresholds")
    top.update_idletasks()

    thresh1 = tk.IntVar()
    thresh2 = tk.IntVar()
    polys = tk.IntVar()

    #3 labels and 3 textbox selectors
    label1 = tk.Label(top, text = "Warning (Maybe an object)")
    label1.grid(row = 1, column = 0, sticky ="nsew")
    threshold1 = tk.Entry(top, textvariable=thresh1)
    threshold1.grid(row = 1, column = 1, sticky ="nsew")

    label2 = tk.Label(top, text = "Error (Most defintely an object)")
    label2.grid(row = 2, column = 0, sticky ="nsew")
    threshold2 = tk.Entry(top, textvariable = thresh2)
    threshold2.grid(row = 2, column = 1, sticky ="nsew")

    label3 = tk.Label(top, text = "Number of Polygons")
    label3.grid(row = 3, column = 0, sticky ="nsew")
    threshold3 = tk.Entry(top, textvariable = polys)
    threshold3.grid(row = 3, column = 1, sticky ="nsew")

    def returnVars():
        top.destroy()
        if len(values_list) != 3:
            values_list.append(thresh1.get())
            values_list.append(thresh2.get())
            values_list.append(polys.get())
        else:
            values_list[0] = thresh1.get()
            values_list[1] = thresh2.get()
            values_list[2] = polys.get()

    button = tk.Button(top, text = "Set Values", command = returnVars)
    button.grid(row = 4, column = 0, columnspan = 2, sticky ="nsew")

    threshold1.delete(0, "end")
    threshold2.delete(0, "end")
    threshold3.delete(0, "end")

    if len(values_list) == 3:
        threshold1.insert(0, values_list[0])
        threshold2.insert(0, values_list[1])
        threshold3.insert(0, values_list[2])
    else:
        threshold1.insert(0, 50)
        threshold2.insert(0, 200)
        threshold3.insert(0, 2)

    width = 300
    height = 100
    x = (top.winfo_screenwidth() // 2) - (width // 2)
    y = (top.winfo_screenheight() // 2) - (height // 2)
    top.geometry('{}x{}+{}+{}'.format(width, height, x, y))

    top.grid_columnconfigure(1, weight=1)
    top.grid_rowconfigure(0, weight=1)
    top.grid_rowconfigure(1, weight=1)
    top.grid()
    top.update()
    parent.wait_window(top)

if __name__=='__main__':
    root = tk.Tk()
    vals = []
    get_thresholds_widget(root, vals)
    print 'obtained values', vals
    get_thresholds_widget(root, vals)
    print 'obtained values', vals
    root.mainloop()


"""Truncated SVD for sparse matrices, aka latent semantic analysis (LSA).
"""

# Author: Lars Buitinck <L.J.Buitinck@uva.nl>
#         Olivier Grisel <olivier.grisel@ensta.org>
#         Michael Becker <mike@beckerfuffle.com>
# License: 3-clause BSD.

import numpy as np
import scipy.sparse as sp

try:
    from scipy.sparse.linalg import svds
except ImportError:
    from ..utils.arpack import svds

from ..base import BaseEstimator, TransformerMixin
from ..utils import check_array, as_float_array, check_random_state
from ..utils.extmath import randomized_svd, safe_sparse_dot, svd_flip
from ..utils.sparsefuncs import mean_variance_axis

__all__ = ["TruncatedSVD"]

class TruncatedSVD(BaseEstimator, TransformerMixin):
    """Dimensionality reduction using truncated SVD (aka LSA).

    This transformer performs linear dimensionality reduction by means of
    truncated singular value decomposition (SVD). It is very similar to PCA,
    but operates on sample vectors directly, instead of on a covariance matrix.
    This means it can work with scipy.sparse matrices efficiently.

    In particular, truncated SVD works on term count/tf-idf matrices as
    returned by the vectorizers in sklearn.feature_extraction.text. In that
    context, it is known as latent semantic analysis (LSA).

    This estimator supports two algorithm: a fast randomized SVD solver, and
    a "naive" algorithm that uses ARPACK as an eigensolver on (X * X.T) or
    (X.T * X), whichever is more efficient.

    Read more in the :ref:`User Guide <LSA>`.

    Parameters
    ----------
    n_components : int, default = 2
        Desired dimensionality of output data.
        Must be strictly less than the number of features.
        The default value is useful for visualisation. For LSA, a value of
        100 is recommended.

    algorithm : string, default = "randomized"
        SVD solver to use. Either "arpack" for the ARPACK wrapper in SciPy
        (scipy.sparse.linalg.svds), or "randomized" for the randomized
        algorithm due to Halko (2009).

    n_iter : int, optional
        Number of iterations for randomized SVD solver. Not used by ARPACK.

    random_state : int or RandomState, optional
        (Seed for) pseudo-random number generator. If not given, the
        numpy.random singleton is used.

    tol : float, optional
        Tolerance for ARPACK. 0 means machine precision. Ignored by randomized
        SVD solver.

    Attributes
    ----------
    components_ : array, shape (n_components, n_features)

    explained_variance_ratio_ : array, [n_components]
        Percentage of variance explained by each of the selected components.

    explained_variance_ : array, [n_components]
        The variance of the training samples transformed by a projection to
        each component.

    Examples
    --------
    >>> from sklearn.decomposition import TruncatedSVD
    >>> from sklearn.random_projection import sparse_random_matrix
    >>> X = sparse_random_matrix(100, 100, density=0.01, random_state=42)
    >>> svd = TruncatedSVD(n_components=5, random_state=42)
    >>> svd.fit(X) # doctest: +NORMALIZE_WHITESPACE
    TruncatedSVD(algorithm='randomized', n_components=5, n_iter=5,
            random_state=42, tol=0.0)
    >>> print(svd.explained_variance_ratio_) # doctest: +ELLIPSIS
    [ 0.07825... 0.05528... 0.05445... 0.04997... 0.04134...]
    >>> print(svd.explained_variance_ratio_.sum()) # doctest: +ELLIPSIS
    0.27930...

    See also
    --------
    PCA
    RandomizedPCA

    References
    ----------
    Finding structure with randomness: Stochastic algorithms for constructing
    approximate matrix decompositions
    Halko, et al., 2009 (arXiv:909) http://arxiv.org/pdf/0909.4061

    Notes
    -----
    SVD suffers from a problem called "sign indeterminancy", which means the
    sign of the ``components_`` and the output from transform depend on the
    algorithm and random state. To work around this, fit instances of this
    class to data once, then keep the instance around to do transformations.

    """
    def __init__(self, n_components=2, algorithm="randomized", n_iter=5,
                 random_state=None, tol=0.):
        self.algorithm = algorithm
        self.n_components = n_components
        self.n_iter = n_iter
        self.random_state = random_state
        self.tol = tol

    def fit(self, X, y=None):
        """Fit LSI model on training data X.

        Parameters
        ----------
        X : {array-like, sparse matrix}, shape (n_samples, n_features)
            Training data.

        Returns
        -------
        self : object
            Returns the transformer object.
        """
        self.fit_transform(X)
        return self

    def fit_transform(self, X, y=None):
        """Fit LSI model to X and perform dimensionality reduction on X.

        Parameters
        ----------
        X : {array-like, sparse matrix}, shape (n_samples, n_features)
            Training data.

        Returns
        -------
        X_new : array, shape (n_samples, n_components)
            Reduced version of X. This will always be a dense array.
        """
        X = as_float_array(X, copy=False)
        random_state = check_random_state(self.random_state)

        # If sparse and not csr or csc, convert to csr
        if sp.issparse(X) and X.getformat() not in ["csr", "csc"]:
            X = X.tocsr()

        if self.algorithm == "arpack":
            U, Sigma, VT = svds(X, k=self.n_components, tol=self.tol)
            # svds doesn't abide by scipy.linalg.svd/randomized_svd
            # conventions, so reverse its outputs.
            Sigma = Sigma[::-1]
            U, VT = svd_flip(U[:, ::-1], VT[::-1])

        elif self.algorithm == "randomized":
            k = self.n_components
            n_features = X.shape[1]
            if k >= n_features:
                raise ValueError("n_components must be < n_features;"
                                 " got %d >= %d" % (k, n_features))
            U, Sigma, VT = randomized_svd(X, self.n_components,
                                          n_iter=self.n_iter,
                                          random_state=random_state)
        else:
            raise ValueError("unknown algorithm %r" % self.algorithm)

        self.components_ = VT

        # Calculate explained variance & explained variance ratio
        X_transformed = np.dot(U, np.diag(Sigma))
        self.explained_variance_ = exp_var = np.var(X_transformed, axis=0)
        if sp.issparse(X):
            _, full_var = mean_variance_axis(X, axis=0)
            full_var = full_var.sum()
        else:
            full_var = np.var(X, axis=0).sum()
        self.explained_variance_ratio_ = exp_var / full_var
        return X_transformed

    def transform(self, X):
        """Perform dimensionality reduction on X.

        Parameters
        ----------
        X : {array-like, sparse matrix}, shape (n_samples, n_features)
            New data.

        Returns
        -------
        X_new : array, shape (n_samples, n_components)
            Reduced version of X. This will always be a dense array.
        """
        X = check_array(X, accept_sparse='csr')
        return safe_sparse_dot(X, self.components_.T)

    def inverse_transform(self, X):
        """Transform X back to its original space.

        Returns an array X_original whose transform would be X.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_components)
            New data.

        Returns
        -------
        X_original : array, shape (n_samples, n_features)
            Note that this is always a dense array.
        """
        X = check_array(X)
        return np.dot(X, self.components_)


import numpy as np
import unittest2

from instr.base import SourceMeter

class Keithley2636A(SourceMeter):
    def __init__(self, rsrc=None, timeout_sec=600, reset=True):
        self._smu = 'a'
        idn = 'Keithley Instruments Inc., Model 2636A'
        super().__init__(rsrc, idn, timeout_sec, reset)

    @property
    def smu(self):
        return self._smu

    @smu.setter
    def smu(self, value):
        if value not in ['a', 'b']:
            raise ValueError
        self._smu = value

    def check_error(self):
        if self._debug_mode:
            super().check_error()
        tmp = self.q('print(errorqueue.next())')
        if tmp != '0.00000e+00\tQueue Is Empty\t0.00000e+00\n':
            raise RuntimeError('Error on Keithley 2636A.')

    def reset(self):
        self.w('reset()', True)
        # self.w('smua.reset(); smub.reset()', True)

    def iv_sweep(self, v_start=0.0, v_end=10e-3, v_step=1e-3,
                 v_points=None, i_limit=1e-6, settle_time=0.0, reset=True):
        """
        Reference manual 3-31
        TODO: when aborted?

        :return: vis, is_aborted
        """
        if reset:
            self.reset()

        if v_points is None:
            v_points = self._v_step_to_points(v_start, v_end, v_step)

        lim = 'smu{}.source.limiti = {}'.format(self.smu, i_limit)
        self.w(lim, True)

        meas = 'SweepVLinMeasureI(smu{}, {}, {}, {}, {})'. \
            format(self.smu, v_start, v_end, settle_time, v_points)
        self.w(meas, True)

        prnt = 'printbuffer(1, {}, smu{}.nvbuffer1.readings)'. \
            format(v_points, self.smu)
        resp = self.q(prnt, True)
        Is = resp.split(', ')
        Is = np.asarray(Is, np.float64)
        if len(Is) != v_points:
            aborted = True
            v_points = len(Is)
        else:
            aborted = False
        vs = np.linspace(v_start, v_end, v_points)

        vis = np.array([vs, Is]).transpose()
        return vis, aborted

    def iv_sweep_double(self, v_max, v_step=1e-3, v_points=None,
                        i_limit=1e-3, settle_time=0.0, reset=True):
        vis1, aborted = self.iv_sweep(0, v_max, v_step,
                                      v_points, i_limit, settle_time, reset)
        if aborted:
            return vis1, aborted
        vis2, aborted = self.iv_sweep(v_max, 0, v_step,
                                      v_points, i_limit, settle_time, reset)

        ret = np.concatenate((vis1, vis2))
        return ret, aborted

class TestKeithley2636A(unittest2.TestCase):
    def test_iv_sweep(self):
        import matplotlib.pyplot as plt
        ke2636a.reset()
        v_start = 0.0
        v_end = 1e-3

        self.smu = 'a'

        vis, aborted = \
            ke2636a.iv_sweep(v_start, v_end, v_step=v_end / 10, i_limit=1e-9)
        plt.plot(*vis.transpose(), 'o-')
        plt.show()

        vis, aborted = \
            ke2636a.iv_sweep(v_start, v_end, v_points=101, i_limit=1e-6)
        plt.plot(*vis.transpose(), 'o-')
        plt.show()

        # v_step ignored
        vis, aborted = \
            ke2636a.iv_sweep(v_start, v_end, v_step=1, v_points=11)
        plt.plot(*vis.transpose(), 'o-')
        plt.show()

        vis, aborted = ke2636a.iv_sweep_double(10e-3)
        plt.plot(*vis.transpose(), 'o-')
        plt.show()

        self.smu = 'b'
        vis, aborted = ke2636a.iv_sweep(v_start, v_end, v_points=11)
        plt.plot(*vis.transpose(), 'o-')
        plt.show()

if __name__ == '__main__':
    import visa

    rm = visa.ResourceManager()

    # ke2636a_rsrc = rm.open_resource('visa://169.254.136.196/GPIB0::20::INSTR')
    ke2636a_rsrc = rm.open_resource('TCPIP::169.254.000.001::INSTR')
    ke2636a = Keithley2636A(ke2636a_rsrc)

    unittest2.main()

    pass


import numpy as np

class BaseStrategy(object):
    """
    Base class of exploration strategy.
    """

    def get_action(self, obs, policy):

        raise NotImplementedError

    def reset(self):

        pass

class OUStrategy(BaseStrategy):
    """
    Ornstein-Uhlenbeck process: dxt = theta * (mu - xt) * dt + sigma * dWt
    where Wt denotes the Wiener process.
    """

    def __init__(self, env_spec, mu=0, theta=0.15, sigma=0.3):

        self.mu = mu
        self.theta = theta
        self.sigma = sigma
        self.action_space = env_spec.action_space
        self.state = np.ones(self.action_space.flat_dim) * self.mu

    def evolve_state(self):

        x = self.state
        dx = self.theta * (self.mu - x) + self.sigma * np.random.randn(len(x))
        self.state = x + dx

        return self.state

    def reset(self):

        self.state = np.ones(self.action_space.flat_dim) * self.mu

    def get_action(self, obs, policy):

        # get_action accepts a 2D tensor with one row
    	obs = obs.reshape((1, -1))
        action = policy.get_action(obs)
        increment = self.evolve_state()

        return np.clip(action + increment,
                       self.action_space.low,
                       self.action_space.high)

if __name__ == "__main__":

    class Env1(object):

        def __init__(self):
            self.action_space = Env2()

    class Env2(object):

        def __init__(self):
            self.flat_dim = 2

    env_spec = Env1()
    test = OUStrategy(env_spec)

    states = []
    for i in range(1000):
        states.append(test.evolve_state()[0])
    import matplotlib.pyplot as plt

    plt.plot(states)
    plt.show()


#!/usr/bin/env python
"""
A comparison of multilabel target formats and metrics over them
"""
from __future__ import division
from __future__ import print_function

from timeit import timeit
from functools import partial
import itertools
import argparse
import sys

import matplotlib.pyplot as plt
import scipy.sparse as sp
import numpy as np

from sklearn.datasets import make_multilabel_classification
from sklearn.metrics import (f1_score, accuracy_score, hamming_loss,
                             jaccard_similarity_score)
from sklearn.utils.testing import ignore_warnings

METRICS = {
    'f1': partial(f1_score, average='micro'),
    'f1-by-sample': partial(f1_score, average='samples'),
    'accuracy': accuracy_score,
    'hamming': hamming_loss,
    'jaccard': jaccard_similarity_score,
}

FORMATS = {
    'sequences': lambda y: [list(np.flatnonzero(s)) for s in y],
    'dense': lambda y: y,
    'csr': lambda y: sp.csr_matrix(y),
    'csc': lambda y: sp.csc_matrix(y),
}

@ignore_warnings
def benchmark(metrics=tuple(v for k, v in sorted(METRICS.items())),
              formats=tuple(v for k, v in sorted(FORMATS.items())),
              samples=1000, classes=4, density=.2,
              n_times=5):
    """Times metric calculations for a number of inputs

    Parameters
    ----------
    metrics : array-like of callables (1d or 0d)
        The metric functions to time.

    formats : array-like of callables (1d or 0d)
        These may transform a dense indicator matrix into multilabel
        representation.

    samples : array-like of ints (1d or 0d)
        The number of samples to generate as input.

    classes : array-like of ints (1d or 0d)
        The number of classes in the input.

    density : array-like of ints (1d or 0d)
        The density of positive labels in the input.

    n_times : int
        Time calling the metric n_times times.

    Returns
    -------
    array of floats shaped like (metrics, formats, samples, classes, density)
        Time in seconds.
    """
    metrics = np.atleast_1d(metrics)
    samples = np.atleast_1d(samples)
    classes = np.atleast_1d(classes)
    density = np.atleast_1d(density)
    formats = np.atleast_1d(formats)
    out = np.zeros((len(metrics), len(formats), len(samples), len(classes),
                    len(density)), dtype=float)
    it = itertools.product(samples, classes, density)
    for i, (s, c, d) in enumerate(it):
        _, y_true = make_multilabel_classification(n_samples=s, n_features=1,
                                                   n_classes=c, n_labels=d * c,
                                                   random_state=42)
        _, y_pred = make_multilabel_classification(n_samples=s, n_features=1,
                                                   n_classes=c, n_labels=d * c,
                                                   random_state=84)
        for j, f in enumerate(formats):
            f_true = f(y_true)
            f_pred = f(y_pred)
            for k, metric in enumerate(metrics):
                t = timeit(partial(metric, f_true, f_pred), number=n_times)

                out[k, j].flat[i] = t
    return out

def _tabulate(results, metrics, formats):
    """Prints results by metric and format

    Uses the last ([-1]) value of other fields
    """
    column_width = max(max(len(k) for k in formats) + 1, 8)
    first_width = max(len(k) for k in metrics)
    head_fmt = ('{:<{fw}s}' + '{:>{cw}s}' * len(formats))
    row_fmt = ('{:<{fw}s}' + '{:>{cw}.3f}' * len(formats))
    print(head_fmt.format('Metric', *formats,
                          cw=column_width, fw=first_width))
    for metric, row in zip(metrics, results[:, :, -1, -1, -1]):
        print(row_fmt.format(metric, *row,
                             cw=column_width, fw=first_width))

def _plot(results, metrics, formats, title, x_ticks, x_label,
          format_markers=('x', '|', 'o', '+'),
          metric_colors=('c', 'm', 'y', 'k', 'g', 'r', 'b')):
    """
    Plot the results by metric, format and some other variable given by
    x_label
    """
    fig = plt.figure('scikit-learn multilabel metrics benchmarks')
    plt.title(title)
    ax = fig.add_subplot(111)
    for i, metric in enumerate(metrics):
        for j, format in enumerate(formats):
            ax.plot(x_ticks, results[i, j].flat,
                    label='{}, {}'.format(metric, format),
                    marker=format_markers[j],
                    color=metric_colors[i % len(metric_colors)])
    ax.set_xlabel(x_label)
    ax.set_ylabel('Time (s)')
    ax.legend()
    plt.show()

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument('metrics', nargs='*', default=sorted(METRICS),
                    help='Specifies metrics to benchmark, defaults to all. '
                         'Choices are: {}'.format(sorted(METRICS)))
    ap.add_argument('--formats', nargs='+', choices=sorted(FORMATS),
                    help='Specifies multilabel formats to benchmark '
                         '(defaults to all).')
    ap.add_argument('--samples', type=int, default=1000,
                    help='The number of samples to generate')
    ap.add_argument('--classes', type=int, default=10,
                    help='The number of classes')
    ap.add_argument('--density', type=float, default=.2,
                    help='The average density of labels per sample')
    ap.add_argument('--plot', choices=['classes', 'density', 'samples'],
                    default=None,
                    help='Plot time with respect to this parameter varying '
                         'up to the specified value')
    ap.add_argument('--n-steps', default=10, type=int,
                    help='Plot this many points for each metric')
    ap.add_argument('--n-times',
                    default=5, type=int,
                    help="Time performance over n_times trials")
    args = ap.parse_args()

    if args.plot is not None:
        max_val = getattr(args, args.plot)
        if args.plot in ('classes', 'samples'):
            min_val = 2
        else:
            min_val = 0
        steps = np.linspace(min_val, max_val, num=args.n_steps + 1)[1:]
        if args.plot in ('classes', 'samples'):
            steps = np.unique(np.round(steps).astype(int))
        setattr(args, args.plot, steps)

    if args.metrics is None:
        args.metrics = sorted(METRICS)
    if args.formats is None:
        args.formats = sorted(FORMATS)

    results = benchmark([METRICS[k] for k in args.metrics],
                        [FORMATS[k] for k in args.formats],
                        args.samples, args.classes, args.density,
                        args.n_times)

    _tabulate(results, args.metrics, args.formats)

    if args.plot is not None:
        print('Displaying plot', file=sys.stderr)
        title = ('Multilabel metrics with %s' %
                 ', '.join('{0}={1}'.format(field, getattr(args, field))
                           for field in ['samples', 'classes', 'density']
                           if args.plot != field))
        _plot(results, args.metrics, args.formats, title, steps, args.plot)


"""
Table Schema builders

http://specs.frictionlessdata.io/json-table-schema/
"""
from pandas.core.dtypes.common import (
    is_integer_dtype, is_timedelta64_dtype, is_numeric_dtype,
    is_bool_dtype, is_datetime64_dtype, is_datetime64tz_dtype,
    is_categorical_dtype, is_period_dtype, is_string_dtype
)

def as_json_table_type(x):
    """
    Convert a NumPy / pandas type to its corresponding json_table.

    Parameters
    ----------
    x : array or dtype

    Returns
    -------
    t : str
        the Table Schema data types

    Notes
    -----
    This table shows the relationship between NumPy / pandas dtypes,
    and Table Schema dtypes.

    ==============  =================
    Pandas type     Table Schema type
    ==============  =================
    int64           integer
    float64         number
    bool            boolean
    datetime64[ns]  datetime
    timedelta64[ns] duration
    object          str
    categorical     any
    =============== =================
    """
    if is_integer_dtype(x):
        return 'integer'
    elif is_bool_dtype(x):
        return 'boolean'
    elif is_numeric_dtype(x):
        return 'number'
    elif (is_datetime64_dtype(x) or is_datetime64tz_dtype(x) or
          is_period_dtype(x)):
        return 'datetime'
    elif is_timedelta64_dtype(x):
        return 'duration'
    elif is_categorical_dtype(x):
        return 'any'
    elif is_string_dtype(x):
        return 'string'
    else:
        return 'any'

def set_default_names(data):
    """Sets index names to 'index' for regular, or 'level_x' for Multi"""
    if all(name is not None for name in data.index.names):
        return data

    data = data.copy()
    if data.index.nlevels > 1:
        names = [name if name is not None else 'level_{}'.format(i)
                 for i, name in enumerate(data.index.names)]
        data.index.names = names
    else:
        data.index.name = data.index.name or 'index'
    return data

def make_field(arr, dtype=None):
    dtype = dtype or arr.dtype
    if arr.name is None:
        name = 'values'
    else:
        name = arr.name
    field = {'name': name,
             'type': as_json_table_type(dtype)}

    if is_categorical_dtype(arr):
        if hasattr(arr, 'categories'):
            cats = arr.categories
            ordered = arr.ordered
        else:
            cats = arr.cat.categories
            ordered = arr.cat.ordered
        field['constraints'] = {"enum": list(cats)}
        field['ordered'] = ordered
    elif is_period_dtype(arr):
        field['freq'] = arr.freqstr
    elif is_datetime64tz_dtype(arr):
        if hasattr(arr, 'dt'):
            field['tz'] = arr.dt.tz.zone
        else:
            field['tz'] = arr.tz.zone
    return field

def build_table_schema(data, index=True, primary_key=None, version=True):
    """
    Create a Table schema from ``data``.

    Parameters
    ----------
    data : Series, DataFrame
    index : bool, default True
        Whether to include ``data.index`` in the schema.
    primary_key : bool or None, default True
        column names to designate as the primary key.
        The default `None` will set `'primaryKey'` to the index
        level or levels if the index is unique.
    version : bool, default True
        Whether to include a field `pandas_version` with the version
        of pandas that generated the schema.

    Returns
    -------
    schema : dict

    Examples
    --------
    >>> df = pd.DataFrame(
    ...     {'A': [1, 2, 3],
    ...      'B': ['a', 'b', 'c'],
    ...      'C': pd.date_range('2016-01-01', freq='d', periods=3),
    ...     }, index=pd.Index(range(3), name='idx'))
    >>> build_table_schema(df)
    {'fields': [{'name': 'idx', 'type': 'integer'},
    {'name': 'A', 'type': 'integer'},
    {'name': 'B', 'type': 'string'},
    {'name': 'C', 'type': 'datetime'}],
    'pandas_version': '0.20.0',
    'primaryKey': ['idx']}

    Notes
    -----
    See `_as_json_table_type` for conversion types.
    Timedeltas as converted to ISO8601 duration format with
    9 decimal places after the secnods field for nanosecond precision.

    Categoricals are converted to the `any` dtype, and use the `enum` field
    constraint to list the allowed values. The `ordered` attribute is included
    in an `ordered` field.
    """
    if index is True:
        data = set_default_names(data)

    schema = {}
    fields = []

    if index:
        if data.index.nlevels > 1:
            for level in data.index.levels:
                fields.append(make_field(level))
        else:
            fields.append(make_field(data.index))

    if data.ndim > 1:
        for column, s in data.iteritems():
            fields.append(make_field(s))
    else:
        fields.append(make_field(data))

    schema['fields'] = fields
    if index and data.index.is_unique and primary_key is None:
        if data.index.nlevels == 1:
            schema['primaryKey'] = [data.index.name]
        else:
            schema['primaryKey'] = data.index.names
    elif primary_key is not None:
        schema['primaryKey'] = primary_key

    if version:
        schema['pandas_version'] = '0.20.0'
    return schema