################################################################################
#  Licensed to the Apache Software Foundation (ASF) under one
#  or more contributor license agreements.  See the NOTICE file
#  distributed with this work for additional information
#  regarding copyright ownership.  The ASF licenses this file
#  to you under the Apache License, Version 2.0 (the
#  "License"); you may not use this file except in compliance
#  with the License.  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
# limitations under the License.
################################################################################
import datetime
import decimal

from pandas.testing import assert_frame_equal

from pyflink.common import Row
from pyflink.table.types import DataTypes
from pyflink.testing import source_sink_utils
from pyflink.testing.test_case_utils import PyFlinkBatchTableTestCase, \
    PyFlinkStreamTableTestCase


class PandasConversionTestBase(object):

    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.data = [(1, 1, 1, 1, True, 1.1, 1.2, 'hello', bytearray(b"aaa"),
                     decimal.Decimal('1000000000000000000.01'), datetime.date(2014, 9, 13),
                     datetime.time(hour=1, minute=0, second=1),
                     datetime.datetime(1970, 1, 1, 0, 0, 0, 123000), ['hello', '中文'],
                     Row(a=1, b='hello', c=datetime.datetime(1970, 1, 1, 0, 0, 0, 123000),
                         d=[1, 2])),
                    (1, 2, 2, 2, False, 2.1, 2.2, 'world', bytearray(b"bbb"),
                     decimal.Decimal('1000000000000000000.02'), datetime.date(2014, 9, 13),
                     datetime.time(hour=1, minute=0, second=1),
                     datetime.datetime(1970, 1, 1, 0, 0, 0, 123000), ['hello', '中文'],
                     Row(a=1, b='hello', c=datetime.datetime(1970, 1, 1, 0, 0, 0, 123000),
                         d=[1, 2]))]
        cls.data_type = DataTypes.ROW(
            [DataTypes.FIELD("f1", DataTypes.TINYINT()),
             DataTypes.FIELD("f2", DataTypes.SMALLINT()),
             DataTypes.FIELD("f3", DataTypes.INT()),
             DataTypes.FIELD("f4", DataTypes.BIGINT()),
             DataTypes.FIELD("f5", DataTypes.BOOLEAN()),
             DataTypes.FIELD("f6", DataTypes.FLOAT()),
             DataTypes.FIELD("f7", DataTypes.DOUBLE()),
             DataTypes.FIELD("f8", DataTypes.STRING()),
             DataTypes.FIELD("f9", DataTypes.BYTES()),
             DataTypes.FIELD("f10", DataTypes.DECIMAL(38, 18)),
             DataTypes.FIELD("f11", DataTypes.DATE()),
             DataTypes.FIELD("f12", DataTypes.TIME()),
             DataTypes.FIELD("f13", DataTypes.TIMESTAMP(3)),
             DataTypes.FIELD("f14", DataTypes.ARRAY(DataTypes.STRING())),
             DataTypes.FIELD("f15", DataTypes.ROW(
                 [DataTypes.FIELD("a", DataTypes.INT()),
                  DataTypes.FIELD("b", DataTypes.STRING()),
                  DataTypes.FIELD("c", DataTypes.TIMESTAMP(3)),
                  DataTypes.FIELD("d", DataTypes.ARRAY(DataTypes.INT()))]))], False)
        cls.pdf = cls.create_pandas_data_frame()

    @classmethod
    def create_pandas_data_frame(cls):
        data_dict = {}
        for j, name in enumerate(cls.data_type.names):
            data_dict[name] = [cls.data[i][j] for i in range(len(cls.data))]
        # need convert to numpy types
        import numpy as np
        data_dict["f1"] = np.int8(data_dict["f1"])
        data_dict["f2"] = np.int16(data_dict["f2"])
        data_dict["f3"] = np.int32(data_dict["f3"])
        data_dict["f4"] = np.int64(data_dict["f4"])
        data_dict["f6"] = np.float32(data_dict["f6"])
        data_dict["f7"] = np.float64(data_dict["f7"])
        data_dict["f15"] = [row.as_dict() for row in data_dict["f15"]]
        import pandas as pd
        return pd.DataFrame(data=data_dict,
                            index=[2., 3.],
                            columns=['f1', 'f2', 'f3', 'f4', 'f5', 'f6', 'f7', 'f8', 'f9',
                                     'f10', 'f11', 'f12', 'f13', 'f14', 'f15'])


class PandasConversionTests(PandasConversionTestBase):

    def test_from_pandas_with_incorrect_schema(self):
        fields = self.data_type.fields.copy()
        fields[0], fields[7] = fields[7], fields[0]  # swap str with tinyint
        wrong_schema = DataTypes.ROW(fields)  # should be DataTypes.STRING()
        with self.assertRaisesRegex(Exception, "Expected a string.*got int8"):
            self.t_env.from_pandas(self.pdf, schema=wrong_schema)

    def test_from_pandas_with_names(self):
        # skip decimal as currently only decimal(38, 18) is supported
        pdf = self.pdf.drop(['f10', 'f11', 'f12', 'f13', 'f14', 'f15'], axis=1)
        new_names = list(map(str, range(len(pdf.columns))))
        table = self.t_env.from_pandas(pdf, schema=new_names)
        self.assertEqual(new_names, table.get_schema().get_field_names())
        table = self.t_env.from_pandas(pdf, schema=tuple(new_names))
        self.assertEqual(new_names, table.get_schema().get_field_names())

    def test_from_pandas_with_types(self):
        new_types = self.data_type.field_types()
        new_types[0] = DataTypes.BIGINT()
        table = self.t_env.from_pandas(self.pdf, schema=new_types)
        self.assertEqual(new_types, table.get_schema().get_field_data_types())
        table = self.t_env.from_pandas(self.pdf, schema=tuple(new_types))
        self.assertEqual(new_types, table.get_schema().get_field_data_types())


