from unittest import TestCase
from datetime import datetime
import pyarrow as pa
import numpy as np
import pandas as pd
from h1st.schema import SchemaInferrer


class SchemaInferrerTestCase(TestCase):
    def test_infer_python(self):
        inferrer = SchemaInferrer()

        self.assertEqual(inferrer.infer_schema(1), pa.int64())
        self.assertEqual(inferrer.infer_schema(1.1), pa.float64())

        self.assertEqual(inferrer.infer_schema({
            'test1': 1,
            'test2': "hello",
            'test3': b"hello",
            'today': datetime.now(),
        }), {
            'type': dict,
            'fields': {
                'test1': pa.int64(),
                'test2': pa.string(),
                'test3': pa.binary(),
                'today': pa.date64(),
            }
        })

        self.assertEqual(inferrer.infer_schema((
            1, 2, 3
        )), pa.list_(pa.int64()))

        self.assertEqual(inferrer.infer_schema((
            1.2, 1.3, 1.4
        )), pa.list_(pa.float64()))

        table = pa.Table.from_arrays(
            [pa.array([1, 2, 3]), pa.array(["a", "b", "c"])],
            ['c1', 'c2']
        )
        self.assertEqual(inferrer.infer_schema(table), table.schema)

    def test_infer_numpy(self):
        inferrer = SchemaInferrer()
        self.assertEqual(inferrer.infer_schema(np.random.random((100, 28, 28))), {
            'type': np.ndarray,
            'item': pa.float64(),
            'shape': (None, 28, 28)
        })

        self.assertEqual(inferrer.infer_schema(np.array(["1", "2", "3"])), {
            'type': np.ndarray,
            'item': pa.string()
        })

    def test_infer_dataframe(self):
        inferrer = SchemaInferrer()
        df = pd.DataFrame({
            'f1': [1, 2, 3],
            'f2': ['a', 'b', 'c'],
            'f3': [0.1, 0.2, 0.9]
        })

        self.assertEqual(inferrer.infer_schema(df), {
            'type': pd.DataFrame,
            'fields': {
                'f1': pa.int64(),
                'f2': pa.string(),
                'f3': pa.float64()
            }
        })

        df = pd.DataFrame({
            'Timestamp': [1.1, 2.2, 3.1],
            'CarSpeed': [0.1, 0.2, 0.9],
            'Gx': [0.1, 0.2, 0.9],
            'Gy': [0.1, 0.2, 0.9],
            'Label': ['1', '0', '1']
        })

        self.assertEqual(inferrer.infer_schema(df), {
            'type': pd.DataFrame,
            'fields': {
                'Timestamp': pa.float64(),
                'CarSpeed': pa.float64(),
                'Gx': pa.float64(),
                'Gy': pa.float64(),
                'Label': pa.string(),
            }
        })

        self.assertEqual(inferrer.infer_schema(pd.Series([1, 2, 3])), {
            'type': pd.Series,
            'item': pa.int64()
        })

    def test_infer_dict(self):
        inferrer = SchemaInferrer()
        self.assertEqual(inferrer.infer_schema({
            'test': 123,
        }), {
            'type': dict,
            'fields': {
                'test': pa.int64(),
            }
        })

        self.assertEqual(inferrer.infer_schema({
            'test': 123,
            'indices': [1, 2, 3]
        }), {
            'type': dict,
            'fields': {
                'test': pa.int64(),
                'indices': pa.list_(pa.int64())
            }
        })

        self.assertEqual(inferrer.infer_schema({
            'results': pd.DataFrame({
                'CarSpeed': [0, 1, 2],
                'Label': ['a', 'b', 'c']
            })
        }), {
            'type': dict,
            'fields': {
                'results': {
                    'type': pd.DataFrame,
                    'fields': {
                        'CarSpeed': pa.int64(),
                        'Label': pa.string(),
                    }
                }
            }
        })

    def test_infer_list(self):
        inferrer = SchemaInferrer()
        self.assertEqual(inferrer.infer_schema([
            {'test': 123},
            {'test': 345},
        ]), {
            'type': list,
            'item': {
                'type': dict,
                'fields': {
                    'test': pa.int64()
                }
            }
        })

