import pytest
import datetime

from pymapd._loaders import _build_input_rows
from pymapd import _pandas_loaders
from mapd.MapD import TStringRow, TStringValue, TColumn, TColumnData

from .utils import assert_columnar_equal


class TestLoaders(object):

    def test_build_input_rows(self):
        data = [(1, 'a'), (2, 'b')]
        result = _build_input_rows(data)
        expected = [TStringRow(cols=[TStringValue(str_val='1', is_null=None),
                                     TStringValue(str_val='a', is_null=None)]),
                    TStringRow(cols=[TStringValue(str_val='2', is_null=None),
                                     TStringValue(str_val='b', is_null=None)])]

        assert result == expected

    def test_build_table_columnar(self):
        pd = pytest.importorskip("pandas")
        pytest.importorskip("pyarrow")
        from pymapd._pandas_loaders import build_input_columnar

        data = pd.DataFrame({"a": [1, 2, 3], "b": [1.1, 2.2, 3.3]})
        nulls = [False] * 3
        result = build_input_columnar(data, preserve_index=False)
        expected = [
            TColumn(TColumnData(int_col=[1, 2, 3]), nulls=nulls),
            TColumn(TColumnData(real_col=[1.1, 2.2, 3.3]), nulls=nulls)
        ]
        assert_columnar_equal(result, expected)

    def test_build_table_columnar_pandas(self):
        import pandas as pd
        import numpy as np

        data = pd.DataFrame({
            "boolean_": [True, False],
            "smallint_": np.array([0, 1], dtype=np.int8),
            "int_": np.array([0, 1], dtype=np.int32),
            "bigint_": np.array([0, 1], dtype=np.int64),
            "float_": np.array([0, 1], dtype=np.float32),
            "double_": np.array([0, 1], dtype=np.float64),
            "varchar_": ["a", "b"],
            "text_": ['a', 'b'],
            "time_": [datetime.time(0, 11, 59), datetime.time(13)],
            "timestamp_": [pd.Timestamp("2016"), pd.Timestamp("2017")],
            "date_": [datetime.date(2016, 1, 1), datetime.date(2017, 1, 1)],
        }, columns=['boolean_', 'smallint_', 'int_', 'bigint_', 'float_',
                    'double_', 'varchar_', 'text_', 'time_', 'timestamp_',
                    'date_'])
        result = _pandas_loaders.build_input_columnar(data,
                                                      preserve_index=False)

        nulls = [False, False]
        expected = [
            TColumn(TColumnData(int_col=[True, False]), nulls=nulls),
            TColumn(TColumnData(int_col=np.array([0, 1], dtype=np.int8)), nulls=nulls),  # noqa
            TColumn(TColumnData(int_col=np.array([0, 1], dtype=np.int32)), nulls=nulls),  # noqa
            TColumn(TColumnData(int_col=np.array([0, 1], dtype=np.int64)), nulls=nulls),  # noqa
            TColumn(TColumnData(real_col=np.array([0, 1], dtype=np.float32)), nulls=nulls),  # noqa
            TColumn(TColumnData(real_col=np.array([0, 1], dtype=np.float64)), nulls=nulls),  # noqa
            TColumn(TColumnData(str_col=['a', 'b']), nulls=nulls),
            TColumn(TColumnData(str_col=['a', 'b']), nulls=nulls),
            TColumn(TColumnData(int_col=[719, 46800]), nulls=nulls),
            TColumn(TColumnData(int_col=[1451606400, 1483228800]), nulls=nulls),  # noqa
            TColumn(TColumnData(int_col=[1451606400, 1483228800]), nulls=nulls)
        ]
        assert_columnar_equal(result, expected)

    def test_build_table_columnar_nulls(self):
        import pandas as pd
        import numpy as np

        data = pd.DataFrame({
            "boolean_": [True, False, None],
            "bigint_": np.array([0, 1, None], dtype=np.object),
            "double_": np.array([0, 1, None], dtype=np.float64),
            "varchar_": ["a", "b", None],
            "text_": ['a', 'b', None],
            "time_": [datetime.time(0, 11, 59), datetime.time(13), None],
            "timestamp_": [pd.Timestamp("2016"), pd.Timestamp("2017"), None],
            "date_": [datetime.date(2016, 1, 1), datetime.date(2017, 1, 1),
                      None],
        }, columns=['boolean_', 'bigint_',
                    'double_', 'varchar_', 'text_', 'time_', 'timestamp_',
                    'date_'])
        result = _pandas_loaders.build_input_columnar(data,
                                                      preserve_index=False)

        nulls = [False, False, True]
        int_na = -2147483648
        bigint_na = -9223372036854775808
        ns_na = -9223372037

        expected = [
            TColumn(TColumnData(int_col=[1, 0, int_na]), nulls=nulls),
            TColumn(TColumnData(int_col=np.array([0, 1, int_na], dtype=np.int64)), nulls=nulls),  # noqa
            TColumn(TColumnData(real_col=np.array([0, 1, np.nan], dtype=np.float64)), nulls=nulls),  # noqa
            TColumn(TColumnData(str_col=['a', 'b', '']), nulls=nulls),
            TColumn(TColumnData(str_col=['a', 'b', '']), nulls=nulls),
            TColumn(TColumnData(int_col=[719, 46800, bigint_na]), nulls=nulls),
            TColumn(TColumnData(int_col=[1451606400, 1483228800, ns_na]), nulls=nulls),  # noqa
            TColumn(TColumnData(int_col=[1451606400, 1483228800, bigint_na]), nulls=nulls)  # noqa
        ]
        assert_columnar_equal(result, expected)

    def test_build_row_desc(self):
        pd = pytest.importorskip("pandas")
        import numpy as np
        from mapd.ttypes import TTypeInfo, TColumnType

        data = pd.DataFrame({
            "boolean_": [True, False],
            "smallint_": np.array([0, 1], dtype=np.int8),
            "int_": np.array([0, 1], dtype=np.int32),
            "bigint_": np.array([0, 1], dtype=np.int64),
            "float_": np.array([0, 1], dtype=np.float32),
            "double_": np.array([0, 1], dtype=np.float64),
            "varchar_": ["a", "b"],
            "text_": ['a', 'b'],
            "time_": [datetime.time(0, 11, 59), datetime.time(13)],
            "timestamp_": [pd.Timestamp("2016"), pd.Timestamp("2017")],
            "date_": [datetime.date(2016, 1, 1), datetime.date(2017, 1, 1)],
        }, columns=['boolean_', 'smallint_', 'int_', 'bigint_', 'float_',
                    'double_', 'varchar_', 'text_', 'time_', 'timestamp_',
                    'date_'])
        result = _pandas_loaders.build_row_desc(data)
        expected = [
            TColumnType(col_name='boolean_',
                        col_type=TTypeInfo(type=10),
                        is_reserved_keyword=None),
            TColumnType(col_name='smallint_',
                        col_type=TTypeInfo(type=0),
                        is_reserved_keyword=None),
            TColumnType(col_name='int_',
                        col_type=TTypeInfo(type=1),
                        is_reserved_keyword=None),
            TColumnType(col_name='bigint_',
                        col_type=TTypeInfo(type=2)),
            TColumnType(col_name='float_',
                        col_type=TTypeInfo(type=3)),
            TColumnType(col_name='double_',
                        col_type=TTypeInfo(type=5)),
            TColumnType(col_name='varchar_',
                        col_type=TTypeInfo(type=6)),
            TColumnType(col_name='text_',
                        col_type=TTypeInfo(type=6)),
            TColumnType(col_name='time_',
                        col_type=TTypeInfo(type=7)),
            TColumnType(col_name='timestamp_',
                        col_type=TTypeInfo(type=8)),
            TColumnType(col_name='date_',
                        col_type=TTypeInfo(type=9))
        ]

        assert result == expected

        data.index.name = 'idx'
        result = _pandas_loaders.build_row_desc(data, preserve_index=True)
        expected.insert(0, TColumnType(col_name='idx',
                                       col_type=TTypeInfo(type=2)))

        assert result == expected

    def test_create_non_pandas_raises(self):
        with pytest.raises(TypeError) as m:
            _pandas_loaders.build_row_desc([(1, 'a'), (2, 'b')])

        assert m.match('is not supported for type ')

    @pytest.mark.parametrize('create', ['infer', True, False])
    def test_check_create(self, create):
        from pymapd.connection import _check_create

        result = _check_create(create)
        assert result == create

    def test_check_create_raises(self):
        from pymapd.connection import _check_create

        with pytest.raises(ValueError):
            _check_create('foo')
