#
# 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 typing
import unittest
import warnings

import pandas as pd

import apache_beam as beam
from apache_beam import coders
from apache_beam import metrics
from apache_beam.dataframe import convert
from apache_beam.dataframe import expressions
from apache_beam.dataframe import frame_base
from apache_beam.dataframe import transforms
from apache_beam.runners.portability.fn_api_runner import fn_runner
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to


def check_correct(expected, actual):
  if actual is None:
    raise AssertionError('Empty frame but expected: \n\n%s' % (expected))
  if isinstance(expected, pd.core.generic.NDFrame):
    expected = expected.sort_index()
    actual = actual.sort_index()

    if isinstance(expected, pd.Series):
      pd.testing.assert_series_equal(expected, actual)
    elif isinstance(expected, pd.DataFrame):
      pd.testing.assert_frame_equal(expected, actual)
    else:
      raise ValueError(
          f"Expected value is a {type(expected)},"
          "not a Series or DataFrame.")
  else:
    if actual != expected:
      raise AssertionError('Scalars not equal: %s != %s' % (actual, expected))


def concat(parts):
  if len(parts) > 1:
    return pd.concat(parts)
  elif len(parts) == 1:
    return parts[0]
  else:
    return None


def df_equal_to(expected):
  return lambda actual: check_correct(expected, concat(actual))


AnimalSpeed = typing.NamedTuple(
    'AnimalSpeed', [('Animal', str), ('Speed', int)])
coders.registry.register_coder(AnimalSpeed, coders.RowCoder)
Nested = typing.NamedTuple(
    'Nested', [('id', int), ('animal_speed', AnimalSpeed)])
coders.registry.register_coder(Nested, coders.RowCoder)


