# -*- coding: utf-8 -*-
# pylint: disable=W0102

from datetime import datetime, date
import operator
import sys
import pytest
import numpy as np

import re
from distutils.version import LooseVersion
import itertools
from pandas import (Index, MultiIndex, DataFrame, DatetimeIndex,
                    Series, Categorical, TimedeltaIndex, SparseArray)
from pandas.compat import OrderedDict, lrange
from pandas.core.internals import (SingleBlockManager,
                                   make_block, BlockManager)
import pandas.core.algorithms as algos
import pandas.util.testing as tm
import pandas as pd
from pandas._libs.internals import BlockPlacement
from pandas.util.testing import (assert_almost_equal, assert_frame_equal,
                                 randn, assert_series_equal)
from pandas.compat import zip, u

# in 3.6.1 a c-api slicing function changed, see src/compat_helper.h
PY361 = LooseVersion(sys.version) >= LooseVersion('3.6.1')


@pytest.fixture
def mgr():
    return create_mgr(
        'a: f8; b: object; c: f8; d: object; e: f8;'
        'f: bool; g: i8; h: complex; i: datetime-1; j: datetime-2;'
        'k: M8[ns, US/Eastern]; l: M8[ns, CET];')


def assert_block_equal(left, right):
    tm.assert_numpy_array_equal(left.values, right.values)
    assert left.dtype == right.dtype
    assert isinstance(left.mgr_locs, BlockPlacement)
    assert isinstance(right.mgr_locs, BlockPlacement)
    tm.assert_numpy_array_equal(left.mgr_locs.as_array,
                                right.mgr_locs.as_array)


def get_numeric_mat(shape):
    arr = np.arange(shape[0])
    return np.lib.stride_tricks.as_strided(x=arr, shape=shape, strides=(
        arr.itemsize, ) + (0, ) * (len(shape) - 1)).copy()


N = 10


def create_block(typestr, placement, item_shape=None, num_offset=0):
    """
    Supported typestr:

        * float, f8, f4, f2
        * int, i8, i4, i2, i1
        * uint, u8, u4, u2, u1
        * complex, c16, c8
        * bool
        * object, string, O
        * datetime, dt, M8[ns], M8[ns, tz]
        * timedelta, td, m8[ns]
        * sparse (SparseArray with fill_value=0.0)
        * sparse_na (SparseArray with fill_value=np.nan)
        * category, category2

    """
    placement = BlockPlacement(placement)
    num_items = len(placement)

    if item_shape is None:
        item_shape = (N, )

    shape = (num_items, ) + item_shape

    mat = get_numeric_mat(shape)

    if typestr in ('float', 'f8', 'f4', 'f2', 'int', 'i8', 'i4', 'i2', 'i1',
                   'uint', 'u8', 'u4', 'u2', 'u1'):
        values = mat.astype(typestr) + num_offset
    elif typestr in ('complex', 'c16', 'c8'):
        values = 1.j * (mat.astype(typestr) + num_offset)
    elif typestr in ('object', 'string', 'O'):
        values = np.reshape(['A%d' % i for i in mat.ravel() + num_offset],
                            shape)
    elif typestr in ('b', 'bool', ):
        values = np.ones(shape, dtype=np.bool_)
    elif typestr in ('datetime', 'dt', 'M8[ns]'):
        values = (mat * 1e9).astype('M8[ns]')
    elif typestr.startswith('M8[ns'):
        # datetime with tz
        m = re.search(r'M8\[ns,\s*(\w+\/?\w*)\]', typestr)
        assert m is not None, "incompatible typestr -> {0}".format(typestr)
        tz = m.groups()[0]
        assert num_items == 1, "must have only 1 num items for a tz-aware"
        values = DatetimeIndex(np.arange(N) * 1e9, tz=tz)
    elif typestr in ('timedelta', 'td', 'm8[ns]'):
        values = (mat * 1).astype('m8[ns]')
    elif typestr in ('category', ):
        values = Categorical([1, 1, 2, 2, 3, 3, 3, 3, 4, 4])
    elif typestr in ('category2', ):
        values = Categorical(['a', 'a', 'a', 'a', 'b', 'b', 'c', 'c', 'c', 'd'
                              ])
    elif typestr in ('sparse', 'sparse_na'):
        # FIXME: doesn't support num_rows != 10
        assert shape[-1] == 10
        assert all(s == 1 for s in shape[:-1])
        if typestr.endswith('_na'):
            fill_value = np.nan
        else:
            fill_value = 0.0
        values = SparseArray([fill_value, fill_value, 1, 2, 3, fill_value,
                              4, 5, fill_value, 6], fill_value=fill_value)
        arr = values.sp_values.view()
        arr += (num_offset - 1)
    else:
        raise ValueError('Unsupported typestr: "%s"' % typestr)

    return make_block(values, placement=placement, ndim=len(shape))


