#
# 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.
#

"""Tests for schemas."""

# pytype: skip-file

import typing
import unittest

import numpy as np
import pandas as pd
from parameterized import parameterized

import apache_beam as beam
from apache_beam.coders import RowCoder
from apache_beam.coders.typecoders import registry as coders_registry
from apache_beam.dataframe import schemas
from apache_beam.dataframe import transforms
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
from apache_beam.typehints import row_type
from apache_beam.typehints import typehints
from apache_beam.typehints.native_type_compatibility import match_is_named_tuple

Simple = typing.NamedTuple(
    'Simple', [('name', str), ('id', int), ('height', float)])
coders_registry.register_coder(Simple, RowCoder)
Animal = typing.NamedTuple(
    'Animal', [('animal', str), ('max_speed', typing.Optional[float])])
coders_registry.register_coder(Animal, RowCoder)


def matches_df(expected):
  def check_df_pcoll_equal(actual):
    actual = pd.concat(actual)
    sorted_actual = actual.sort_values(by=list(actual.columns)).reset_index(
        drop=True)
    sorted_expected = expected.sort_values(
        by=list(expected.columns)).reset_index(drop=True)
    pd.testing.assert_frame_equal(sorted_actual, sorted_expected)

  return check_df_pcoll_equal


# Test data for all supported types that can be easily tested.
# Excludes bytes because it's difficult to create a series and dataframe bytes
# dtype. For example:
#   pd.Series([b'abc'], dtype=bytes).dtype != 'S'
#   pd.Series([b'abc'], dtype=bytes).astype(bytes).dtype == 'S'
# (test data, pandas_type, column_name, beam_type)
COLUMNS: typing.List[typing.Tuple[typing.List[typing.Any],
                                  typing.Any,
                                  str,
                                  typing.Any]] = [
                                      ([375, 24, 0, 10, 16],
                                       np.int32,
                                       'i32',
                                       np.int32),
                                      ([375, 24, 0, 10, 16],
                                       np.int64,
                                       'i64',
                                       np.int64),
                                      ([375, 24, None, 10, 16],
                                       pd.Int32Dtype(),
                                       'i32_nullable',
                                       typing.Optional[np.int32]),
                                      ([375, 24, None, 10, 16],
                                       pd.Int64Dtype(),
                                       'i64_nullable',
                                       typing.Optional[np.int64]),
                                      ([375., 24., None, 10., 16.],
                                       np.float64,
                                       'f64',
                                       typing.Optional[np.float64]),
                                      ([375., 24., None, 10., 16.],
                                       np.float32,
                                       'f32',
                                       typing.Optional[np.float32]),
                                      ([True, False, True, True, False],
                                       bool,
                                       'bool',
                                       bool),
                                      (['Falcon', 'Ostrich', None, 3.14, 0],
                                       object,
                                       'any',
                                       typing.Any),
                                      ([True, False, True, None, False],
                                       pd.BooleanDtype(),
                                       'bool_nullable',
                                       typing.Optional[bool]),
                                      ([
                                          'Falcon',
                                          'Ostrich',
                                          None,
                                          'Aardvark',
                                          'Elephant'
                                      ],
                                       pd.StringDtype(),
                                       'strdtype',
                                       typing.Optional[str]),
                                  ]

NICE_TYPES_DF = pd.DataFrame(columns=[name for _, _, name, _ in COLUMNS])
for arr, dtype, name, _ in COLUMNS:
  NICE_TYPES_DF[name] = pd.Series(arr, dtype=dtype, name=name).astype(dtype)

NICE_TYPES_PROXY = NICE_TYPES_DF[:0]

SERIES_TESTS = [(pd.Series(arr, dtype=dtype, name=name), arr, beam_type)
                for (arr, dtype, name, beam_type) in COLUMNS]

_TEST_ARRAYS: typing.List[typing.List[typing.Any]] = [
    arr for (arr, _, _, _) in COLUMNS
]
DF_RESULT = list(zip(*_TEST_ARRAYS))
BEAM_SCHEMA = typing.NamedTuple(  # type: ignore
    'BEAM_SCHEMA', [(name, beam_type) for _, _, name, beam_type in COLUMNS])
INDEX_DF_TESTS = [(
    NICE_TYPES_DF.set_index([name for _, _, name, _ in COLUMNS[:i]]),
    DF_RESULT,
    BEAM_SCHEMA) for i in range(1, len(COLUMNS) + 1)]

NOINDEX_DF_TESTS = [(NICE_TYPES_DF, DF_RESULT, BEAM_SCHEMA)]

