import numpy as np
from pandas.util import testing as tm
from pandas import Panel, date_range, DataFrame


class TestPanel(tm.TestCase):

    def test_iloc_getitem_panel(self):

        # GH 7189
        p = Panel(np.arange(4 * 3 * 2).reshape(4, 3, 2),
                  items=['A', 'B', 'C', 'D'],
                  major_axis=['a', 'b', 'c'],
                  minor_axis=['one', 'two'])

        result = p.iloc[1]
        expected = p.loc['B']
        tm.assert_frame_equal(result, expected)

        result = p.iloc[1, 1]
        expected = p.loc['B', 'b']
        tm.assert_series_equal(result, expected)

        result = p.iloc[1, 1, 1]
        expected = p.loc['B', 'b', 'two']
        self.assertEqual(result, expected)

        # slice
        result = p.iloc[1:3]
        expected = p.loc[['B', 'C']]
        tm.assert_panel_equal(result, expected)

        result = p.iloc[:, 0:2]
        expected = p.loc[:, ['a', 'b']]
        tm.assert_panel_equal(result, expected)

        # list of integers
        result = p.iloc[[0, 2]]
        expected = p.loc[['A', 'C']]
        tm.assert_panel_equal(result, expected)

        # neg indicies
        result = p.iloc[[-1, 1], [-1, 1]]
        expected = p.loc[['D', 'B'], ['c', 'b']]
        tm.assert_panel_equal(result, expected)

        # dups indicies
        result = p.iloc[[-1, -1, 1], [-1, 1]]
        expected = p.loc[['D', 'D', 'B'], ['c', 'b']]
        tm.assert_panel_equal(result, expected)

        # combined
        result = p.iloc[0, [True, True], [0, 1]]
        expected = p.loc['A', ['a', 'b'], ['one', 'two']]
        tm.assert_frame_equal(result, expected)

        # out-of-bounds exception
        self.assertRaises(IndexError, p.iloc.__getitem__, tuple([10, 5]))

        def f():
            p.iloc[0, [True, True], [0, 1, 2]]

        self.assertRaises(IndexError, f)

        # trying to use a label
        self.assertRaises(ValueError, p.iloc.__getitem__, tuple(['j', 'D']))

        # GH
        p = Panel(
            np.random.rand(4, 3, 2), items=['A', 'B', 'C', 'D'],
            major_axis=['U', 'V', 'W'], minor_axis=['X', 'Y'])
        expected = p['A']

        result = p.iloc[0, :, :]
        tm.assert_frame_equal(result, expected)

        result = p.iloc[0, [True, True, True], :]
        tm.assert_frame_equal(result, expected)

        result = p.iloc[0, [True, True, True], [0, 1]]
        tm.assert_frame_equal(result, expected)

        def f():
            p.iloc[0, [True, True, True], [0, 1, 2]]

        self.assertRaises(IndexError, f)

        def f():
            p.iloc[0, [True, True, True], [2]]

        self.assertRaises(IndexError, f)

    def test_iloc_panel_issue(self):

        # GH 3617
        p = Panel(np.random.randn(4, 4, 4))

        self.assertEqual(p.iloc[:3, :3, :3].shape, (3, 3, 3))
        self.assertEqual(p.iloc[1, :3, :3].shape, (3, 3))
        self.assertEqual(p.iloc[:3, 1, :3].shape, (3, 3))
        self.assertEqual(p.iloc[:3, :3, 1].shape, (3, 3))
        self.assertEqual(p.iloc[1, 1, :3].shape, (3, ))
        self.assertEqual(p.iloc[1, :3, 1].shape, (3, ))
        self.assertEqual(p.iloc[:3, 1, 1].shape, (3, ))

    def test_panel_getitem(self):
        # GH4016, date selection returns a frame when a partial string
        # selection
        ind = date_range(start="2000", freq="D", periods=1000)
        df = DataFrame(
            np.random.randn(
                len(ind), 5), index=ind, columns=list('ABCDE'))
        panel = Panel(dict([('frame_' + c, df) for c in list('ABC')]))

        test2 = panel.ix[:, "2002":"2002-12-31"]
        test1 = panel.ix[:, "2002"]
        tm.assert_panel_equal(test1, test2)

        # GH8710
        # multi-element getting with a list
        panel = tm.makePanel()

        expected = panel.iloc[[0, 1]]

        result = panel.loc[['ItemA', 'ItemB']]
        tm.assert_panel_equal(result, expected)

        result = panel.loc[['ItemA', 'ItemB'], :, :]
        tm.assert_panel_equal(result, expected)

        result = panel[['ItemA', 'ItemB']]
        tm.assert_panel_equal(result, expected)

        result = panel.loc['ItemA':'ItemB']
        tm.assert_panel_equal(result, expected)

        result = panel.ix['ItemA':'ItemB']
        tm.assert_panel_equal(result, expected)

        result = panel.ix[['ItemA', 'ItemB']]
        tm.assert_panel_equal(result, expected)

        # with an object-like
        # GH 9140
        class TestObject:

            def __str__(self):
                return "TestObject"

        obj = TestObject()

        p = Panel(np.random.randn(1, 5, 4), items=[obj],
                  major_axis=date_range('1/1/2000', periods=5),
                  minor_axis=['A', 'B', 'C', 'D'])

        expected = p.iloc[0]
        result = p[obj]
        tm.assert_frame_equal(result, expected)

    def test_panel_setitem(self):

        # GH 7763
        # loc and setitem have setting differences
        np.random.seed(0)
        index = range(3)
        columns = list('abc')

        panel = Panel({'A': DataFrame(np.random.randn(3, 3),
                                      index=index, columns=columns),
                       'B': DataFrame(np.random.randn(3, 3),
                                      index=index, columns=columns),
                       'C': DataFrame(np.random.randn(3, 3),
                                      index=index, columns=columns)})

        replace = DataFrame(np.eye(3, 3), index=range(3), columns=columns)
        expected = Panel({'A': replace, 'B': replace, 'C': replace})

        p = panel.copy()
        for idx in list('ABC'):
            p[idx] = replace
        tm.assert_panel_equal(p, expected)

        p = panel.copy()
        for idx in list('ABC'):
            p.loc[idx, :, :] = replace
        tm.assert_panel_equal(p, expected)

    def test_panel_assignment(self):
        # GH3777
        wp = Panel(np.random.randn(2, 5, 4), items=['Item1', 'Item2'],
                   major_axis=date_range('1/1/2000', periods=5),
                   minor_axis=['A', 'B', 'C', 'D'])
        wp2 = Panel(np.random.randn(2, 5, 4), items=['Item1', 'Item2'],
                    major_axis=date_range('1/1/2000', periods=5),
                    minor_axis=['A', 'B', 'C', 'D'])

        # TODO: unused?
        # expected = wp.loc[['Item1', 'Item2'], :, ['A', 'B']]

        def f():
            wp.loc[['Item1', 'Item2'], :, ['A', 'B']] = wp2.loc[
                ['Item1', 'Item2'], :, ['A', 'B']]

        self.assertRaises(NotImplementedError, f)

        # to_assign = wp2.loc[['Item1', 'Item2'], :, ['A', 'B']]
        # wp.loc[['Item1', 'Item2'], :, ['A', 'B']] = to_assign
        # result = wp.loc[['Item1', 'Item2'], :, ['A', 'B']]
        # tm.assert_panel_equal(result,expected)
