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

from __future__ import print_function
from collections import deque
from datetime import datetime
import operator

import pytest

from numpy import nan, random
import numpy as np

from pandas.compat import lrange, range
from pandas import compat
from pandas import (DataFrame, Series, MultiIndex, Timestamp,
                    date_range)
import pandas.core.common as com
import pandas.io.formats.printing as printing
import pandas as pd

from pandas.util.testing import (assert_numpy_array_equal,
                                 assert_series_equal,
                                 assert_frame_equal)

import pandas.util.testing as tm

from pandas.tests.frame.common import (TestData, _check_mixed_float,
                                       _check_mixed_int)


class TestDataFrameOperators(TestData):

    def test_operators(self):
        garbage = random.random(4)
        colSeries = Series(garbage, index=np.array(self.frame.columns))

        idSum = self.frame + self.frame
        seriesSum = self.frame + colSeries

        for col, series in compat.iteritems(idSum):
            for idx, val in compat.iteritems(series):
                origVal = self.frame[col][idx] * 2
                if not np.isnan(val):
                    assert val == origVal
                else:
                    assert np.isnan(origVal)

        for col, series in compat.iteritems(seriesSum):
            for idx, val in compat.iteritems(series):
                origVal = self.frame[col][idx] + colSeries[col]
                if not np.isnan(val):
                    assert val == origVal
                else:
                    assert np.isnan(origVal)

        added = self.frame2 + self.frame2
        expected = self.frame2 * 2
        assert_frame_equal(added, expected)

        df = DataFrame({'a': ['a', None, 'b']})
        assert_frame_equal(df + df, DataFrame({'a': ['aa', np.nan, 'bb']}))

        # Test for issue #10181
        for dtype in ('float', 'int64'):
            frames = [
                DataFrame(dtype=dtype),
                DataFrame(columns=['A'], dtype=dtype),
                DataFrame(index=[0], dtype=dtype),
            ]
            for df in frames:
                assert (df + df).equals(df)
                assert_frame_equal(df + df, df)

    def test_ops_np_scalar(self):
        vals, xs = np.random.rand(5, 3), [nan, 7, -23, 2.718, -3.14, np.inf]
        f = lambda x: DataFrame(x, index=list('ABCDE'),
                                columns=['jim', 'joe', 'jolie'])

        df = f(vals)

        for x in xs:
            assert_frame_equal(df / np.array(x), f(vals / x))
            assert_frame_equal(np.array(x) * df, f(vals * x))
            assert_frame_equal(df + np.array(x), f(vals + x))
            assert_frame_equal(np.array(x) - df, f(x - vals))

    def test_operators_boolean(self):

        # GH 5808
        # empty frames, non-mixed dtype

        result = DataFrame(index=[1]) & DataFrame(index=[1])
        assert_frame_equal(result, DataFrame(index=[1]))

        result = DataFrame(index=[1]) | DataFrame(index=[1])
        assert_frame_equal(result, DataFrame(index=[1]))

        result = DataFrame(index=[1]) & DataFrame(index=[1, 2])
        assert_frame_equal(result, DataFrame(index=[1, 2]))

        result = DataFrame(index=[1], columns=['A']) & DataFrame(
            index=[1], columns=['A'])
        assert_frame_equal(result, DataFrame(index=[1], columns=['A']))

        result = DataFrame(True, index=[1], columns=['A']) & DataFrame(
            True, index=[1], columns=['A'])
        assert_frame_equal(result, DataFrame(True, index=[1], columns=['A']))

        result = DataFrame(True, index=[1], columns=['A']) | DataFrame(
            True, index=[1], columns=['A'])
        assert_frame_equal(result, DataFrame(True, index=[1], columns=['A']))

        # boolean ops
        result = DataFrame(1, index=[1], columns=['A']) | DataFrame(
            True, index=[1], columns=['A'])
        assert_frame_equal(result, DataFrame(1, index=[1], columns=['A']))

        def f():
            DataFrame(1.0, index=[1], columns=['A']) | DataFrame(
                True, index=[1], columns=['A'])
        pytest.raises(TypeError, f)

        def f():
            DataFrame('foo', index=[1], columns=['A']) | DataFrame(
                True, index=[1], columns=['A'])
        pytest.raises(TypeError, f)

    def test_operators_none_as_na(self):
        df = DataFrame({"col1": [2, 5.0, 123, None],
                        "col2": [1, 2, 3, 4]}, dtype=object)

        ops = [operator.add, operator.sub, operator.mul, operator.truediv]

        # since filling converts dtypes from object, changed expected to be
        # object
        for op in ops:
            filled = df.fillna(np.nan)
            result = op(df, 3)
            expected = op(filled, 3).astype(object)
            expected[com.isna(expected)] = None
            assert_frame_equal(result, expected)

            result = op(df, df)
            expected = op(filled, filled).astype(object)
            expected[com.isna(expected)] = None
            assert_frame_equal(result, expected)

            result = op(df, df.fillna(7))
            assert_frame_equal(result, expected)

            result = op(df.fillna(7), df)
            assert_frame_equal(result, expected, check_dtype=False)

    def test_comparison_invalid(self):

        def check(df, df2):

            for (x, y) in [(df, df2), (df2, df)]:
                pytest.raises(TypeError, lambda: x == y)
                pytest.raises(TypeError, lambda: x != y)
                pytest.raises(TypeError, lambda: x >= y)
                pytest.raises(TypeError, lambda: x > y)
                pytest.raises(TypeError, lambda: x < y)
                pytest.raises(TypeError, lambda: x <= y)

        # GH4968
        # invalid date/int comparisons
        df = DataFrame(np.random.randint(10, size=(10, 1)), columns=['a'])
        df['dates'] = date_range('20010101', periods=len(df))

        df2 = df.copy()
        df2['dates'] = df['a']
        check(df, df2)

        df = DataFrame(np.random.randint(10, size=(10, 2)), columns=['a', 'b'])
        df2 = DataFrame({'a': date_range('20010101', periods=len(
            df)), 'b': date_range('20100101', periods=len(df))})
        check(df, df2)

    def test_timestamp_compare(self):
        # make sure we can compare Timestamps on the right AND left hand side
        # GH4982
        df = DataFrame({'dates1': date_range('20010101', periods=10),
                        'dates2': date_range('20010102', periods=10),
                        'intcol': np.random.randint(1000000000, size=10),
                        'floatcol': np.random.randn(10),
                        'stringcol': list(tm.rands(10))})
        df.loc[np.random.rand(len(df)) > 0.5, 'dates2'] = pd.NaT
        ops = {'gt': 'lt', 'lt': 'gt', 'ge': 'le', 'le': 'ge', 'eq': 'eq',
               'ne': 'ne'}

        for left, right in ops.items():
            left_f = getattr(operator, left)
            right_f = getattr(operator, right)

            # no nats
            expected = left_f(df, Timestamp('20010109'))
            result = right_f(Timestamp('20010109'), df)
            assert_frame_equal(result, expected)

            # nats
            expected = left_f(df, Timestamp('nat'))
            result = right_f(Timestamp('nat'), df)
            assert_frame_equal(result, expected)

    def test_modulo(self):
        # GH3590, modulo as ints
        p = DataFrame({'first': [3, 4, 5, 8], 'second': [0, 0, 0, 3]})

        # this is technically wrong as the integer portion is coerced to float
        # ###
        expected = DataFrame({'first': Series([0, 0, 0, 0], dtype='float64'),
                              'second': Series([np.nan, np.nan, np.nan, 0])})
        result = p % p
        assert_frame_equal(result, expected)

        # numpy has a slightly different (wrong) treatement
        with np.errstate(all='ignore'):
            arr = p.values % p.values
        result2 = DataFrame(arr, index=p.index,
                            columns=p.columns, dtype='float64')
        result2.iloc[0:3, 1] = np.nan
        assert_frame_equal(result2, expected)

        result = p % 0
        expected = DataFrame(np.nan, index=p.index, columns=p.columns)
        assert_frame_equal(result, expected)

        # numpy has a slightly different (wrong) treatement
        with np.errstate(all='ignore'):
            arr = p.values.astype('float64') % 0
        result2 = DataFrame(arr, index=p.index, columns=p.columns)
        assert_frame_equal(result2, expected)

        # not commutative with series
        p = DataFrame(np.random.randn(10, 5))
        s = p[0]
        res = s % p
        res2 = p % s
        assert not res.fillna(0).equals(res2.fillna(0))

    def test_div(self):

        # integer div, but deal with the 0's (GH 9144)
        p = DataFrame({'first': [3, 4, 5, 8], 'second': [0, 0, 0, 3]})
        result = p / p

        expected = DataFrame({'first': Series([1.0, 1.0, 1.0, 1.0]),
                              'second': Series([nan, nan, nan, 1])})
        assert_frame_equal(result, expected)

        with np.errstate(all='ignore'):
            arr = p.values.astype('float') / p.values
        result2 = DataFrame(arr, index=p.index,
                            columns=p.columns)
        assert_frame_equal(result2, expected)

        result = p / 0
        expected = DataFrame(np.inf, index=p.index, columns=p.columns)
        expected.iloc[0:3, 1] = nan
        assert_frame_equal(result, expected)

        # numpy has a slightly different (wrong) treatement
        with np.errstate(all='ignore'):
            arr = p.values.astype('float64') / 0
        result2 = DataFrame(arr, index=p.index,
                            columns=p.columns)
        assert_frame_equal(result2, expected)

        p = DataFrame(np.random.randn(10, 5))
        s = p[0]
        res = s / p
        res2 = p / s
        assert not res.fillna(0).equals(res2.fillna(0))

    def test_logical_operators(self):

        def _check_bin_op(op):
            result = op(df1, df2)
            expected = DataFrame(op(df1.values, df2.values), index=df1.index,
                                 columns=df1.columns)
            assert result.values.dtype == np.bool_
            assert_frame_equal(result, expected)

        def _check_unary_op(op):
            result = op(df1)
            expected = DataFrame(op(df1.values), index=df1.index,
                                 columns=df1.columns)
            assert result.values.dtype == np.bool_
            assert_frame_equal(result, expected)

        df1 = {'a': {'a': True, 'b': False, 'c': False, 'd': True, 'e': True},
               'b': {'a': False, 'b': True, 'c': False,
                     'd': False, 'e': False},
               'c': {'a': False, 'b': False, 'c': True,
                     'd': False, 'e': False},
               'd': {'a': True, 'b': False, 'c': False, 'd': True, 'e': True},
               'e': {'a': True, 'b': False, 'c': False, 'd': True, 'e': True}}

        df2 = {'a': {'a': True, 'b': False, 'c': True, 'd': False, 'e': False},
               'b': {'a': False, 'b': True, 'c': False,
                     'd': False, 'e': False},
               'c': {'a': True, 'b': False, 'c': True, 'd': False, 'e': False},
               'd': {'a': False, 'b': False, 'c': False,
                     'd': True, 'e': False},
               'e': {'a': False, 'b': False, 'c': False,
                     'd': False, 'e': True}}

        df1 = DataFrame(df1)
        df2 = DataFrame(df2)

        _check_bin_op(operator.and_)
        _check_bin_op(operator.or_)
        _check_bin_op(operator.xor)

        # operator.neg is deprecated in numpy >= 1.9
        _check_unary_op(operator.inv)

    @pytest.mark.parametrize('op,res', [('__eq__', False),
                                        ('__ne__', True)])
    def test_logical_typeerror_with_non_valid(self, op, res):
        # we are comparing floats vs a string
        result = getattr(self.frame, op)('foo')
        assert bool(result.all().all()) is res

    def test_logical_with_nas(self):
        d = DataFrame({'a': [np.nan, False], 'b': [True, True]})

        # GH4947
        # bool comparisons should return bool
        result = d['a'] | d['b']
        expected = Series([False, True])
        assert_series_equal(result, expected)

        # GH4604, automatic casting here
        result = d['a'].fillna(False) | d['b']
        expected = Series([True, True])
        assert_series_equal(result, expected)

        result = d['a'].fillna(False, downcast=False) | d['b']
        expected = Series([True, True])
        assert_series_equal(result, expected)

    def test_neg(self):
        # what to do?
        assert_frame_equal(-self.frame, -1 * self.frame)

    def test_invert(self):
        assert_frame_equal(-(self.frame < 0), ~(self.frame < 0))

    def test_arith_flex_frame(self):
        ops = ['add', 'sub', 'mul', 'div', 'truediv', 'pow', 'floordiv', 'mod']
        if not compat.PY3:
            aliases = {}
        else:
            aliases = {'div': 'truediv'}

        for op in ops:
            try:
                alias = aliases.get(op, op)
                f = getattr(operator, alias)
                result = getattr(self.frame, op)(2 * self.frame)
                exp = f(self.frame, 2 * self.frame)
                assert_frame_equal(result, exp)

                # vs mix float
                result = getattr(self.mixed_float, op)(2 * self.mixed_float)
                exp = f(self.mixed_float, 2 * self.mixed_float)
                assert_frame_equal(result, exp)
                _check_mixed_float(result, dtype=dict(C=None))

                # vs mix int
                if op in ['add', 'sub', 'mul']:
                    result = getattr(self.mixed_int, op)(2 + self.mixed_int)
                    exp = f(self.mixed_int, 2 + self.mixed_int)

                    # no overflow in the uint
                    dtype = None
                    if op in ['sub']:
                        dtype = dict(B='uint64', C=None)
                    elif op in ['add', 'mul']:
                        dtype = dict(C=None)
                    assert_frame_equal(result, exp)
                    _check_mixed_int(result, dtype=dtype)

                    # rops
                    r_f = lambda x, y: f(y, x)
                    result = getattr(self.frame, 'r' + op)(2 * self.frame)
                    exp = r_f(self.frame, 2 * self.frame)
                    assert_frame_equal(result, exp)

                    # vs mix float
                    result = getattr(self.mixed_float, op)(
                        2 * self.mixed_float)
                    exp = f(self.mixed_float, 2 * self.mixed_float)
                    assert_frame_equal(result, exp)
                    _check_mixed_float(result, dtype=dict(C=None))

                    result = getattr(self.intframe, op)(2 * self.intframe)
                    exp = f(self.intframe, 2 * self.intframe)
                    assert_frame_equal(result, exp)

                    # vs mix int
                    if op in ['add', 'sub', 'mul']:
                        result = getattr(self.mixed_int, op)(
                            2 + self.mixed_int)
                        exp = f(self.mixed_int, 2 + self.mixed_int)

                        # no overflow in the uint
                        dtype = None
                        if op in ['sub']:
                            dtype = dict(B='uint64', C=None)
                        elif op in ['add', 'mul']:
                            dtype = dict(C=None)
                        assert_frame_equal(result, exp)
                        _check_mixed_int(result, dtype=dtype)
            except:
                printing.pprint_thing("Failing operation %r" % op)
                raise

            # ndim >= 3
            ndim_5 = np.ones(self.frame.shape + (3, 4, 5))
            msg = "Unable to coerce to Series/DataFrame"
            with tm.assert_raises_regex(ValueError, msg):
                f(self.frame, ndim_5)

            with tm.assert_raises_regex(ValueError, msg):
                getattr(self.frame, op)(ndim_5)

        # res_add = self.frame.add(self.frame)
        # res_sub = self.frame.sub(self.frame)
        # res_mul = self.frame.mul(self.frame)
        # res_div = self.frame.div(2 * self.frame)

        # assert_frame_equal(res_add, self.frame + self.frame)
        # assert_frame_equal(res_sub, self.frame - self.frame)
        # assert_frame_equal(res_mul, self.frame * self.frame)
        # assert_frame_equal(res_div, self.frame / (2 * self.frame))

        const_add = self.frame.add(1)
        assert_frame_equal(const_add, self.frame + 1)

        # corner cases
        result = self.frame.add(self.frame[:0])
        assert_frame_equal(result, self.frame * np.nan)

        result = self.frame[:0].add(self.frame)
        assert_frame_equal(result, self.frame * np.nan)
        with tm.assert_raises_regex(NotImplementedError, 'fill_value'):
            self.frame.add(self.frame.iloc[0], fill_value=3)
        with tm.assert_raises_regex(NotImplementedError, 'fill_value'):
            self.frame.add(self.frame.iloc[0], axis='index', fill_value=3)

    def test_binary_ops_align(self):

        # test aligning binary ops

        # GH 6681
        index = MultiIndex.from_product([list('abc'),
                                         ['one', 'two', 'three'],
                                         [1, 2, 3]],
                                        names=['first', 'second', 'third'])

        df = DataFrame(np.arange(27 * 3).reshape(27, 3),
                       index=index,
                       columns=['value1', 'value2', 'value3']).sort_index()

        idx = pd.IndexSlice
        for op in ['add', 'sub', 'mul', 'div', 'truediv']:
            opa = getattr(operator, op, None)
            if opa is None:
                continue

            x = Series([1.0, 10.0, 100.0], [1, 2, 3])
            result = getattr(df, op)(x, level='third', axis=0)

            expected = pd.concat([opa(df.loc[idx[:, :, i], :], v)
                                  for i, v in x.iteritems()]).sort_index()
            assert_frame_equal(result, expected)

            x = Series([1.0, 10.0], ['two', 'three'])
            result = getattr(df, op)(x, level='second', axis=0)

            expected = (pd.concat([opa(df.loc[idx[:, i], :], v)
                                   for i, v in x.iteritems()])
                        .reindex_like(df).sort_index())
            assert_frame_equal(result, expected)

        # GH9463 (alignment level of dataframe with series)

        midx = MultiIndex.from_product([['A', 'B'], ['a', 'b']])
        df = DataFrame(np.ones((2, 4), dtype='int64'), columns=midx)
        s = pd.Series({'a': 1, 'b': 2})

        df2 = df.copy()
        df2.columns.names = ['lvl0', 'lvl1']
        s2 = s.copy()
        s2.index.name = 'lvl1'

        # different cases of integer/string level names:
        res1 = df.mul(s, axis=1, level=1)
        res2 = df.mul(s2, axis=1, level=1)
        res3 = df2.mul(s, axis=1, level=1)
        res4 = df2.mul(s2, axis=1, level=1)
        res5 = df2.mul(s, axis=1, level='lvl1')
        res6 = df2.mul(s2, axis=1, level='lvl1')

        exp = DataFrame(np.array([[1, 2, 1, 2], [1, 2, 1, 2]], dtype='int64'),
                        columns=midx)

        for res in [res1, res2]:
            assert_frame_equal(res, exp)

        exp.columns.names = ['lvl0', 'lvl1']
        for res in [res3, res4, res5, res6]:
            assert_frame_equal(res, exp)

    def test_arith_mixed(self):

        left = DataFrame({'A': ['a', 'b', 'c'],
                          'B': [1, 2, 3]})

        result = left + left
        expected = DataFrame({'A': ['aa', 'bb', 'cc'],
                              'B': [2, 4, 6]})
        assert_frame_equal(result, expected)

    def test_arith_getitem_commute(self):
        df = DataFrame({'A': [1.1, 3.3], 'B': [2.5, -3.9]})

        self._test_op(df, operator.add)
        self._test_op(df, operator.sub)
        self._test_op(df, operator.mul)
        self._test_op(df, operator.truediv)
        self._test_op(df, operator.floordiv)
        self._test_op(df, operator.pow)

        self._test_op(df, lambda x, y: y + x)
        self._test_op(df, lambda x, y: y - x)
        self._test_op(df, lambda x, y: y * x)
        self._test_op(df, lambda x, y: y / x)
        self._test_op(df, lambda x, y: y ** x)

        self._test_op(df, lambda x, y: x + y)
        self._test_op(df, lambda x, y: x - y)
        self._test_op(df, lambda x, y: x * y)
        self._test_op(df, lambda x, y: x / y)
        self._test_op(df, lambda x, y: x ** y)

    @staticmethod
    def _test_op(df, op):
        result = op(df, 1)

        if not df.columns.is_unique:
            raise ValueError("Only unique columns supported by this test")

        for col in result.columns:
            assert_series_equal(result[col], op(df[col], 1))

    def test_bool_flex_frame(self):