def create_single_mgr(typestr, num_rows=None):
    if num_rows is None:
        num_rows = N

    return SingleBlockManager(
        create_block(typestr, placement=slice(0, num_rows), item_shape=()),
        np.arange(num_rows))


def create_mgr(descr, item_shape=None):
    """
    Construct BlockManager from string description.

    String description syntax looks similar to np.matrix initializer.  It looks
    like this::

        a,b,c: f8; d,e,f: i8

    Rules are rather simple:

    * see list of supported datatypes in `create_block` method
    * components are semicolon-separated
    * each component is `NAME,NAME,NAME: DTYPE_ID`
    * whitespace around colons & semicolons are removed
    * components with same DTYPE_ID are combined into single block
    * to force multiple blocks with same dtype, use '-SUFFIX'::

        'a:f8-1; b:f8-2; c:f8-foobar'

    """
    if item_shape is None:
        item_shape = (N, )

    offset = 0
    mgr_items = []
    block_placements = OrderedDict()
    for d in descr.split(';'):
        d = d.strip()
        if not len(d):
            continue
        names, blockstr = d.partition(':')[::2]
        blockstr = blockstr.strip()
        names = names.strip().split(',')

        mgr_items.extend(names)
        placement = list(np.arange(len(names)) + offset)
        try:
            block_placements[blockstr].extend(placement)
        except KeyError:
            block_placements[blockstr] = placement
        offset += len(names)

    mgr_items = Index(mgr_items)

    blocks = []
    num_offset = 0
    for blockstr, placement in block_placements.items():
        typestr = blockstr.split('-')[0]
        blocks.append(create_block(typestr,
                                   placement,
                                   item_shape=item_shape,
                                   num_offset=num_offset, ))
        num_offset += len(placement)

    return BlockManager(sorted(blocks, key=lambda b: b.mgr_locs[0]),
                        [mgr_items] + [np.arange(n) for n in item_shape])


class TestBlock(object):

    def setup_method(self, method):
        # self.fblock = get_float_ex()  # a,c,e
        # self.cblock = get_complex_ex() #
        # self.oblock = get_obj_ex()
        # self.bool_block = get_bool_ex()
        # self.int_block = get_int_ex()

        self.fblock = create_block('float', [0, 2, 4])
        self.cblock = create_block('complex', [7])
        self.oblock = create_block('object', [1, 3])
        self.bool_block = create_block('bool', [5])
        self.int_block = create_block('int', [6])

    def test_constructor(self):
        int32block = create_block('i4', [0])
        assert int32block.dtype == np.int32

    def test_pickle(self):
        def _check(blk):
            assert_block_equal(tm.round_trip_pickle(blk), blk)

        _check(self.fblock)
        _check(self.cblock)
        _check(self.oblock)
        _check(self.bool_block)

    def test_mgr_locs(self):
        assert isinstance(self.fblock.mgr_locs, BlockPlacement)
        tm.assert_numpy_array_equal(self.fblock.mgr_locs.as_array,
                                    np.array([0, 2, 4], dtype=np.int64))

    def test_attrs(self):
        assert self.fblock.shape == self.fblock.values.shape
        assert self.fblock.dtype == self.fblock.values.dtype
        assert len(self.fblock) == len(self.fblock.values)

    def test_merge(self):
        avals = randn(2, 10)
        bvals = randn(2, 10)

        ref_cols = Index(['e', 'a', 'b', 'd', 'f'])

        ablock = make_block(avals, ref_cols.get_indexer(['e', 'b']))
        bblock = make_block(bvals, ref_cols.get_indexer(['a', 'd']))
        merged = ablock.merge(bblock)
        tm.assert_numpy_array_equal(merged.mgr_locs.as_array,
                                    np.array([0, 1, 2, 3], dtype=np.int64))
        tm.assert_numpy_array_equal(merged.values[[0, 2]], np.array(avals))
        tm.assert_numpy_array_equal(merged.values[[1, 3]], np.array(bvals))

        # TODO: merge with mixed type?

    def test_copy(self):
        cop = self.fblock.copy()
        assert cop is not self.fblock
        assert_block_equal(self.fblock, cop)

    def test_reindex_index(self):
        pass

    def test_reindex_cast(self):
        pass

    def test_insert(self):
        pass

    def test_delete(self):
        newb = self.fblock.copy()
        newb.delete(0)
        assert isinstance(newb.mgr_locs, BlockPlacement)
        tm.assert_numpy_array_equal(newb.mgr_locs.as_array,
                                    np.array([2, 4], dtype=np.int64))
        assert (newb.values[0] == 1).all()

        newb = self.fblock.copy()
        newb.delete(1)
        assert isinstance(newb.mgr_locs, BlockPlacement)
        tm.assert_numpy_array_equal(newb.mgr_locs.as_array,
                                    np.array([0, 4], dtype=np.int64))
        assert (newb.values[1] == 2).all()

        newb = self.fblock.copy()
        newb.delete(2)
        tm.assert_numpy_array_equal(newb.mgr_locs.as_array,
                                    np.array([0, 2], dtype=np.int64))
        assert (newb.values[1] == 1).all()

        newb = self.fblock.copy()
        with pytest.raises(Exception):
            newb.delete(3)

    def test_make_block_same_class(self):
        # issue 19431
        block = create_block('M8[ns, US/Eastern]', [3])
        with tm.assert_produces_warning(DeprecationWarning,
                                        check_stacklevel=False):
            block.make_block_same_class(block.values.values,
                                        dtype=block.values.dtype)


