#
# 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 inspect
import warnings
import weakref
from collections.abc import Iterable
from typing import Any
from typing import Optional
from typing import Union

import pandas as pd

import apache_beam as beam
from apache_beam import pvalue
from apache_beam.dataframe import expressions
from apache_beam.dataframe import frame_base
from apache_beam.dataframe import transforms
from apache_beam.dataframe.schemas import element_typehint_from_dataframe_proxy
from apache_beam.dataframe.schemas import generate_proxy
from apache_beam.typehints.pandas_type_compatibility import dtype_to_fieldtype


# TODO: Or should this be called as_dataframe?
def to_dataframe(
    pcoll: pvalue.PCollection,
    proxy: Optional[pd.core.generic.NDFrame] = None,
    label: Optional[str] = None,
) -> frame_base.DeferredFrame:
  """Converts a PCollection to a deferred dataframe-like object, which can
  manipulated with pandas methods like `filter` and `groupby`.

  For example, one might write::

    pcoll = ...
    df = to_dataframe(pcoll, proxy=...)
    result = df.groupby('col').sum()
    pcoll_result = to_pcollection(result)

  A proxy object must be given if the schema for the PCollection is not known.
  """
  if proxy is None:
    if pcoll.element_type is None:
      raise ValueError(
          "Cannot infer a proxy because the input PCollection does not have a "
          "schema defined. Please make sure a schema type is specified for "
          "the input PCollection, or provide a proxy.")
    # If no proxy is given, assume this is an element-wise schema-aware
    # PCollection that needs to be batched.
    if label is None:
      # Attempt to come up with a reasonable, stable label by retrieving
      # the name of these variables in the calling context.
      label = 'BatchElements(%s)' % _var_name(pcoll, 2)
    proxy = generate_proxy(pcoll.element_type)

    shim_dofn: beam.DoFn
    if isinstance(proxy, pd.DataFrame):
      shim_dofn = RowsToDataFrameFn()
    elif isinstance(proxy, pd.Series):
      shim_dofn = ElementsToSeriesFn()
    else:
      raise AssertionError("Unknown proxy type: %s" % proxy)

    pcoll = pcoll | label >> beam.ParDo(shim_dofn)
  return frame_base.DeferredFrame.wrap(
      expressions.PlaceholderExpression(proxy, pcoll))


# PCollections generated by to_pcollection are memoized, keyed by expression id.
# WeakValueDictionary is used so the caches are cleaned up with the parent
# pipelines
# Note that the pipeline (indirectly) holds references to the transforms which
# keeps both the PCollections and expressions alive. This ensures the
# expression's ids are never accidentally re-used.
TO_PCOLLECTION_CACHE: 'weakref.WeakValueDictionary[str, pvalue.PCollection]' = (
    weakref.WeakValueDictionary())
UNBATCHED_CACHE: 'weakref.WeakValueDictionary[str, pvalue.PCollection]' = (
    weakref.WeakValueDictionary())


class RowsToDataFrameFn(beam.DoFn):
  @beam.DoFn.yields_elements
  def process_batch(self, batch: pd.DataFrame) -> Iterable[pd.DataFrame]:
    yield batch


class ElementsToSeriesFn(beam.DoFn):
  @beam.DoFn.yields_elements
  def process_batch(self, batch: pd.Series) -> Iterable[pd.Series]:
    yield batch


def _make_unbatched_pcoll(
    pc: pvalue.PCollection, expr: expressions.Expression,
    include_indexes: bool):
  label = f"Unbatch '{expr._id}'"
  if include_indexes:
    label += " with indexes"

  if label not in UNBATCHED_CACHE:
    proxy = expr.proxy()
    shim_dofn: beam.DoFn
    if isinstance(proxy, pd.DataFrame):
      shim_dofn = DataFrameToRowsFn(proxy, include_indexes)
    elif isinstance(proxy, pd.Series):
      if include_indexes:
        warnings.warn(
            "Pipeline is converting a DeferredSeries to PCollection "
            "with include_indexes=True. Note that this parameter is "
            "_not_ respected for DeferredSeries conversion. To "
            "include the index with your data, produce a"
            "DeferredDataFrame instead.")

      shim_dofn = SeriesToElementsFn(proxy)
    else:
      raise TypeError(f"Proxy '{proxy}' has unsupported type '{type(proxy)}'")

    UNBATCHED_CACHE[label] = pc | label >> beam.ParDo(shim_dofn)

  # Note unbatched cache is keyed by the expression id as well as parameters
  # for the unbatching (i.e. include_indexes)
  return UNBATCHED_CACHE[label]


class DataFrameToRowsFn(beam.DoFn):
  def __init__(self, proxy, include_indexes):
    self._proxy = proxy
    self._include_indexes = include_indexes

  @beam.DoFn.yields_batches
  def process(self, element: pd.DataFrame) -> Iterable[pd.DataFrame]:
    yield element

  def infer_output_type(self, input_element_type):
    return element_typehint_from_dataframe_proxy(
        self._proxy, self._include_indexes)


class SeriesToElementsFn(beam.DoFn):
  def __init__(self, proxy):
    self._proxy = proxy

  @beam.DoFn.yields_batches
  def process(self, element: pd.Series) -> Iterable[pd.Series]:
    yield element

  def infer_output_type(self, input_element_type):
    return dtype_to_fieldtype(self._proxy.dtype)