[((253, 269), 'h1st.schema.SchemaInferrer', 'SchemaInferrer', ([], {}), '()\n', (267, 269), False, 'from h1st.schema import SchemaInferrer\n'), ((1286, 1302), 'h1st.schema.SchemaInferrer', 'SchemaInferrer', ([], {}), '()\n', (1300, 1302), False, 'from h1st.schema import SchemaInferrer\n'), ((1708, 1724), 'h1st.schema.SchemaInferrer', 'SchemaInferrer', ([], {}), '()\n', (1722, 1724), False, 'from h1st.schema import SchemaInferrer\n'), ((1738, 1815), 'pandas.DataFrame', 'pd.DataFrame', (["{'f1': [1, 2, 3], 'f2': ['a', 'b', 'c'], 'f3': [0.1, 0.2, 0.9]}"], {}), "({'f1': [1, 2, 3], 'f2': ['a', 'b', 'c'], 'f3': [0.1, 0.2, 0.9]})\n", (1750, 1815), True, 'import pandas as pd\n'), ((2118, 2267), 'pandas.DataFrame', 'pd.DataFrame', (["{'Timestamp': [1.1, 2.2, 3.1], 'CarSpeed': [0.1, 0.2, 0.9], 'Gx': [0.1, 0.2,\n    0.9], 'Gy': [0.1, 0.2, 0.9], 'Label': ['1', '0', '1']}"], {}), "({'Timestamp': [1.1, 2.2, 3.1], 'CarSpeed': [0.1, 0.2, 0.9],\n    'Gx': [0.1, 0.2, 0.9], 'Gy': [0.1, 0.2, 0.9], 'Label': ['1', '0', '1']})\n", (2130, 2267), True, 'import pandas as pd\n'), ((2864, 2880), 'h1st.schema.SchemaInferrer', 'SchemaInferrer', ([], {}), '()\n', (2878, 2880), False, 'from h1st.schema import SchemaInferrer\n'), ((3921, 3937), 'h1st.schema.SchemaInferrer', 'SchemaInferrer', ([], {}), '()\n', (3935, 3937), False, 'from h1st.schema import SchemaInferrer\n'), ((322, 332), 'pyarrow.int64', 'pa.int64', ([], {}), '()\n', (330, 332), True, 'import pyarrow as pa\n'), ((387, 399), 'pyarrow.float64', 'pa.float64', ([], {}), '()\n', (397, 399), True, 'import pyarrow as pa\n'), ((904, 914), 'pyarrow.int64', 'pa.int64', ([], {}), '()\n', (912, 914), True, 'import pyarrow as pa\n'), ((1014, 1026), 'pyarrow.float64', 'pa.float64', ([], {}), '()\n', (1024, 1026), True, 'import pyarrow as pa\n'), ((1081, 1100), 'pyarrow.array', 'pa.array', (['[1, 2, 3]'], {}), '([1, 2, 3])\n', (1089, 1100), True, 'import pyarrow as pa\n'), ((1102, 1127), 'pyarrow.array', 'pa.array', (["['a', 'b', 'c']"], {}), "(['a', 'b', 'c'])\n", (1110, 1127), True, 'import pyarrow as pa\n'), ((1350, 1381), 'numpy.random.random', 'np.random.random', (['(100, 28, 28)'], {}), '((100, 28, 28))\n', (1366, 1381), True, 'import numpy as np\n'), ((1438, 1450), 'pyarrow.float64', 'pa.float64', ([], {}), '()\n', (1448, 1450), True, 'import pyarrow as pa\n'), ((1547, 1572), 'numpy.array', 'np.array', (["['1', '2', '3']"], {}), "(['1', '2', '3'])\n", (1555, 1572), True, 'import numpy as np\n'), ((1629, 1640), 'pyarrow.string', 'pa.string', ([], {}), '()\n', (1638, 1640), True, 'import pyarrow as pa\n'), ((2715, 2735), 'pandas.Series', 'pd.Series', (['[1, 2, 3]'], {}), '([1, 2, 3])\n', (2724, 2735), True, 'import pandas as pd\n'), ((2791, 2801), 'pyarrow.int64', 'pa.int64', ([], {}), '()\n', (2799, 2801), True, 'import pyarrow as pa\n'), ((557, 571), 'datetime.datetime.now', 'datetime.now', ([], {}), '()\n', (569, 571), False, 'from datetime import datetime\n'), ((662, 672), 'pyarrow.int64', 'pa.int64', ([], {}), '()\n', (670, 672), True, 'import pyarrow as pa\n'), ((699, 710), 'pyarrow.string', 'pa.string', ([], {}), '()\n', (708, 710), True, 'import pyarrow as pa\n'), ((737, 748), 'pyarrow.binary', 'pa.binary', ([], {}), '()\n', (746, 748), True, 'import pyarrow as pa\n'), ((775, 786), 'pyarrow.date64', 'pa.date64', ([], {}), '()\n', (784, 786), True, 'import pyarrow as pa\n'), ((1997, 2007), 'pyarrow.int64', 'pa.int64', ([], {}), '()\n', (2005, 2007), True, 'import pyarrow as pa\n'), ((2031, 2042), 'pyarrow.string', 'pa.string', ([], {}), '()\n', (2040, 2042), True, 'import pyarrow as pa\n'), ((2066, 2078), 'pyarrow.float64', 'pa.float64', ([], {}), '()\n', (2076, 2078), True, 'import pyarrow as pa\n'), ((2476, 2488), 'pyarrow.float64', 'pa.float64', ([], {}), '()\n', (2486, 2488), True, 'import pyarrow as pa\n'), ((2518, 2530), 'pyarrow.float64', 'pa.float64', ([], {}), '()\n', (2528, 2530), True, 'import pyarrow as pa\n'), ((2554, 2566), 'pyarrow.float64', 'pa.float64', ([], {}), '()\n', (2564, 2566), True, 'import pyarrow as pa\n'), ((2590, 2602), 'pyarrow.float64', 'pa.float64', ([], {}), '()\n', (2600, 2602), True, 'import pyarrow as pa\n'), ((2629, 2640), 'pyarrow.string', 'pa.string', ([], {}), '()\n', (2638, 2640), True, 'import pyarrow as pa\n'), ((3043, 3053), 'pyarrow.int64', 'pa.int64', ([], {}), '()\n', (3051, 3053), True, 'import pyarrow as pa\n'), ((3276, 3286), 'pyarrow.int64', 'pa.int64', ([], {}), '()\n', (3284, 3286), True, 'import pyarrow as pa\n'), ((3434, 3497), 'pandas.DataFrame', 'pd.DataFrame', (["{'CarSpeed': [0, 1, 2], 'Label': ['a', 'b', 'c']}"], {}), "({'CarSpeed': [0, 1, 2], 'Label': ['a', 'b', 'c']})\n", (3446, 3497), True, 'import pandas as pd\n'), ((3324, 3334), 'pyarrow.int64', 'pa.int64', ([], {}), '()\n', (3332, 3334), True, 'import pyarrow as pa\n'), ((4189, 4199), 'pyarrow.int64', 'pa.int64', ([], {}), '()\n', (4197, 4199), True, 'import pyarrow as pa\n'), ((3747, 3757), 'pyarrow.int64', 'pa.int64', ([], {}), '()\n', (3755, 3757), True, 'import pyarrow as pa\n'), ((3792, 3803), 'pyarrow.string', 'pa.string', ([], {}), '()\n', (3801, 3803), True, 'import pyarrow as pa\n')]