"""Tests for Bundle.

"""

import numpy as np
import pytest

import datreant as dtr


def do_stuff(treant):
    return treant.name


def return_nothing(treant):
    b = treant.name


class CollectionsTests(object):
    """Mixin tests for collections"""

    def test_parents(self, collection, tmpdir):
        with tmpdir.as_cwd():
            t1 = dtr.Treant('free-associate/lark')
            t2 = dtr.Treant('free-associate/hark')
            t3 = dtr.Treant('characters/linus')

        col = collection(t1, t2, t3)

        assert len(col.parents()) == 2
        assert 'free-associate' in col.parents().names
        assert 'linus' not in col.parents().names
        assert 'characters' in col.parents().names

    class TestGetitem(object):
        @pytest.mark.parametrize('slx', (
            [1, 2],
            np.array([1, 2]),
        ))
        def test_fancy_index(self, filled_collection, slx):
            b, (t1, t2, t3) = filled_collection
            sl = b[slx]
            assert len(sl) == 2
            assert t2 == sl[0]
            assert t3 == sl[1]

        @pytest.mark.parametrize('slx', (
            [False, False, True],
            np.array([False, False, True]),
        ))
        def test_boolean_index(self, filled_collection, slx):
            b, (t1, t2, t3) = filled_collection
            sl = b[slx]
            assert len(sl) == 1
            assert t3 == sl[0]

        @pytest.mark.parametrize('slx', (
            slice(0, 1, None),
            slice(1, None, None),
            slice(None, None, -1),
            slice(None, None, 2),
        ))
        def test_slice_index(self, filled_collection, slx):
            b, ts = filled_collection
            sl = b[slx]
            ref = ts[slx]
            for x, y in zip(sl, ref):
                assert x == y

        def test_getitem_IE(self, filled_collection):
            bundle = filled_collection[0]
            with pytest.raises(IndexError):
                bundle[4.0]

    class TestSetOperations(object):
        def test_sub_single(self, filled_collection):
            b, (t1, t2, t3) = filled_collection
            b1 = b[[0, 1, 2]]
            b2 = b[1]
            b3 = b1 - b2
            assert len(b3) == 2
            assert t1 in b3
            assert t2 not in b3
            assert t3 in b3

        def test_sub_many(self, filled_collection):
            b, (t1, t2, t3) = filled_collection
            b1 = b[[0, 1]]
            b2 = b[[1, 2]]
            b3 = b1 - b2
            assert len(b3) == 1
            assert t1 in b3
            assert t2 not in b3

        def test_or(self, filled_collection):
            b, (t1, t2, t3) = filled_collection
            b1 = b[[0, 1]]
            b2 = b[[1, 2]]
            b3 = b1 | b2
            assert t1 in b3
            assert t2 in b3
            assert t3 in b3

        def test_and(self, filled_collection):
            b, (t1, t2, t3) = filled_collection
            b1 = b[[0, 1]]
            b2 = b[[1, 2]]
            b3 = b1 & b2
            assert t1 not in b3

        def test_xor(self, filled_collection):
            b, (t1, t2, t3) = filled_collection
            b1 = b[[0, 1]]
            b2 = b[[1, 2]]
            b3 = b1 ^ b2
            assert len(b3) == 2
            assert t1 in b3
            assert t2 not in b3
            assert t3 in b3

        def test_sub_TypeError(self, filled_collection):
            b = filled_collection[0]
            with pytest.raises(TypeError):
                b - ['this']

        def test_or_TypeError(self, filled_collection):
            b = filled_collection[0]
            with pytest.raises(TypeError):
                b | ['this']

        def test_and_TypeError(self, filled_collection):
            b = filled_collection[0]
            with pytest.raises(TypeError):
                b & ['this']

        def test_xor_TypeError(self, filled_collection):
            b = filled_collection[0]
            with pytest.raises(TypeError):
                b ^ ['this']

    class TestAddition(object):
        def test_add_many(self, filled_collection):
            b, (t1, t2, t3) = filled_collection
            b1 = b[[0, 1]]
            b2 = b[[1, 2]]
            b3 = b1 + b2
            assert len(b3) == 3
            assert t1 in b3
            assert t2 in b3
            assert t3 in b3

        def test_add_singular(self, filled_collection):
            b, (t1, t2, t3) = filled_collection
            b1 = b[[0, 1]]
            b2 = b[2]
            b3 = b1 + b2
            assert len(b3) == 3
            assert t1 in b3
            assert t2 in b3
            assert t3 in b3

        def test_add(self, filled_collection):
            b = filled_collection[0]
            with pytest.raises(TypeError):
                b + 25