# TODO: Or should this be called from_dataframe?


def to_pcollection(
    *dataframes: Union[frame_base.DeferredFrame, pd.DataFrame, pd.Series],
    label=None,
    always_return_tuple=False,
    yield_elements='schemas',
    include_indexes=False,
    pipeline=None) -> Union[pvalue.PCollection, tuple[pvalue.PCollection, ...]]:
  """Converts one or more deferred dataframe-like objects back to a PCollection.

  This method creates and applies the actual Beam operations that compute
  the given deferred dataframes, returning a PCollection of their results. By
  default the resulting PCollections are schema-aware PCollections where each
  element is one row from the output dataframes, excluding indexes. This
  behavior can be modified with the `yield_elements` and `include_indexes`
  arguments.

  Also accepts non-deferred pandas dataframes, which are converted to deferred,
  schema'd PCollections. In this case the contents of the entire dataframe are
  serialized into the graph, so for large amounts of data it is preferable to
  write them to disk and read them with one of the read methods.

  If more than one (related) result is desired, it can be more efficient to
  pass them all at the same time to this method.

  Args:
    label: (optional, default "ToPCollection(...)"") the label to use for the
        conversion transform.
    always_return_tuple: (optional, default: False) If true, always return
        a tuple of PCollections, even if there's only one output.
    yield_elements: (optional, default: "schemas") If set to "pandas", return
        PCollections containing the raw Pandas objects (DataFrames or Series),
        if set to "schemas", return an element-wise PCollection, where DataFrame
        and Series instances are expanded to one element per row. DataFrames are
        converted to schema-aware PCollections, where column values can be
        accessed by attribute.
    include_indexes: (optional, default: False) When yield_elements="schemas",
        if include_indexes=True, attempt to include index columns in the output
        schema for expanded DataFrames. Raises an error if any of the index
        levels are unnamed (name=None), or if any of the names are not unique
        among all column and index names.
    pipeline: (optional, unless non-deferred dataframes are passed) Used when
        creating a PCollection from a non-deferred dataframe.
  """
  if not yield_elements in ("pandas", "schemas"):
    raise ValueError(
        "Invalid value for yield_elements argument, '%s'. "
        "Allowed values are 'pandas' and 'schemas'" % yield_elements)
  if label is None:
    # Attempt to come up with a reasonable, stable label by retrieving the name
    # of these variables in the calling context.
    label = 'ToPCollection(%s)' % ', '.join(_var_name(e, 3) for e in dataframes)

  # Support for non-deferred dataframes.
  deferred_dataframes = []
  for ix, df in enumerate(dataframes):
    if isinstance(df, frame_base.DeferredBase):
      # TODO(robertwb): Maybe extract pipeline object?
      deferred_dataframes.append(df)
    elif isinstance(df, (pd.Series, pd.DataFrame)):
      if pipeline is None:
        raise ValueError(
            'Pipeline keyword required for non-deferred dataframe conversion.')
      deferred = pipeline | '%s_Defer%s' % (label, ix) >> beam.Create([df])
      deferred_dataframes.append(
          frame_base.DeferredFrame.wrap(
              expressions.PlaceholderExpression(df.iloc[:0], deferred)))
    else:
      raise TypeError(
          'Unable to convert objects of type %s to a PCollection' % type(df))
  dataframes = tuple(deferred_dataframes)

  def extract_input(placeholder):
    if not isinstance(placeholder._reference, pvalue.PCollection):
      raise TypeError(
          'Expression roots must have been created with to_dataframe.')
    return placeholder._reference

  placeholders = frozenset.union(
      frozenset(), *[df._expr.placeholders() for df in dataframes])

  # Exclude any dataframes that have already been converted to PCollections.
  # We only want to convert each DF expression once, then re-use.
  new_dataframes = [
      df for df in dataframes if df._expr._id not in TO_PCOLLECTION_CACHE
  ]
  if len(new_dataframes):
    new_results: dict[Any, pvalue.PCollection] = {
        p: extract_input(p)
        for p in placeholders
    } | label >> transforms._DataframeExpressionsTransform(
        {ix: df._expr
         for (ix, df) in enumerate(new_dataframes)})

    TO_PCOLLECTION_CACHE.update({
        new_dataframes[ix]._expr._id: pc
        for ix, pc in new_results.items()
    })

  raw_results = {
      ix: TO_PCOLLECTION_CACHE[df._expr._id]
      for ix, df in enumerate(dataframes)
  }

  if yield_elements == "schemas":

    def maybe_unbatch(pc, value):
      if isinstance(value, frame_base._DeferredScalar):
        return pc
      else:
        return _make_unbatched_pcoll(pc, value._expr, include_indexes)

    results = {
        ix: maybe_unbatch(pc, dataframes[ix])
        for (ix, pc) in raw_results.items()
    }
  else:
    results = raw_results

  if len(results) == 1 and not always_return_tuple:
    return results[0]
  else:
    return tuple(value for key, value in sorted(results.items()))


def _var_name(obj, level):
  frame = inspect.currentframe()
  for _ in range(level):
    if frame is None:
      return '...'
    frame = frame.f_back
  for key, value in frame.f_locals.items():
    if obj is value:
      return key
  for key, value in frame.f_globals.items():
    if obj is value:
      return key
  return '...'