class TransformTest(unittest.TestCase):
  def run_scenario(self, input, func):
    expected = func(input)

    empty = input.iloc[0:0]
    input_placeholder = expressions.PlaceholderExpression(empty)
    input_deferred = frame_base.DeferredFrame.wrap(input_placeholder)
    actual_deferred = func(input_deferred)._expr.evaluate_at(
        expressions.Session({input_placeholder: input}))

    check_correct(expected, actual_deferred)

    with beam.Pipeline() as p:
      input_pcoll = p | beam.Create([input.iloc[::2], input.iloc[1::2]])
      input_df = convert.to_dataframe(input_pcoll, proxy=empty)
      output_df = func(input_df)

      output_proxy = output_df._expr.proxy()
      if isinstance(output_proxy, pd.core.generic.NDFrame):
        self.assertTrue(
            output_proxy.iloc[:0].equals(expected.iloc[:0]),
            (
                'Output proxy is incorrect:\n'
                f'Expected:\n{expected.iloc[:0]}\n\n'
                f'Actual:\n{output_proxy.iloc[:0]}'))
      else:
        self.assertEqual(type(output_proxy), type(expected))

      output_pcoll = convert.to_pcollection(output_df, yield_elements='pandas')

      assert_that(
          output_pcoll, lambda actual: check_correct(expected, concat(actual)))

  def test_identity(self):
    df = pd.DataFrame({
        'Animal': ['Falcon', 'Falcon', 'Parrot', 'Parrot'],
        'Speed': [380., 370., 24., 26.]
    })
    self.run_scenario(df, lambda x: x)

  def test_groupby_sum_mean(self):
    df = pd.DataFrame({
        'Animal': ['Falcon', 'Falcon', 'Parrot', 'Parrot'],
        'Speed': [380., 370., 24., 26.]
    })
    self.run_scenario(df, lambda df: df.groupby('Animal').sum())
    with expressions.allow_non_parallel_operations():
      self.run_scenario(df, lambda df: df.groupby('Animal').mean())
    self.run_scenario(
        df, lambda df: df.loc[df.Speed > 25].groupby('Animal').sum())

  def test_groupby_apply(self):
    df = pd.DataFrame({
        'group': ['a' if i % 5 == 0 or i % 3 == 0 else 'b' for i in range(100)],
        'foo': [None if i % 11 == 0 else i for i in range(100)],
        'bar': [None if i % 7 == 0 else 99 - i for i in range(100)],
        'baz': [None if i % 13 == 0 else i * 2 for i in range(100)],
    })

    def median_sum_fn(x):
      with warnings.catch_warnings():
        warnings.filterwarnings("ignore", message="Mean of empty slice")
        return (x.foo + x.bar).median()

    describe = lambda df: df.describe()

    self.run_scenario(df, lambda df: df.groupby('group').foo.apply(describe))
    self.run_scenario(
        df, lambda df: df.groupby('group')[['foo', 'bar']].apply(describe))
    self.run_scenario(df, lambda df: df.groupby('group').apply(median_sum_fn))
    self.run_scenario(
        df,
        lambda df: df.set_index('group').foo.groupby(level=0).apply(describe))
    self.run_scenario(df, lambda df: df.groupby(level=0).apply(median_sum_fn))
    self.run_scenario(
        df, lambda df: df.groupby(lambda x: x % 3).apply(describe))

  def test_filter(self):
    df = pd.DataFrame({
        'Animal': ['Aardvark', 'Ant', 'Elephant', 'Zebra'],
        'Speed': [5, 2, 35, 40]
    })
    self.run_scenario(df, lambda df: df.filter(items=['Animal']))
    self.run_scenario(df, lambda df: df.filter(regex='Anim.*'))
    self.run_scenario(
        df, lambda df: df.set_index('Animal').filter(regex='F.*', axis='index'))

    with expressions.allow_non_parallel_operations():
      a = pd.DataFrame({'col': [1, 2, 3]})
      self.run_scenario(a, lambda a: a.agg(sum))
      self.run_scenario(a, lambda a: a.agg(['mean', 'min', 'max']))

  def test_scalar(self):
    with expressions.allow_non_parallel_operations():
      a = pd.Series([1, 2, 6])
      self.run_scenario(a, lambda a: a.agg(sum))
      self.run_scenario(a, lambda a: a / a.agg(sum))
      self.run_scenario(a, lambda a: a / (a.max() - a.min()))
      self.run_scenario(a, lambda a: a / (a.sum() - 1))

      # Tests scalar being used as an input to a downstream stage.
      df = pd.DataFrame({'key': ['a', 'a', 'b'], 'val': [1, 2, 6]})
      self.run_scenario(
          df, lambda df: df.groupby('key').sum().val / df.val.agg(sum))

  def test_getitem_projection(self):
    df = pd.DataFrame({
        'Animal': ['Aardvark', 'Ant', 'Elephant', 'Zebra'],
        'Speed': [5, 2, 35, 40],
        'Size': ['Small', 'Extra Small', 'Large', 'Medium']
    })
    self.run_scenario(df, lambda df: df[['Speed', 'Size']])

  def test_offset_elementwise(self):
    s = pd.Series(range(10)).astype(float)
    df = pd.DataFrame({'value': s, 'square': s * s, 'cube': s * s * s})
    # Only those values that are both squares and cubes will intersect.
    self.run_scenario(
        df,
        lambda df: df.set_index('square').value + df.set_index('cube').value)

  def test_batching_named_tuple_input(self):
    with beam.Pipeline() as p:
      result = (
          p | beam.Create([
              AnimalSpeed('Aardvark', 5),
              AnimalSpeed('Ant', 2),
              AnimalSpeed('Elephant', 35),
              AnimalSpeed('Zebra', 40)
          ]).with_output_types(AnimalSpeed)
          | transforms.DataframeTransform(lambda df: df.filter(regex='Anim.*')))

      assert_that(
          result,
          equal_to([('Aardvark', ), ('Ant', ), ('Elephant', ), ('Zebra', )]))

  def test_batching_beam_row_input(self):
    with beam.Pipeline() as p:
      result = (
          p
          | beam.Create([('Falcon', 380.), ('Falcon', 370.), ('Parrot', 24.),
                         ('Parrot', 26.)])
          | beam.Map(lambda tpl: beam.Row(Animal=tpl[0], Speed=tpl[1]))
          | transforms.DataframeTransform(
              lambda df: df.groupby('Animal').mean(), include_indexes=True))

      assert_that(result, equal_to([('Falcon', 375.), ('Parrot', 25.)]))

  def test_batching_beam_row_to_dataframe(self):
    with beam.Pipeline() as p:
      df = convert.to_dataframe(
          p
          | beam.Create([('Falcon', 380.), ('Falcon', 370.), ('Parrot', 24.), (
              'Parrot', 26.)])
          | beam.Map(lambda tpl: beam.Row(Animal=tpl[0], Speed=tpl[1])))

      result = convert.to_pcollection(
          df.groupby('Animal').mean(), include_indexes=True)

      assert_that(result, equal_to([('Falcon', 375.), ('Parrot', 25.)]))

  def test_batching_passthrough_nested_schema(self):
    with beam.Pipeline() as p:
      nested_schema_pc = (
          p | beam.Create([Nested(1, AnimalSpeed('Aardvark', 5))
                           ]).with_output_types(Nested))
      result = nested_schema_pc | transforms.DataframeTransform(  # pylint: disable=expression-not-assigned
          lambda df: df.filter(items=['animal_speed']))

      assert_that(result, equal_to([(('Aardvark', 5), )]))

  def test_batching_passthrough_nested_array(self):
    Array = typing.NamedTuple(
        'Array', [('id', int), ('business_numbers', typing.Sequence[int])])
    coders.registry.register_coder(Array, coders.RowCoder)

    with beam.Pipeline() as p:
      array_schema_pc = (p | beam.Create([Array(1, [7, 8, 9])]))
      result = array_schema_pc | transforms.DataframeTransform(  # pylint: disable=expression-not-assigned
            lambda df: df.filter(items=['business_numbers']))

      assert_that(result, equal_to([([7, 8, 9], )]))

  def test_unbatching_series(self):
    with beam.Pipeline() as p:
      result = (
          p
          | beam.Create([('Falcon', 380.), ('Falcon', 370.), ('Parrot', 24.),
                         ('Parrot', 26.)])
          | beam.Map(lambda tpl: beam.Row(Animal=tpl[0], Speed=tpl[1]))
          | transforms.DataframeTransform(lambda df: df.Animal))

      assert_that(result, equal_to(['Falcon', 'Falcon', 'Parrot', 'Parrot']))

  def test_input_output_polymorphism(self):
    one_series = pd.Series([1])
    two_series = pd.Series([2])
    three_series = pd.Series([3])
    proxy = one_series[:0]

    def equal_to_series(expected):
      def check(actual):
        actual = pd.concat(actual)
        if not expected.equals(actual):
          raise AssertionError(
              'Series not equal: \n%s\n%s\n' % (expected, actual))

      return check

    with beam.Pipeline() as p:
      one = p | 'One' >> beam.Create([one_series])
      two = p | 'Two' >> beam.Create([two_series])

      assert_that(
          one | 'PcollInPcollOut' >> transforms.DataframeTransform(
              lambda x: 3 * x, proxy=proxy, yield_elements='pandas'),
          equal_to_series(three_series),
          label='CheckPcollInPcollOut')

      assert_that(
          (one, two)
          | 'TupleIn' >> transforms.DataframeTransform(
              lambda x, y: (x + y), (proxy, proxy), yield_elements='pandas'),
          equal_to_series(three_series),
          label='CheckTupleIn')

      assert_that(
          dict(x=one, y=two)
          | 'DictIn' >> transforms.DataframeTransform(
              lambda x, y: (x + y),
              proxy=dict(x=proxy, y=proxy),
              yield_elements='pandas'),
          equal_to_series(three_series),
          label='CheckDictIn')

      double, triple = one | 'TupleOut' >> transforms.DataframeTransform(
              lambda x: (2*x, 3*x), proxy, yield_elements='pandas')
      assert_that(double, equal_to_series(two_series), 'CheckTupleOut0')
      assert_that(triple, equal_to_series(three_series), 'CheckTupleOut1')

      res = one | 'DictOut' >> transforms.DataframeTransform(
          lambda x: {'res': 3 * x}, proxy, yield_elements='pandas')
      assert_that(res['res'], equal_to_series(three_series), 'CheckDictOut')

  def test_multiple_dataframes_transforms(self):
    expected_output = ["Bryan", "DKER2"]

    def transform_func(a, b):
      b["name"] = "DKER2"
      return a, b

    with beam.Pipeline() as p:
      pcol1 = p | "Create1" >> beam.Create([beam.Row(name="Bryan")])
      pcol2 = p | "Create2" >> beam.Create([beam.Row(name="common")])

      result = ({
          "a": pcol1, "b": pcol2
      }
                |
                "TransformDF" >> transforms.DataframeTransform(transform_func)
                | "Flatten" >> beam.Flatten()
                | transforms.DataframeTransform(lambda df: df.name))
      assert_that(result, equal_to(expected_output))

  def test_cat(self):
    # verify that cat works with a List[Series] since this is
    # missing from doctests
    df = pd.DataFrame({
        'one': ['A', 'B', 'C'],
        'two': ['BB', 'CC', 'A'],
        'three': ['CCC', 'AA', 'B'],
    })
    self.run_scenario(df, lambda df: df.two.str.cat([df.three], join='outer'))
    self.run_scenario(
        df, lambda df: df.one.str.cat([df.two, df.three], join='outer'))

  def test_repeat(self):
    # verify that repeat works with a Series since this is
    # missing from doctests
    df = pd.DataFrame({
        'strings': ['A', 'B', 'C', 'D', 'E'],
        'repeats': [3, 1, 4, 5, 2],
    })
    self.run_scenario(df, lambda df: df.strings.str.repeat(df.repeats))

  def test_rename(self):
    df = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})
    self.run_scenario(
        df, lambda df: df.rename(columns={'B': 'C'}, index={
            0: 2, 2: 0
        }))

    with expressions.allow_non_parallel_operations():
      self.run_scenario(
          df, lambda df: df.rename(
              columns={'B': 'C'}, index={
                  0: 2, 2: 0
              }, errors='raise'))


