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()
                }
            }
        })

[((11, 19, 11, 35), 'h1st.schema.SchemaInferrer', 'SchemaInferrer', ({}, {}), '()', False, 'from h1st.schema import SchemaInferrer\n'), ((46, 19, 46, 35), 'h1st.schema.SchemaInferrer', 'SchemaInferrer', ({}, {}), '()', False, 'from h1st.schema import SchemaInferrer\n'), ((59, 19, 59, 35), 'h1st.schema.SchemaInferrer', 'SchemaInferrer', ({}, {}), '()', False, 'from h1st.schema import SchemaInferrer\n'), ((60, 13, 64, 10), 'pandas.DataFrame', 'pd.DataFrame', ({(60, 26, 64, 9): "{'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]})", True, 'import pandas as pd\n'), ((75, 13, 81, 10), 'pandas.DataFrame', 'pd.DataFrame', ({(75, 26, 81, 9): "{'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']})", True, 'import pandas as pd\n'), ((100, 19, 100, 35), 'h1st.schema.SchemaInferrer', 'SchemaInferrer', ({}, {}), '()', False, 'from h1st.schema import SchemaInferrer\n'), ((140, 19, 140, 35), 'h1st.schema.SchemaInferrer', 'SchemaInferrer', ({}, {}), '()', False, 'from h1st.schema import SchemaInferrer\n'), ((13, 51, 13, 61), 'pyarrow.int64', 'pa.int64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((14, 53, 14, 65), 'pyarrow.float64', 'pa.float64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((33, 21, 33, 31), 'pyarrow.int64', 'pa.int64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((37, 21, 37, 33), 'pyarrow.float64', 'pa.float64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((40, 13, 40, 32), 'pyarrow.array', 'pa.array', ({(40, 22, 40, 31): '[1, 2, 3]'}, {}), '([1, 2, 3])', True, 'import pyarrow as pa\n'), ((40, 34, 40, 59), 'pyarrow.array', 'pa.array', ({(40, 43, 40, 58): "['a', 'b', 'c']"}, {}), "(['a', 'b', 'c'])", True, 'import pyarrow as pa\n'), ((47, 47, 47, 78), 'numpy.random.random', 'np.random.random', ({(47, 64, 47, 77): '(100, 28, 28)'}, {}), '((100, 28, 28))', True, 'import numpy as np\n'), ((49, 20, 49, 32), 'pyarrow.float64', 'pa.float64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((53, 47, 53, 72), 'numpy.array', 'np.array', ({(53, 56, 53, 71): "['1', '2', '3']"}, {}), "(['1', '2', '3'])", True, 'import numpy as np\n'), ((55, 20, 55, 31), 'pyarrow.string', 'pa.string', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((94, 47, 94, 67), 'pandas.Series', 'pd.Series', ({(94, 57, 94, 66): '[1, 2, 3]'}, {}), '([1, 2, 3])', True, 'import pandas as pd\n'), ((96, 20, 96, 30), 'pyarrow.int64', 'pa.int64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((20, 21, 20, 35), 'datetime.datetime.now', 'datetime.now', ({}, {}), '()', False, 'from datetime import datetime\n'), ((24, 25, 24, 35), 'pyarrow.int64', 'pa.int64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((25, 25, 25, 36), 'pyarrow.string', 'pa.string', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((26, 25, 26, 36), 'pyarrow.binary', 'pa.binary', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((27, 25, 27, 36), 'pyarrow.date64', 'pa.date64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((69, 22, 69, 32), 'pyarrow.int64', 'pa.int64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((70, 22, 70, 33), 'pyarrow.string', 'pa.string', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((71, 22, 71, 34), 'pyarrow.float64', 'pa.float64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((86, 29, 86, 41), 'pyarrow.float64', 'pa.float64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((87, 28, 87, 40), 'pyarrow.float64', 'pa.float64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((88, 22, 88, 34), 'pyarrow.float64', 'pa.float64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((89, 22, 89, 34), 'pyarrow.float64', 'pa.float64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((90, 25, 90, 36), 'pyarrow.string', 'pa.string', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((106, 24, 106, 34), 'pyarrow.int64', 'pa.int64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((116, 24, 116, 34), 'pyarrow.int64', 'pa.int64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((122, 23, 125, 14), 'pandas.DataFrame', 'pd.DataFrame', ({(122, 36, 125, 13): "{'CarSpeed': [0, 1, 2], 'Label': ['a', 'b', 'c']}"}, {}), "({'CarSpeed': [0, 1, 2], 'Label': ['a', 'b', 'c']})", True, 'import pandas as pd\n'), ((117, 36, 117, 46), 'pyarrow.int64', 'pa.int64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((149, 28, 149, 38), 'pyarrow.int64', 'pa.int64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((132, 36, 132, 46), 'pyarrow.int64', 'pa.int64', ({}, {}), '()', True, 'import pyarrow as pa\n'), ((133, 33, 133, 44), 'pyarrow.string', 'pa.string', ({}, {}), '()', True, 'import pyarrow as pa\n')]