# 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 os
import shutil
import tempfile
import unittest

import pandas as pd
import pyarrow as pa

from pypaimon import CatalogFactory, Schema
from pypaimon.tests.predicates_test import (_check_filtered_result,
                                            _random_format)


class AOPredicatePy36Test(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.tempdir = tempfile.mkdtemp()
        cls.warehouse = os.path.join(cls.tempdir, 'warehouse')
        cls.catalog = CatalogFactory.create({'warehouse': cls.warehouse})
        cls.catalog.create_database('default', False)
        pa_schema = pa.schema([
            ('f0', pa.int64()),
            ('f1', pa.string()),
        ])
        cls.catalog.create_table('default.test_append', Schema.from_pyarrow_schema(
            pa_schema, options={'file.format': _random_format(), 'metadata.stats-mode': 'full'}), False)

        df = pd.DataFrame({
            'f0': [1, 2, 3, 4, 5],
            'f1': ['abc', 'abbc', 'bc', 'd', None],
        })

        append_table = cls.catalog.get_table('default.test_append')
        write_builder = append_table.new_batch_write_builder()
        write = write_builder.new_write()
        commit = write_builder.new_commit()
        write.write_pandas(df)
        commit.commit(write.prepare_commit())
        write.close()
        commit.close()

        cls.df = df

    @classmethod
    def tearDownClass(cls):
        shutil.rmtree(cls.tempdir, ignore_errors=True)

    def test_wrong_field_name(self):
        table = self.catalog.get_table('default.test_append')
        predicate_builder = table.new_read_builder().new_predicate_builder()
        with self.assertRaises(ValueError) as e:
            predicate_builder.equal('f2', 'a')
        self.assertEqual(str(e.exception), "The field f2 is not in field list ['f0', 'f1'].")

    def test_append_with_duplicate(self):
        pa_schema = pa.schema([
            ('f0', pa.int64()),
            ('f1', pa.string()),
        ])
        self.catalog.create_table('default.test_append_with_duplicate', Schema.from_pyarrow_schema(pa_schema), False)

        df = pd.DataFrame({
            'f0': [1, 1, 2, 2],
            'f1': ['a', 'b', 'c', 'd'],
        })

        table = self.catalog.get_table('default.test_append_with_duplicate')
        write_builder = table.new_batch_write_builder()
        write = write_builder.new_write()
        commit = write_builder.new_commit()
        write.write_pandas(df)
        commit.commit(write.prepare_commit())
        write.close()
        commit.close()

        predicate_builder = table.new_read_builder().new_predicate_builder()

        predicate = predicate_builder.equal('f0', 1)
        _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[0:1])

        predicate = predicate_builder.equal('f0', 0)
        read_builder = table.new_read_builder().with_filter(predicate)
        scan = read_builder.new_scan()
        read = read_builder.new_read()
        actual_df = read.to_pandas(scan.plan().splits())
        self.assertEqual(len(actual_df), 0)

    def test_all_field_types_with_equal(self):
        pa_schema = pa.schema([
            # int
            ('_tinyint', pa.int8()),
            ('_smallint', pa.int16()),
            ('_int', pa.int32()),
            ('_bigint', pa.int64()),
            # float
            ('_float16', pa.float32()),  # NOTE: cannot write pa.float16() data into Paimon
            ('_float32', pa.float32()),
            ('_double', pa.float64()),
            # string
            ('_string', pa.string()),
            # bool
            ('_boolean', pa.bool_())
        ])
        self.catalog.create_table('default.test_all_field_types',
                                  Schema.from_pyarrow_schema(pa_schema, options={'file.format': _random_format()}),
                                  False)
        table = self.catalog.get_table('default.test_all_field_types')
        write_builder = table.new_batch_write_builder()
        write = write_builder.new_write()
        commit = write_builder.new_commit()

        df = pd.DataFrame({
            '_tinyint': pd.Series([1, 2], dtype='int8'),
            '_smallint': pd.Series([10, 20], dtype='int16'),
            '_int': pd.Series([100, 200], dtype='int32'),
            '_bigint': pd.Series([1000, 2000], dtype='int64'),
            '_float16': pd.Series([1.0, 2.0], dtype='float32'),
            '_float32': pd.Series([1.00, 2.00], dtype='float32'),
            '_double': pd.Series([1.000, 2.000], dtype='double'),
            '_string': pd.Series(['A', 'B'], dtype='object'),
            '_boolean': [True, False]
        })
        record_batch = pa.RecordBatch.from_pandas(df, schema=pa_schema)

        write.write_arrow_batch(record_batch)
        commit.commit(write.prepare_commit())
        write.close()
        commit.close()

        predicate_builder = table.new_read_builder().new_predicate_builder()

        predicate = predicate_builder.equal('_tinyint', 1)
        _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[0]])

        predicate = predicate_builder.equal('_smallint', 20)
        _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[1]])

        predicate = predicate_builder.equal('_int', 100)
        _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[0]])

        predicate = predicate_builder.equal('_bigint', 2000)
        _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[1]])

        predicate = predicate_builder.equal('_float16', 1.0)
        _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[0]])

        predicate = predicate_builder.equal('_float32', 2.00)
        _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[1]])

        predicate = predicate_builder.equal('_double', 1.000)
        _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[0]])

        predicate = predicate_builder.equal('_string', 'B')
        _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[1]])

        predicate = predicate_builder.equal('_boolean', True)
        _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[0]])

    def test_not_equal_append(self):
        table = self.catalog.get_table('default.test_append')
        predicate_builder = table.new_read_builder().new_predicate_builder()
        predicate = predicate_builder.not_equal('f0', 1)
        _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[1:4])

    def test_less_than_append(self):
        table = self.catalog.get_table('default.test_append')
        predicate_builder = table.new_read_builder().new_predicate_builder()
        predicate = predicate_builder.less_than('f0', 3)
        _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:1])

    def test_less_or_equal_append(self):
        table = self.catalog.get_table('default.test_append')
        predicate_builder = table.new_read_builder().new_predicate_builder()
        predicate = predicate_builder.less_or_equal('f0', 3)
        _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:2])

    def test_greater_than_append(self):
        table = self.catalog.get_table('default.test_append')
        predicate_builder = table.new_read_builder().new_predicate_builder()
        predicate = predicate_builder.greater_than('f0', 3)
        _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[3:4])

    def test_greater_or_equal_append(self):
        table = self.catalog.get_table('default.test_append')
        predicate_builder = table.new_read_builder().new_predicate_builder()
        predicate = predicate_builder.greater_or_equal('f0', 3)
        _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[2:4])

    def test_is_null_append(self):
        table = self.catalog.get_table('default.test_append')
        predicate_builder = table.new_read_builder().new_predicate_builder()
        predicate = predicate_builder.is_null('f1')
        _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[[4]])

    def test_is_not_null_append(self):
        table = self.catalog.get_table('default.test_append')
        predicate_builder = table.new_read_builder().new_predicate_builder()
        predicate = predicate_builder.is_not_null('f1')
        _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:3])

    def test_is_in_append(self):
        table = self.catalog.get_table('default.test_append')
        predicate_builder = table.new_read_builder().new_predicate_builder()
        predicate = predicate_builder.is_in('f0', [1, 2])
        _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:1])

    def test_is_not_in_append(self):
        table = self.catalog.get_table('default.test_append')
        predicate_builder = table.new_read_builder().new_predicate_builder()
        predicate = predicate_builder.is_not_in('f0', [1, 2])
        _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[2:4])

    def test_between_append(self):
        table = self.catalog.get_table('default.test_append')
        predicate_builder = table.new_read_builder().new_predicate_builder()
        predicate = predicate_builder.between('f0', 1, 3)
        _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:2])

    def test_and_predicates(self):
        table = self.catalog.get_table('default.test_append')
        predicate_builder = table.new_read_builder().new_predicate_builder()
        predicate1 = predicate_builder.greater_than('f0', 1)
        predicate2 = predicate_builder.less_than('f0', 4)
        predicate = predicate_builder.and_predicates([predicate1, predicate2])
        _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[1:2])

    def test_or_predicates(self):
        table = self.catalog.get_table('default.test_append')
        predicate_builder = table.new_read_builder().new_predicate_builder()
        predicate1 = predicate_builder.greater_than('f0', 3)
        predicate2 = predicate_builder.less_than('f0', 2)
        predicate = predicate_builder.or_predicates([predicate1, predicate2])
        _check_filtered_result(table.new_read_builder().with_filter(predicate),
                               self.df.loc[[0, 3, 4]])