class FusionTest(unittest.TestCase):
  @staticmethod
  def fused_stages(p):
    return p.result.monitoring_metrics().query(
        metrics.MetricsFilter().with_name(
            fn_runner.FnApiRunner.NUM_FUSED_STAGES_COUNTER)
    )['counters'][0].result

  @staticmethod
  def create_animal_speed_input(p):
    return p | beam.Create([
        AnimalSpeed('Aardvark', 5),
        AnimalSpeed('Ant', 2),
        AnimalSpeed('Elephant', 35),
        AnimalSpeed('Zebra', 40)
    ],
                           reshuffle=False)

  def test_loc_filter(self):
    # TODO(https://github.com/apache/beam/issues/34549): This test relies on
    # monitoring_metrics property of the FnApiRunner which does not exist on
    # other runners like Prism.
    # https://github.com/apache/beam/blob/5f9cd73b7c9a2f37f83971ace3a399d633201dd1/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py#L1590
    with beam.Pipeline('FnApiRunner') as p:
      _ = (
          self.create_animal_speed_input(p)
          | transforms.DataframeTransform(lambda df: df[df.Speed > 10]))
    self.assertEqual(self.fused_stages(p), 1)

  def test_column_manipulation(self):
    def set_column(df, name, s):
      df[name] = s
      return df

    # TODO(https://github.com/apache/beam/issues/34549): This test relies on
    # monitoring_metrics property of the FnApiRunner which does not exist on
    # other runners like Prism.
    # https://github.com/apache/beam/blob/5f9cd73b7c9a2f37f83971ace3a399d633201dd1/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py#L1590
    with beam.Pipeline('FnApiRunner') as p:
      _ = (
          self.create_animal_speed_input(p)
          | transforms.DataframeTransform(
              lambda df: set_column(df, 'x', df.Speed + df.Animal.str.len())))
    self.assertEqual(self.fused_stages(p), 1)