class TestView(CollectionsTests):
    """Tests for Views"""

    @pytest.fixture
    def collection(self):
        return dtr.View

    @pytest.fixture
    def filled_collection(self, tmpdir):
        # returns (a bundle of [t1, t2, t3], then individal references to each)
        with tmpdir.as_cwd():
            t1 = dtr.Tree('larry')
            t2 = dtr.Leaf('curly')
            t3 = dtr.Treant('moe')
            b = dtr.View(t1, t2, t3)
            return b, (t1, t2, t3)

    class TestGetitem(CollectionsTests.TestGetitem):
        def test_getitem_name_string(self, filled_collection):
            b, (t1, t2, t3) = filled_collection
            n = t1.name

            b_new = b[n]
            assert isinstance(b_new, dtr.View)
            assert b_new[0] == t1

    class TestAddition(CollectionsTests.TestAddition):
        def test_tree_addition(self, filled_collection):
            b, (t1, t2, t3) = filled_collection
            b1 = b[[1, 2]]
            b3 = b1 + t1
            assert len(b3) == 3
            assert isinstance(b3, dtr.View)
            assert t1 in b3

        def test_leaf_addition(self, filled_collection):
            b, (t1, t2, t3) = filled_collection
            b1 = b[[0, 2]]
            b3 = b1 + t2
            assert len(b3) == 3
            assert isinstance(b3, dtr.View)
            assert t2 in b3

        def test_treant_addition(self, filled_collection):
            b, (t1, t2, t3) = filled_collection
            b1 = b[[0, 1]]
            b3 = b1 + t3
            assert len(b3) == 3
            assert isinstance(b3, dtr.View)
            assert t2 in b3

    def test_exists(self, collection, tmpdir):
        pass