class TestDatetimeBlock(object):

    def test_try_coerce_arg(self):
        block = create_block('datetime', [0])

        # coerce None
        none_coerced = block._try_coerce_args(block.values, None)[2]
        assert pd.Timestamp(none_coerced) is pd.NaT

        # coerce different types of date bojects
        vals = (np.datetime64('2010-10-10'), datetime(2010, 10, 10),
                date(2010, 10, 10))
        for val in vals:
            coerced = block._try_coerce_args(block.values, val)[2]
            assert np.int64 == type(coerced)
            assert pd.Timestamp('2010-10-10') == pd.Timestamp(coerced)


class TestBlockManager(object):

    def test_constructor_corner(self):
        pass

    def test_attrs(self):
        mgr = create_mgr('a,b,c: f8-1; d,e,f: f8-2')
        assert mgr.nblocks == 2
        assert len(mgr) == 6

    def test_is_mixed_dtype(self):
        assert not create_mgr('a,b:f8').is_mixed_type
        assert not create_mgr('a:f8-1; b:f8-2').is_mixed_type

        assert create_mgr('a,b:f8; c,d: f4').is_mixed_type
        assert create_mgr('a,b:f8; c,d: object').is_mixed_type

    def test_duplicate_ref_loc_failure(self):
        tmp_mgr = create_mgr('a:bool; a: f8')

        axes, blocks = tmp_mgr.axes, tmp_mgr.blocks

        blocks[0].mgr_locs = np.array([0])
        blocks[1].mgr_locs = np.array([0])

        # test trying to create block manager with overlapping ref locs
        with pytest.raises(AssertionError):
            BlockManager(blocks, axes)

        blocks[0].mgr_locs = np.array([0])
        blocks[1].mgr_locs = np.array([1])
        mgr = BlockManager(blocks, axes)
        mgr.iget(1)

    def test_contains(self, mgr):
        assert 'a' in mgr
        assert 'baz' not in mgr

    def test_pickle(self, mgr):

        mgr2 = tm.round_trip_pickle(mgr)
        assert_frame_equal(DataFrame(mgr), DataFrame(mgr2))

        # share ref_items
        # assert mgr2.blocks[0].ref_items is mgr2.blocks[1].ref_items

        # GH2431
        assert hasattr(mgr2, "_is_consolidated")
        assert hasattr(mgr2, "_known_consolidated")

        # reset to False on load
        assert not mgr2._is_consolidated
        assert not mgr2._known_consolidated

    def test_non_unique_pickle(self):

        mgr = create_mgr('a,a,a:f8')
        mgr2 = tm.round_trip_pickle(mgr)
        assert_frame_equal(DataFrame(mgr), DataFrame(mgr2))

        mgr = create_mgr('a: f8; a: i8')
        mgr2 = tm.round_trip_pickle(mgr)
        assert_frame_equal(DataFrame(mgr), DataFrame(mgr2))

    def test_categorical_block_pickle(self):
        mgr = create_mgr('a: category')
        mgr2 = tm.round_trip_pickle(mgr)
        assert_frame_equal(DataFrame(mgr), DataFrame(mgr2))

        smgr = create_single_mgr('category')
        smgr2 = tm.round_trip_pickle(smgr)
        assert_series_equal(Series(smgr), Series(smgr2))

    def test_get(self):
        cols = Index(list('abc'))
        values = np.random.rand(3, 3)
        block = make_block(values=values.copy(), placement=np.arange(3))
        mgr = BlockManager(blocks=[block], axes=[cols, np.arange(3)])

        assert_almost_equal(mgr.get('a', fastpath=False), values[0])
        assert_almost_equal(mgr.get('b', fastpath=False), values[1])
        assert_almost_equal(mgr.get('c', fastpath=False), values[2])
        assert_almost_equal(mgr.get('a').internal_values(), values[0])
        assert_almost_equal(mgr.get('b').internal_values(), values[1])
        assert_almost_equal(mgr.get('c').internal_values(), values[2])

    def test_set(self):
        mgr = create_mgr('a,b,c: int', item_shape=(3, ))

        mgr.set('d', np.array(['foo'] * 3))
        mgr.set('b', np.array(['bar'] * 3))
        tm.assert_numpy_array_equal(mgr.get('a').internal_values(),
                                    np.array([0] * 3))
        tm.assert_numpy_array_equal(mgr.get('b').internal_values(),
                                    np.array(['bar'] * 3, dtype=np.object_))
        tm.assert_numpy_array_equal(mgr.get('c').internal_values(),
                                    np.array([2] * 3))
        tm.assert_numpy_array_equal(mgr.get('d').internal_values(),
                                    np.array(['foo'] * 3, dtype=np.object_))

    def test_set_change_dtype(self, mgr):
        mgr.set('baz', np.zeros(N, dtype=bool))

        mgr.set('baz', np.repeat('foo', N))
        assert mgr.get('baz').dtype == np.object_

        mgr2 = mgr.consolidate()
        mgr2.set('baz', np.repeat('foo', N))
        assert mgr2.get('baz').dtype == np.object_

        mgr2.set('quux', randn(N).astype(int))
        assert mgr2.get('quux').dtype == np.int_

        mgr2.set('quux', randn(N))
        assert mgr2.get('quux').dtype == np.float_

    def test_set_change_dtype_slice(self):  # GH8850
        cols = MultiIndex.from_tuples([('1st', 'a'), ('2nd', 'b'), ('3rd', 'c')
                                       ])
        df = DataFrame([[1.0, 2, 3], [4.0, 5, 6]], columns=cols)
        df['2nd'] = df['2nd'] * 2.0

        blocks = df._to_dict_of_blocks()
        assert sorted(blocks.keys()) == ['float64', 'int64']
        assert_frame_equal(blocks['float64'], DataFrame(
            [[1.0, 4.0], [4.0, 10.0]], columns=cols[:2]))
        assert_frame_equal(blocks['int64'], DataFrame(
            [[3], [6]], columns=cols[2:]))

    def test_copy(self, mgr):
        cp = mgr.copy(deep=False)
        for blk, cp_blk in zip(mgr.blocks, cp.blocks):

            # view assertion
            assert cp_blk.equals(blk)
            if isinstance(blk.values, np.ndarray):
                assert cp_blk.values.base is blk.values.base
            else:
                # DatetimeTZBlock has DatetimeIndex values
                assert cp_blk.values.values.base is blk.values.values.base

        cp = mgr.copy(deep=True)
        for blk, cp_blk in zip(mgr.blocks, cp.blocks):

            # copy assertion we either have a None for a base or in case of
            # some blocks it is an array (e.g. datetimetz), but was copied
            assert cp_blk.equals(blk)
            if not isinstance(cp_blk.values, np.ndarray):
                assert cp_blk.values.values.base is not blk.values.values.base
            else:
                assert cp_blk.values.base is None and blk.values.base is None

    def test_sparse(self):
        mgr = create_mgr('a: sparse-1; b: sparse-2')
        # what to test here?
        assert mgr.as_array().dtype == np.float64

    def test_sparse_mixed(self):
        mgr = create_mgr('a: sparse-1; b: sparse-2; c: f8')
        assert len(mgr.blocks) == 3
        assert isinstance(mgr, BlockManager)

        # what to test here?

    def test_as_array_float(self):
        mgr = create_mgr('c: f4; d: f2; e: f8')
        assert mgr.as_array().dtype == np.float64

        mgr = create_mgr('c: f4; d: f2')
        assert mgr.as_array().dtype == np.float32

    def test_as_array_int_bool(self):
        mgr = create_mgr('a: bool-1; b: bool-2')
        assert mgr.as_array().dtype == np.bool_

        mgr = create_mgr('a: i8-1; b: i8-2; c: i4; d: i2; e: u1')
        assert mgr.as_array().dtype == np.int64

        mgr = create_mgr('c: i4; d: i2; e: u1')
        assert mgr.as_array().dtype == np.int32

    def test_as_array_datetime(self):
        mgr = create_mgr('h: datetime-1; g: datetime-2')
        assert mgr.as_array().dtype == 'M8[ns]'

    def test_as_array_datetime_tz(self):
        mgr = create_mgr('h: M8[ns, US/Eastern]; g: M8[ns, CET]')
        assert mgr.get('h').dtype == 'datetime64[ns, US/Eastern]'
        assert mgr.get('g').dtype == 'datetime64[ns, CET]'
        assert mgr.as_array().dtype == 'object'

    def test_astype(self):
        # coerce all
        mgr = create_mgr('c: f4; d: f2; e: f8')
        for t in ['float16', 'float32', 'float64', 'int32', 'int64']:
            t = np.dtype(t)