class PandasConversionITTests(PandasConversionTestBase):

    def test_from_pandas(self):
        table = self.t_env.from_pandas(self.pdf, self.data_type, 5)
        self.assertEqual(self.data_type, table.get_schema().to_row_data_type())

        table = table.filter(table.f2 < 2)
        sink_table_ddl = """
            CREATE TABLE Results(
            f1 TINYINT,
            f2 SMALLINT,
            f3 INT,
            f4 BIGINT,
            f5 BOOLEAN,
            f6 FLOAT,
            f7 DOUBLE,
            f8 STRING,
            f9 BYTES,
            f10 DECIMAL(38, 18),
            f11 DATE,
            f12 TIME,
            f13 TIMESTAMP(3),
            f14 ARRAY<STRING>,
            f15 ROW<a INT, b STRING, c TIMESTAMP(3), d ARRAY<INT>>)
            WITH ('connector'='test-sink')
        """
        self.t_env.execute_sql(sink_table_ddl)

        table.execute_insert("Results").wait()
        actual = source_sink_utils.results()
        self.assert_equals(actual,
                           ["+I[1, 1, 1, 1, true, 1.1, 1.2, hello, [97, 97, 97], "
                            "1000000000000000000.010000000000000000, 2014-09-13, 01:00:01, "
                            "1970-01-01T00:00:00.123, [hello, 中文], +I[1, hello, "
                            "1970-01-01T00:00:00.123, [1, 2]]]"])

    def test_to_pandas(self):
        table = self.t_env.from_pandas(self.pdf, self.data_type)
        result_pdf = table.to_pandas()
        result_pdf.index = self.pdf.index
        self.assertEqual(2, len(result_pdf))
        expected_arrow = self.pdf.to_records(index=False)
        result_arrow = result_pdf.to_records(index=False)
        for r in range(len(expected_arrow)):
            for e in range(len(expected_arrow[r])):
                self.assert_equal_field(expected_arrow[r][e], result_arrow[r][e])

    def test_empty_to_pandas(self):
        table = self.t_env.from_pandas(self.pdf, self.data_type)
        pdf = table.filter(table.f1 < 0).to_pandas()
        self.assertTrue(pdf.empty)

    def test_to_pandas_for_retract_table(self):
        table = self.t_env.from_pandas(self.pdf, self.data_type)
        result_pdf = table.group_by(table.f1).select(table.f2.max.alias('f2')).to_pandas()
        import pandas as pd
        import numpy as np
        assert_frame_equal(result_pdf, pd.DataFrame(data={'f2': np.int16([2])}))

        result_pdf = table.group_by(table.f2).select(table.f1.max.alias('f2')).to_pandas()
        assert_frame_equal(result_pdf, pd.DataFrame(data={'f2': np.int8([1, 1])}))

    def assert_equal_field(self, expected_field, result_field):
        import numpy as np
        result_type = type(result_field)
        if result_type == dict:
            self.assertEqual(expected_field.keys(), result_field.keys())
            for key in expected_field:
                self.assert_equal_field(expected_field[key], result_field[key])
        elif result_type == np.ndarray:
            self.assertTrue((expected_field == result_field).all())
        else:
            self.assertTrue(expected_field == result_field)


class BatchPandasConversionTests(PandasConversionTests,
                                 PandasConversionITTests,
                                 PyFlinkBatchTableTestCase):
    pass


class StreamPandasConversionTests(PandasConversionITTests,
                                  PyFlinkStreamTableTestCase):
    def test_to_pandas_with_event_time(self):
        self.t_env.get_config().set("parallelism.default", "1")
        # create source file path
        import tempfile
        import os
        tmp_dir = tempfile.gettempdir()
        data = [
            '2018-03-11 03:10:00',
            '2018-03-11 03:10:00',
            '2018-03-11 03:10:00',
            '2018-03-11 03:40:00',
            '2018-03-11 04:20:00',
            '2018-03-11 03:30:00'
        ]
        source_path = tmp_dir + '/test_to_pandas_with_event_time.csv'
        with open(source_path, 'w') as fd:
            for ele in data:
                fd.write(ele + '\n')

        self.t_env.get_config().set(
            "pipeline.time-characteristic", "EventTime")

        source_table = """
            create table source_table(
                rowtime TIMESTAMP(3),
                WATERMARK FOR rowtime AS rowtime - INTERVAL '60' MINUTE
            ) with(
                'connector.type' = 'filesystem',
                'format.type' = 'csv',
                'connector.path' = '%s',
                'format.ignore-first-line' = 'false',
                'format.field-delimiter' = ','
            )
        """ % source_path
        self.t_env.execute_sql(source_table)
        t = self.t_env.from_path("source_table")
        result_pdf = t.to_pandas()
        import pandas as pd
        os.remove(source_path)
        expected_df = pd.DataFrame(data={"rowtime": pd.Series(data, dtype="datetime64[ms]")})
        assert_frame_equal(result_pdf, expected_df)