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, 'h1st.schema.SchemaInferrer', 'SchemaInferrer', ([], {}), False, 'from h1st.schema import SchemaInferrer\n'), (46, 'h1st.schema.SchemaInferrer', 'SchemaInferrer', ([], {}), False, 'from h1st.schema import SchemaInferrer\n'), (59, 'h1st.schema.SchemaInferrer', 'SchemaInferrer', ([], {}), False, 'from h1st.schema import SchemaInferrer\n'), (60, 'pandas.DataFrame', 'pd.DataFrame', (["{'f1': [1, 2, 3], 'f2': ['a', 'b', 'c'], 'f3': [0.1, 0.2, 0.9]}"], {}), True, 'import pandas as pd\n'), (75, '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']}"], {}), True, 'import pandas as pd\n'), (100, 'h1st.schema.SchemaInferrer', 'SchemaInferrer', ([], {}), False, 'from h1st.schema import SchemaInferrer\n'), (140, 'h1st.schema.SchemaInferrer', 'SchemaInferrer', ([], {}), False, 'from h1st.schema import SchemaInferrer\n'), (13, 'pyarrow.int64', 'pa.int64', ([], {}), True, 'import pyarrow as pa\n'), (14, 'pyarrow.float64', 'pa.float64', ([], {}), True, 'import pyarrow as pa\n'), (33, 'pyarrow.int64', 'pa.int64', ([], {}), True, 'import pyarrow as pa\n'), (37, 'pyarrow.float64', 'pa.float64', ([], {}), True, 'import pyarrow as pa\n'), (40, 'pyarrow.array', 'pa.array', (['[1, 2, 3]'], {}), True, 'import pyarrow as pa\n'), (40, 'pyarrow.array', 'pa.array', (["['a', 'b', 'c']"], {}), True, 'import pyarrow as pa\n'), (47, 'numpy.random.random', 'np.random.random', (['(100, 28, 28)'], {}), True, 'import numpy as np\n'), (49, 'pyarrow.float64', 'pa.float64', ([], {}), True, 'import pyarrow as pa\n'), (53, 'numpy.array', 'np.array', (["['1', '2', '3']"], {}), True, 'import numpy as np\n'), (55, 'pyarrow.string', 'pa.string', ([], {}), True, 'import pyarrow as pa\n'), (94, 'pandas.Series', 'pd.Series', (['[1, 2, 3]'], {}), True, 'import pandas as pd\n'), (96, 'pyarrow.int64', 'pa.int64', ([], {}), True, 'import pyarrow as pa\n'), (20, 'datetime.datetime.now', 'datetime.now', ([], {}), False, 'from datetime import datetime\n'), (24, 'pyarrow.int64', 'pa.int64', ([], {}), True, 'import pyarrow as pa\n'), (25, 'pyarrow.string', 'pa.string', ([], {}), True, 'import pyarrow as pa\n'), (26, 'pyarrow.binary', 'pa.binary', ([], {}), True, 'import pyarrow as pa\n'), (27, 'pyarrow.date64', 'pa.date64', ([], {}), True, 'import pyarrow as pa\n'), (69, 'pyarrow.int64', 'pa.int64', ([], {}), True, 'import pyarrow as pa\n'), (70, 'pyarrow.string', 'pa.string', ([], {}), True, 'import pyarrow as pa\n'), (71, 'pyarrow.float64', 'pa.float64', ([], {}), True, 'import pyarrow as pa\n'), (86, 'pyarrow.float64', 'pa.float64', ([], {}), True, 'import pyarrow as pa\n'), (87, 'pyarrow.float64', 'pa.float64', ([], {}), True, 'import pyarrow as pa\n'), (88, 'pyarrow.float64', 'pa.float64', ([], {}), True, 'import pyarrow as pa\n'), (89, 'pyarrow.float64', 'pa.float64', ([], {}), True, 'import pyarrow as pa\n'), (90, 'pyarrow.string', 'pa.string', ([], {}), True, 'import pyarrow as pa\n'), (106, 'pyarrow.int64', 'pa.int64', ([], {}), True, 'import pyarrow as pa\n'), (116, 'pyarrow.int64', 'pa.int64', ([], {}), True, 'import pyarrow as pa\n'), (122, 'pandas.DataFrame', 'pd.DataFrame', (["{'CarSpeed': [0, 1, 2], 'Label': ['a', 'b', 'c']}"], {}), True, 'import pandas as pd\n'), (117, 'pyarrow.int64', 'pa.int64', ([], {}), True, 'import pyarrow as pa\n'), (149, 'pyarrow.int64', 'pa.int64', ([], {}), True, 'import pyarrow as pa\n'), (132, 'pyarrow.int64', 'pa.int64', ([], {}), True, 'import pyarrow as pa\n'), (133, 'pyarrow.string', 'pa.string', ([], {}), True, 'import pyarrow as pa\n')]