class TestBundle(CollectionsTests):
    """Tests for elements of Bundle"""

    @pytest.fixture
    def collection(self):
        return dtr.Bundle

    @pytest.fixture
    def filled_collection(self, tmpdir):
        # returns (a bundle of [t1, t2, t3], then individal references to each)
        with tmpdir.as_cwd():
            t1 = dtr.Treant('larry')
            t2 = dtr.Treant('curly')
            t3 = dtr.Treant('moe')
            b = dtr.Bundle((t1, t2, t3))
            return b, (t1, t2, t3)

    @pytest.fixture
    def testtreant(self, tmpdir, request):
        with tmpdir.as_cwd():
            t = dtr.Treant('dummytreant')
        return t

    @pytest.fixture
    def testtreant2(self, tmpdir, request):
        with tmpdir.as_cwd():
            t = dtr.Treant('dummytreant2')
        return t

    def test_additive(self, tmpdir, testtreant, testtreant2, collection):
        """Test that addition of treants and collections give Bundles.

        """
        with tmpdir.as_cwd():
            assert isinstance(testtreant + testtreant2, dtr.Bundle)
            assert len(testtreant + testtreant2) == 2

            b = collection() + testtreant + testtreant2

            # beating a dead horse
            assert len(b) == 2

    class TestGetitem(CollectionsTests.TestGetitem):
        def test_getitem_name_string(self, filled_collection):
            b, (t1, t2, t3) = filled_collection
            n = t1.name

            b_new = b[n]
            assert isinstance(b_new, dtr.Bundle)
            assert b_new[0] == t1

        def test_getitem_string_KeyError(self, filled_collection):
            b = filled_collection[0]
            with pytest.raises(KeyError):
                b['not there']

    def test_get_members(self, collection, tmpdir):
        """Access members with indexing and slicing"""
        with tmpdir.as_cwd():
            t1 = dtr.Treant('larry')
            t2 = dtr.Treant('curly')
            t3 = dtr.Treant('moe')

            col = collection([[[t1, [t2, [t3]]]]])

            assert col[1] == t2

            t4 = dtr.treants.Treant('shemp')
            col = col + t4

            for member in (t1, t2, t3):
                assert member in col[:3]

            assert t4 not in col[:3]
            assert t4 == col[-1]

    def test_member_attributes(self, collection, tmpdir):
        """Get member names and abspaths"""
        with tmpdir.as_cwd():
            t1 = dtr.Treant('bigger')
            t2 = dtr.Treant('faster')
            t3 = dtr.Treant('stronger')

        col = collection(t1, t2, t3)

        names = [treant.name for treant in [t1, t2, t3]]
        assert col.names == names

        abspaths = [treant.abspath for treant in [t1, t2, t3]]
        assert col.abspaths == abspaths

    def test_map(self, collection, tmpdir):
        with tmpdir.as_cwd():
            t1 = dtr.Treant('lark')
            t2 = dtr.Treant('hark')
            t3 = dtr.Treant('linus')

        col = collection(t1, t2, t3)

        comp = [cont.name for cont in col]
        assert col.map(do_stuff) == comp
        assert col.map(do_stuff, processes=2) == comp

        assert col.map(return_nothing) is None
        assert col.map(return_nothing, processes=2) is None

    class TestGet:
        @pytest.fixture
        def filled_bundle(self, tmpdir):
            with tmpdir.as_cwd():
                t1 = dtr.Treant('one')
                t2 = dtr.Treant('two')
                t3 = dtr.Treant('three')
                t4 = dtr.Treant('four')

                t1.tags.add('odd', 'one')
                t2.tags.add('even', 'two')
                t3.tags.add('odd', 'three')
                t4.tags.add('even', 'four')
                t1.categories.add(is_even=False, value=1)
                t2.categories.add(is_even=True, value=2)
                t3.categories.add(is_even=False, value=3)
                t4.categories.add(is_even=True, value=4)

                return [t1, t2, t3, t4], dtr.Bundle([t1, t2, t3, t4])

        def test_get_blank(self, filled_bundle):
            ref, b = filled_bundle

            new = b.get()

            assert new == b

        def test_get_cats(self, filled_bundle):
            ref, b = filled_bundle

            new = b.get(is_even=True)

            assert len(new) == 2
            assert ref[1] in new
            assert ref[3] in new

        def test_get_cats_double(self, filled_bundle):
            ref, b = filled_bundle

            new = b.get(is_even=True, value=4)

            assert len(new) == 1
            assert ref[3] in new

        def test_get_cats_empty(self, filled_bundle):
            ref, b = filled_bundle

            new = b.get(is_even=True, value=1)

            assert len(new) == 0

        def test_get_cats_KeyError(self, filled_bundle):
            ref, b = filled_bundle

            new = b.get(colour='yellow')

            assert len(new) == 0

        def test_get_tags(self, filled_bundle):
            ref, b = filled_bundle

            new = b.get('odd')

            assert len(new) == 2
            assert ref[0] in new
            assert ref[2] in new

        def test_get_tags_many(self, filled_bundle):
            ref, b = filled_bundle

            new = b.get('odd', 'three')

            assert len(new) == 1
            assert ref[2] in new

        def test_get_tags_empty(self, filled_bundle):
            ref, b = filled_bundle

            new = b.get('magical')

            assert len(new) == 0

        def test_get_cat_and_tag(self, filled_bundle):
            ref, b = filled_bundle

            new = b.get('odd', value=3)

            assert len(new) == 1
            assert ref[2] in new

    class TestAggTags:
        """Test behavior of manipulating tags collectively.

        """
        def test_add_tags(self, collection, testtreant, testtreant2, tmpdir):
            with tmpdir.as_cwd():
                col = collection(testtreant, testtreant2)

                assert len(col.tags) == 0

                col.tags.add('broiled', 'not baked')

                assert len(col.tags) == 2
                for tag in ('broiled', 'not baked'):
                    assert tag in col.tags

        def test_tags_setting(self, collection, testtreant,
                              testtreant2, tmpdir):
            with tmpdir.as_cwd():
                col = collection(testtreant, testtreant2)

                assert len(col.tags) == 0

                # add as list
                col.tags = ['broiled', 'not baked']

                assert len(col.tags) == 2
                for tag in ('broiled', 'not baked'):
                    assert tag in col.tags

                col.tags.clear()

                # add as set
                col.tags = {'broiled', 'not baked'}

                assert len(col.tags) == 2
                for tag in ('broiled', 'not baked'):
                    assert tag in col.tags

                col.tags.clear()

                # add as Tags
                t = dtr.Treant('mark twain')
                t.tags.add('literature', 'quotables')
                col.tags = t.tags

                assert len(col.tags) == 2
                for tag in ('literature', 'quotables'):
                    assert tag in col.tags

        def test_tags_all(self, collection, tmpdir):
            with tmpdir.as_cwd():

                moe = dtr.Treant('moe',
                                 tags=['smartest', 'mean', 'stooge'])
                larry = dtr.Treant('larry',
                                   tags=['weird', 'means well', 'stooge'])
                curly = dtr.Treant('curly',
                                   tags=['dumb', 'nyuk-nyuk', 'stooge'])

                col = collection(moe, larry, curly)

                assert len(col.tags.all) == 1
                assert 'stooge' in col.tags.all

        def test_tags_any(self, collection, testtreant, testtreant2, tmpdir):
            with tmpdir.as_cwd():

                moe = dtr.Treant('moe',
                                 tags=['smartest', 'mean', 'stooge'])
                larry = dtr.Treant('larry',
                                   tags=['weird', 'means well', 'stooge'])
                curly = dtr.Treant('curly',
                                   tags=['dumb', 'nyuk-nyuk', 'stooge'])

                col = collection(moe, larry, curly)

                assert len(col.tags.any) == 7
                for tag in ('smartest', 'mean', 'weird', 'means well',
                            'dumb', 'nyuk-nyuk', 'stooge'):
                    assert tag in col.tags.any

        def test_tags_set_behavior(self, collection, tmpdir):
            with tmpdir.as_cwd():

                t1 = dtr.Treant('maple')
                t2 = dtr.Treant('pine')
                t3 = dtr.Treant('juniper')
                t1.tags.add({'tree', 'new jersey', 'deciduous'})
                t2.tags.add({'tree', 'new york', 'evergreen'})
                t3.tags.add({'shrub', 'new york', 'evergreen'})
                col = collection(t1, t2, t3)
                trees = dtr.Bundle('maple', 'pine')
                evergreens = dtr.Bundle('pine', 'juniper')
                tags = col.tags

                assert len(tags.any) == 6

                # test equality: __eq__ (==)
                assert t1.tags == {'tree', 'new jersey', 'deciduous'}
                assert t2.tags == {'tree', 'new york', 'evergreen'}
                assert t3.tags == {'shrub', 'new york', 'evergreen'}

                # test subset: __lt__ (<)
                assert not t1.tags < {'tree', 'new jersey', 'deciduous'}
                assert tags < {'tree', 'new jersey', 'deciduous'}
                assert t1.tags < tags.any
                assert t2.tags < tags.any
                assert t3.tags < tags.any

                # test difference: __sub__ (-)
                assert t1.tags - {'tree'} == {'new jersey', 'deciduous'}
                assert trees.tags - {'tree'} == set()

                # test right difference: __rsub__ (-)
                evergreen_ny_shrub = {'evergreen', 'new york', 'shrub'}
                dec_nj_sh = {'deciduous', 'new jersey', 'shrub'}
                assert tags.any - t1.tags == evergreen_ny_shrub
                assert tags.any - evergreens.tags - trees.tags == dec_nj_sh
                assert {'tree'} - trees.tags == set()

                # test union: __or__ (|)
                evergreen_ny_shrub = {'evergreen', 'new york', 'shrub'}
                assert evergreens.tags | t3.tags == evergreen_ny_shrub
                assert t1.tags | t2.tags | t3.tags == tags.any

                # test right union: __ror__ (|)
                assert {'shrub'} | evergreens.tags == evergreen_ny_shrub
                assert t3.tags | {'tree'} == {'tree'} | t3.tags

                # test intersection: __and__ (&)
                evergreen_ny = {'evergreen', 'new york'}
                assert evergreens.tags & t3.tags == evergreen_ny
                assert t1.tags & t2.tags & t3.tags == tags.all

                # test right intersection: __rand__ (&)
                assert evergreen_ny_shrub & evergreens.tags == evergreen_ny
                assert t3.tags & {'shrub'} == {'shrub'} & t3.tags

                # test symmetric difference: __xor__ (^)
                evergreen_ny_tree = {'evergreen', 'new york', 'tree'}
                assert trees.tags ^ evergreens.tags == evergreen_ny_tree
                assert evergreens.tags ^ t3.tags == {'shrub'}
                assert t1.tags ^ t2.tags ^ t3.tags == dec_nj_sh

                # test right symmetric difference: __rxor__ (^)
                assert {'new york'} ^ evergreens.tags == {'evergreen'}
                assert {'shrub'} ^ trees.tags == t2.tags ^ t3.tags

                # type_error_msg = "Operands must be AggTags, Tags, or a set."
                with pytest.raises(TypeError) as e:
                    ['tree'] == trees.tags
                # assert e.value.message == type_error_msg
                with pytest.raises(TypeError) as e:
                    trees.tags < ['tree']
                # assert e.value.message == type_error_msg
                with pytest.raises(TypeError) as e:
                    ['tree'] - trees.tags
                # assert e.value.message == type_error_msg
                with pytest.raises(TypeError) as e:
                    trees.tags - ['tree']
                # assert e.value.message == type_error_msg
                with pytest.raises(TypeError) as e:
                    ['tree'] | trees.tags
                # assert e.value.message == type_error_msg
                with pytest.raises(TypeError) as e:
                    trees.tags | ['tree']
                # assert e.value.message == type_error_msg
                with pytest.raises(TypeError) as e:
                    ['tree'] & trees.tags
                # assert e.value.message == type_error_msg
                with pytest.raises(TypeError) as e:
                    trees.tags & ['tree']
                # assert e.value.message == type_error_msg
                with pytest.raises(TypeError) as e:
                    ['tree'] ^ trees.tags
                # assert e.value.message == type_error_msg
                with pytest.raises(TypeError) as e:
                    trees.tags ^ ['tree']
                # assert e.value.message == type_error_msg

        def test_tags_getitem(self, collection, testtreant,
                              testtreant2, tmpdir):
            with tmpdir.as_cwd():

                t1 = dtr.Treant('maple')
                t2 = dtr.Treant('pine')
                t1.tags.add({'tree', 'new jersey', 'deciduous'})
                t2.tags.add({'tree', 'new york', 'evergreen'})
                col = collection(t1, t2)
                tags = col.tags

                assert len(tags.any) == 5
                # test single tags
                assert tags['tree'] == [True, True]
                assert tags['deciduous'] == [True, False]
                assert tags['evergreen'] == [False, True]
                assert tags['new jersey'] == [True, False]
                assert tags['new york'] == [False, True]
                assert tags[{'tree'}] == [False, False]
                assert tags[{'deciduous'}] == [False, True]
                assert tags[{'evergreen'}] == [True, False]
                assert tags[{'new jersey'}] == [False, True]
                assert tags[{'new york'}] == [True, False]

                # test for Treants with ALL the given tags
                assert tags[['tree', 'deciduous']] == [True, False]
                assert tags[['tree', 'evergreen']] == [False, True]
                assert tags[['new jersey', 'evergreen']] == [False, False]

                # test for Treants with ANY of the given tags
                assert tags[('tree', 'deciduous')] == [True, True]
                assert tags[('deciduous', 'evergreen')] == [True, True]
                assert tags[('new york', 'evergreen')] == [False, True]

                # test for Treants without at least one of the given tags
                assert tags[{'deciduous', 'evergreen'}] == [True, True]
                assert tags[{'tree', 'deciduous'}] == [False, True]
                assert tags[{'tree', 'new york', 'evergreen'}] == [True, False]

                # complex logic tests

                # give those that are evergreen or in NY AND also not deciduous
                selection = [('evergreen', 'new york'), {'deciduous'}]
                assert tags[selection] == [False, True]
                # give those that are evergreen or in NY AND also not a tree
                selection = [('evergreen', 'new york'), {'tree'}]
                assert tags[selection] == [False, False]
                # give a tree that's in NJ OR anything that's not evergreen
                selection = (['tree', 'new jersey'], {'evergreen'})
                assert tags[selection] == [True, False]
                # cannot be a tree in NJ, AND must also be deciduous
                # I.e., give all deciduous things that aren't trees in NJ
                selection = [{'tree', 'new jersey'}, 'deciduous']
                assert tags[selection] == [False, False]

        def test_tags_fuzzy(self, collection, testtreant, testtreant2, tmpdir):
            with tmpdir.as_cwd():

                t1 = dtr.Treant('maple')
                t2 = dtr.Treant('pine')
                t1.tags.add({'tree', 'new jersey', 'deciduous'})
                t2.tags.add({'tree', 'new york', 'evergreen'})
                col = collection(t1, t2)
                tags = col.tags

                assert len(tags.any) == 5

                all_tree1 = tags.fuzzy('tree', threshold=80, scope='all')
                all_tree2 = tags.fuzzy('tree')
                assert all_tree1 == all_tree2
                assert all_tree2 == ('tree',)

                any_deciduous = tags.fuzzy('deciduous', scope='any')
                assert any_deciduous == ('deciduous',)
                all_evergreen = tags.fuzzy('evergreen')
                assert all_evergreen == ()

                # check that fuzzy matching is independent of threshold when
                # exact tag is present in all members
                all_tree_strict = tags.fuzzy('tree', threshold=99)
                assert all_tree_strict == ('tree',)
                all_tree_tolerant = tags.fuzzy('tree', threshold=0)
                assert all_tree_tolerant == ('tree',)

                # check that fuzzy matching will give differing tags when
                # members have similar tag names ('new') and the threshold is
                # varied
                all_ny = tags.fuzzy('new york')
                assert all_ny == ()
                any_ny_strict = tags.fuzzy('new york', scope='any')
                assert any_ny_strict == ('new york',)
                any_ny_tol = tags.fuzzy('new york', threshold=50, scope='any')
                assert set(any_ny_tol) == {'new york', 'new jersey'}

                # check fuzzy matching for multiple tags (scope='all')
                new_ever = ['new', 'evergreen']