class TransformPartsTest(unittest.TestCase):
  def test_rebatch(self):
    with beam.Pipeline() as p:
      sA = pd.Series(range(1000))
      sB = sA * sA
      pcA = p | 'CreatePCollA' >> beam.Create([('k0', sA[::3]),
                                               ('k1', sA[1::3]),
                                               ('k2', sA[2::3])])
      pcB = p | 'CreatePCollB' >> beam.Create([('k0', sB[::3]),
                                               ('k1', sB[1::3]),
                                               ('k2', sB[2::3])])
      input = {'A': pcA, 'B': pcB} | beam.CoGroupByKey()
      output = input | beam.ParDo(
          transforms._ReBatch(target_size=sA.memory_usage()))

      # There should be exactly two elements, as the target size will be
      # hit when 2/3 of pcA and 2/3 of pcB is seen, but not before.
      assert_that(output | beam.combiners.Count.Globally(), equal_to([2]))

      # Sanity check that we got all the right values.
      assert_that(
          output | beam.Map(lambda x: x['A'].sum())
          | 'SumA' >> beam.CombineGlobally(sum),
          equal_to([sA.sum()]),
          label='CheckValuesA')
      assert_that(
          output | beam.Map(lambda x: x['B'].sum())
          | 'SumB' >> beam.CombineGlobally(sum),
          equal_to([sB.sum()]),
          label='CheckValuesB')


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