# Get major, minor, bugfix version
PD_VERSION = tuple(map(int, pd.__version__.split('.')[0:3]))


def test_name_func(testcase_func, param_num, params):
  df_or_series, _, _ = params.args
  if isinstance(df_or_series, pd.Series):
    return f"{testcase_func.__name__}_Series[{df_or_series.dtype}]"
  elif isinstance(df_or_series, pd.DataFrame):
    return (
        f"{testcase_func.__name__}_DataFrame"
        f"[{','.join(str(dtype) for dtype in df_or_series.dtypes)}]")
  else:
    raise ValueError(
        f"Encountered unsupported param in {testcase_func.__name__}. "
        "Expected Series or DataFrame, got:\n" + str(df_or_series))


class SchemasTest(unittest.TestCase):
  def test_simple_df(self):
    expected = pd.DataFrame({
        'name': list(str(i) for i in range(5)),
        'id': list(range(5)),
        'height': list(float(i) for i in range(5))
    },
                            columns=['name', 'id', 'height'])

    expected.name = expected.name.astype(pd.StringDtype())

    with TestPipeline() as p:
      res = (
          p
          | beam.Create(
              [Simple(name=str(i), id=i, height=float(i)) for i in range(5)])
          | schemas.BatchRowsAsDataFrame(min_batch_size=10, max_batch_size=10))
      assert_that(res, matches_df(expected))

  def test_simple_df_with_beam_row(self):
    expected = pd.DataFrame({
        'name': list(str(i) for i in range(5)),
        'id': list(range(5)),
        'height': list(float(i) for i in range(5))
    },
                            columns=['name', 'id', 'height'])
    expected.name = expected.name.astype(pd.StringDtype())

    with TestPipeline() as p:
      res = (
          p
          | beam.Create([(str(i), i, float(i)) for i in range(5)])
          | beam.Select(
              name=lambda r: str(r[0]),
              id=lambda r: int(r[1]),
              height=lambda r: float(r[2]))
          | schemas.BatchRowsAsDataFrame(min_batch_size=10, max_batch_size=10))
      assert_that(res, matches_df(expected))

  def test_generate_proxy(self):
    expected = pd.DataFrame({
        'animal': pd.Series(dtype=pd.StringDtype()),
        'max_speed': pd.Series(dtype=np.float64)
    })

    pd.testing.assert_frame_equal(schemas.generate_proxy(Animal), expected)

  def test_generate_proxy_beam_typehint(self):
    expected = pd.Series(dtype=pd.Int32Dtype())

    actual = schemas.generate_proxy(typehints.Optional[np.int32])

    pd.testing.assert_series_equal(actual, expected)

  def test_nice_types_proxy_roundtrip(self):
    roundtripped = schemas.generate_proxy(
        schemas.element_type_from_dataframe(NICE_TYPES_PROXY))
    self.assertTrue(roundtripped.equals(NICE_TYPES_PROXY))

  @unittest.skipIf(
      PD_VERSION == (1, 2, 1),
      "Can't roundtrip bytes in pandas 1.2.1"
      "https://github.com/pandas-dev/pandas/issues/39474")
  def test_bytes_proxy_roundtrip(self):
    proxy = pd.DataFrame({'bytes': []})
    proxy.bytes = proxy.bytes.astype(bytes)

    roundtripped = schemas.generate_proxy(
        schemas.element_type_from_dataframe(proxy))

    self.assertEqual(roundtripped.bytes.dtype.kind, 'S')

  def test_batch_with_df_transform(self):
    with TestPipeline() as p:
      res = (
          p
          | beam.Create([
              Animal('Falcon', 380.0),
              Animal('Falcon', 370.0),
              Animal('Parrot', 24.0),
              Animal('Parrot', 26.0)
          ])
          | schemas.BatchRowsAsDataFrame()
          | transforms.DataframeTransform(
              lambda df: df.groupby('animal').mean(),
              # TODO: Generate proxy in this case as well
              proxy=schemas.generate_proxy(Animal),
              include_indexes=True))
      assert_that(res, equal_to([('Falcon', 375.), ('Parrot', 25.)]))

    # Do the same thing, but use reset_index() to make sure 'animal' is included
    with TestPipeline() as p:
      with beam.dataframe.allow_non_parallel_operations():
        res = (
            p
            | beam.Create([
                Animal('Falcon', 380.0),
                Animal('Falcon', 370.0),
                Animal('Parrot', 24.0),
                Animal('Parrot', 26.0)
            ])
            | schemas.BatchRowsAsDataFrame()
            | transforms.DataframeTransform(
                lambda df: df.groupby('animal').mean().reset_index(),
                # TODO: Generate proxy in this case as well
                proxy=schemas.generate_proxy(Animal)))
        assert_that(res, equal_to([('Falcon', 375.), ('Parrot', 25.)]))

  def assert_typehints_equal(self, left, right):
    def maybe_drop_rowtypeconstraint(typehint):
      if isinstance(typehint, row_type.RowTypeConstraint):
        return typehint.user_type
      else:
        return typehint

    left = maybe_drop_rowtypeconstraint(typehints.normalize(left))
    right = maybe_drop_rowtypeconstraint(typehints.normalize(right))

    if match_is_named_tuple(left):
      self.assertTrue(match_is_named_tuple(right))
      self.assertEqual(left.__annotations__, right.__annotations__)
    else:
      self.assertEqual(left, right)

  @parameterized.expand(
      SERIES_TESTS + NOINDEX_DF_TESTS, name_func=test_name_func)
  def test_unbatch_no_index(self, df_or_series, rows, beam_type):
    proxy = df_or_series[:0]

    with TestPipeline() as p:
      res = (
          p | beam.Create([df_or_series[::2], df_or_series[1::2]])
          | schemas.UnbatchPandas(proxy))

      # Verify that the unbatched PCollection has the expected typehint
      # TODO(https://github.com/apache/beam/issues/19923): typehints should
      # support NamedTuple so we can use typehints.is_consistent_with here
      # instead
      self.assert_typehints_equal(res.element_type, beam_type)

      assert_that(res, equal_to(rows))

  @parameterized.expand(SERIES_TESTS + INDEX_DF_TESTS, name_func=test_name_func)
  def test_unbatch_with_index(self, df_or_series, rows, _):
    proxy = df_or_series[:0]

    if (PD_VERSION < (1, 2) and
        set(['i32_nullable', 'i64_nullable']).intersection(proxy.index.names)):
      self.skipTest(
          "pandas<1.2 incorrectly changes Int64Dtype to int64 when "
          "moved to index.")

    with TestPipeline() as p:
      res = (
          p | beam.Create([df_or_series[::2], df_or_series[1::2]])
          | schemas.UnbatchPandas(proxy, include_indexes=True))

      assert_that(res, equal_to(rows))

  @parameterized.expand(SERIES_TESTS, name_func=test_name_func)
  def test_unbatch_series_with_index_warns(
      self, series, unused_rows, unused_type):
    proxy = series[:0]

    with TestPipeline() as p:
      input_pc = p | beam.Create([series[::2], series[1::2]])
      with self.assertWarns(UserWarning):
        _ = input_pc | schemas.UnbatchPandas(proxy, include_indexes=True)

  def test_unbatch_include_index_unnamed_index_raises(self):
    df = pd.DataFrame({'foo': [1, 2, 3, 4]})
    proxy = df[:0]

    with TestPipeline() as p:
      pc = p | beam.Create([df[::2], df[1::2]])

      with self.assertRaisesRegex(ValueError, 'unnamed'):
        _ = pc | schemas.UnbatchPandas(proxy, include_indexes=True)

  def test_unbatch_include_index_nonunique_index_raises(self):
    df = pd.DataFrame({'foo': [1, 2, 3, 4]})
    df.index = pd.MultiIndex.from_arrays([[1, 2, 3, 4], [4, 3, 2, 1]],
                                         names=['bar', 'bar'])
    proxy = df[:0]

    with TestPipeline() as p:
      pc = p | beam.Create([df[::2], df[1::2]])

      with self.assertRaisesRegex(ValueError, 'bar'):
        _ = pc | schemas.UnbatchPandas(proxy, include_indexes=True)

  def test_unbatch_include_index_column_conflict_raises(self):
    df = pd.DataFrame({'foo': [1, 2, 3, 4]})
    df.index = pd.Index([4, 3, 2, 1], name='foo')
    proxy = df[:0]

    with TestPipeline() as p:
      pc = p | beam.Create([df[::2], df[1::2]])

      with self.assertRaisesRegex(ValueError, 'foo'):
        _ = pc | schemas.UnbatchPandas(proxy, include_indexes=True)

  def test_unbatch_datetime(self):

    s = pd.Series(
        pd.date_range(
            '1/1/2000', periods=100, freq='m', tz='America/Los_Angeles'))
    proxy = s[:0]

    with TestPipeline() as p:
      res = (
          p | beam.Create([s[::2], s[1::2]])
          | schemas.UnbatchPandas(proxy, include_indexes=True))

      assert_that(res, equal_to(list(s)))


if __name__ == '__main__':
  unittest.main()