#
# Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
#

# Licensed to Modin Development Team under one or more contributor license agreements.
# See the NOTICE file distributed with this work for additional information regarding
# copyright ownership.  The Modin Development Team 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.

# Code in this file may constitute partial or total reimplementation, or modification of
# existing code originally distributed by the Modin project, under the Apache License,
# Version 2.0.

"""Implement Window and Rolling public API."""
from typing import Any, Literal, Optional, Union

import numpy as np  # noqa: F401
import pandas.core.window.rolling

from snowflake.snowpark.dataframe import DataFrame as SnowparkDataFrame

# add these two lines to enable doc tests to run
from snowflake.snowpark.modin import pandas as pd  # noqa: F401
from snowflake.snowpark.modin.plugin._internal.telemetry import TelemetryMeta
from snowflake.snowpark.modin.plugin.utils.warning_message import WarningMessage
from snowflake.snowpark.modin.utils import (
    _inherit_docstrings,
    doc_replace_dataframe_with_link,
)


# TODO SNOW-1041934: Add support for more window aggregations
class Window(metaclass=TelemetryMeta):
    def __init__(
        self,
        dataframe,
        window: Any = None,
        min_periods: int = None,
        center: bool = False,
        win_type: str = None,
        on: str = None,
        axis: Union[int, str] = 0,
        closed: str = None,
        step: int = None,
        method: str = "single",
    ) -> None:
        # TODO: SNOW-1063357: Modin upgrade - modin.pandas.window.Window
        self._dataframe = dataframe
        self._query_compiler = dataframe._query_compiler
        self.window_kwargs = {  # pragma: no cover
            "window": window,
            "min_periods": min_periods,
            "center": center,
            "win_type": win_type,
            "on": on,
            "axis": axis,
            "closed": closed,
            "step": step,
            "method": method,
        }
        self.axis = axis

    def mean(self, *args, **kwargs):
        # TODO: SNOW-1063357: Modin upgrade - modin.pandas.window.Window
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.window_mean(
                self.axis, self.window_kwargs, *args, **kwargs
            )
        )

    def sum(self, *args, **kwargs):
        # TODO: SNOW-1063357: Modin upgrade - modin.pandas.window.Window
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.window_sum(
                self.axis, self.window_kwargs, *args, **kwargs
            )
        )

    def var(self, ddof=1, *args, **kwargs):
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.window_var(
                self.axis, self.window_kwargs, ddof, *args, **kwargs
            )
        )

    def std(self, ddof=1, *args, **kwargs):
        # TODO: SNOW-1063357: Modin upgrade - modin.pandas.window.Window
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.window_std(
                self.axis, self.window_kwargs, ddof, *args, **kwargs
            )
        )


@_inherit_docstrings(
    pandas.core.window.rolling.Rolling,
    excluded=[pandas.core.window.rolling.Rolling.__init__],
    modify_doc=doc_replace_dataframe_with_link,
)
# TODO SNOW-1041934: Add support for more window aggregations
class Rolling(metaclass=TelemetryMeta):
    def __init__(
        self,
        dataframe,
        window: Any,
        min_periods: Optional[int] = None,
        center: bool = False,
        win_type: Optional[str] = None,
        on: Optional[str] = None,
        axis: Union[int, str] = 0,
        closed: Optional[str] = None,
        step: Optional[int] = None,
        method: str = "single",
    ) -> None:
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        # Raise ValueError when invalid parameter values/combinations
        if (isinstance(window, int) and window < 0) or window is None:
            raise ValueError("window must be an integer 0 or greater")
        if not isinstance(center, bool):
            raise ValueError("center must be a boolean")
        if min_periods is not None and not isinstance(min_periods, int):
            raise ValueError("min_periods must be an integer")
        if isinstance(min_periods, int) and min_periods < 0:
            raise ValueError("min_periods must be >= 0")
        if (
            isinstance(min_periods, int)
            and isinstance(window, int)
            and min_periods > window
        ):
            raise ValueError(f"min_periods {min_periods} must be <= window {window}")

        self._dataframe = dataframe
        self._query_compiler = dataframe._query_compiler
        self.rolling_kwargs = {
            "window": window,
            "min_periods": min_periods,
            "center": center,
            "win_type": win_type,
            "on": on,
            "axis": axis,
            "closed": closed,
            "step": step,
            "method": method,
        }
        self.axis = axis
        if method != "single":
            WarningMessage.ignored_argument(
                operation="Rolling",
                argument="method",
                message="Snowpark pandas API executes on Snowflake. Ignoring engine related arguments to select a different execution engine.",
            )

    def _call_qc_method(self, method_name, *args, **kwargs):
        """
        Call a query compiler method for the specified rolling aggregation.

        Parameters
        ----------
        method_name : str
            Name of the aggregation.
        *args : tuple
            Positional arguments to pass to the query compiler method.
        **kwargs : dict
            Keyword arguments to pass to the query compiler method.

        Returns
        -------
        BaseQueryCompiler
            QueryCompiler holding the result of the aggregation.
        """
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        qc_method = getattr(self._query_compiler, f"rolling_{method_name}")
        return qc_method(self.axis, self.rolling_kwargs, *args, **kwargs)

    def _aggregate(self, method_name, *args, **kwargs):
        """
        Run the specified rolling aggregation.

        Parameters
        ----------
        method_name : str
            Name of the aggregation.
        *args : tuple
            Positional arguments to pass to the aggregation.
        **kwargs : dict
            Keyword arguments to pass to the aggregation.

        Returns
        -------
        DataFrame or Series
            Result of the aggregation.
        """
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        qc_result = self._call_qc_method(method_name, *args, **kwargs)
        return self._dataframe.__constructor__(query_compiler=qc_result)

    def count(self, numeric_only: bool = False):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(method_name="count", numeric_only=numeric_only)

    def sum(
        self,
        numeric_only: bool = False,
        *args: Any,
        engine: Optional[Literal["cython", "numba"]] = None,
        engine_kwargs: Optional[dict[str, bool]] = None,
        **kwargs: Any,
    ):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(
            method_name="sum",
            numeric_only=numeric_only,
            engine=engine,
            engine_kwargs=engine_kwargs,
            *args,
            **kwargs,
        )

    def mean(
        self,
        numeric_only: bool = False,
        *args: Any,
        engine: Optional[Literal["cython", "numba"]] = None,
        engine_kwargs: Optional[dict[str, bool]] = None,
        **kwargs: Any,
    ):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(
            method_name="mean",
            numeric_only=numeric_only,
            engine=engine,
            engine_kwargs=engine_kwargs,
            *args,
            **kwargs,
        )

    def median(
        self,
        numeric_only: bool = False,
        engine: Optional[Literal["cython", "numba"]] = None,
        engine_kwargs: Optional[dict[str, bool]] = None,
        **kwargs: Any,
    ):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(
            method_name="median",
            numeric_only=numeric_only,
            engine=engine,
            engine_kwargs=engine_kwargs,
            **kwargs,
        )

    def var(
        self,
        ddof: int = 1,
        numeric_only: bool = False,
        *args: Any,
        engine: Optional[Literal["cython", "numba"]] = None,
        engine_kwargs: Optional[dict[str, bool]] = None,
        **kwargs: Any,
    ):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(
            method_name="var",
            ddof=ddof,
            numeric_only=numeric_only,
            engine=engine,
            engine_kwargs=engine_kwargs,
            *args,
            **kwargs,
        )

    def std(
        self,
        ddof: int = 1,
        numeric_only: bool = False,
        *args: Any,
        engine: Optional[Literal["cython", "numba"]] = None,
        engine_kwargs: Optional[dict[str, bool]] = None,
        **kwargs: Any,
    ):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(
            method_name="std",
            ddof=ddof,
            numeric_only=numeric_only,
            engine=engine,
            engine_kwargs=engine_kwargs,
            *args,
            **kwargs,
        )

    def min(
        self,
        numeric_only: bool = False,
        *args: Any,
        engine: Optional[Literal["cython", "numba"]] = None,
        engine_kwargs: Optional[dict[str, bool]] = None,
        **kwargs: Any,
    ):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(
            method_name="min",
            numeric_only=numeric_only,
            engine=engine,
            engine_kwargs=engine_kwargs,
            *args,
            **kwargs,
        )

    def max(
        self,
        numeric_only: bool = False,
        *args: Any,
        engine: Optional[Literal["cython", "numba"]] = None,
        engine_kwargs: Optional[dict[str, bool]] = None,
        **kwargs: Any,
    ):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(
            method_name="max",
            numeric_only=numeric_only,
            engine=engine,
            engine_kwargs=engine_kwargs,
            *args,
            **kwargs,
        )

    def corr(
        self,
        other: Optional[SnowparkDataFrame] = None,
        pairwise: Optional[bool] = None,
        ddof: int = 1,
        numeric_only: bool = False,
        **kwargs: Any,
    ):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(
            method_name="corr",
            other=other,
            pairwise=pairwise,
            ddof=ddof,
            numeric_only=numeric_only,
            **kwargs,
        )

    def cov(
        self,
        other: Optional[SnowparkDataFrame] = None,
        pairwise: Optional[bool] = None,
        ddof: int = 1,
        numeric_only: bool = False,
        **kwargs: Any,
    ):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(
            method_name="cov",
            other=other,
            pairwise=pairwise,
            ddof=ddof,
            numeric_only=numeric_only,
            **kwargs,
        )

    def skew(
        self,
        numeric_only: bool = False,
        **kwargs: Any,
    ):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(method_name="skew", numeric_only=numeric_only, **kwargs)

    def kurt(
        self,
        numeric_only: bool = False,
        **kwargs: Any,
    ):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(method_name="kurt", numeric_only=numeric_only, **kwargs)

    def apply(
        self,
        func: Any,
        raw: bool = False,
        engine: Optional[Literal["cython", "numba"]] = None,
        engine_kwargs: Optional[dict[str, bool]] = None,
        args: Optional[tuple] = None,
        kwargs: Optional[dict] = None,
    ):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(
            method_name="apply",
            func=func,
            raw=raw,
            engine=engine,
            engine_kwargs=engine_kwargs,
            args=args,
            kwargs=kwargs,
        )

    def aggregate(
        self,
        func: Union[str, list, dict],
        *args: Any,
        **kwargs: Any,
    ):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(method_name="aggregate", func=func, *args, **kwargs)

    agg = aggregate

    def quantile(
        self,
        quantile: float,
        interpolation: str = "linear",
        numeric_only: bool = False,
        **kwargs: Any,
    ):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(
            method_name="quantile",
            quantile=quantile,
            interpolation=interpolation,
            numeric_only=numeric_only,
            **kwargs,
        )

    def sem(
        self,
        ddof: int = 1,
        numeric_only: bool = False,
        *args: Any,
        **kwargs: Any,
    ):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(
            method_name="sem", ddof=ddof, numeric_only=numeric_only, *args, **kwargs
        )

    def rank(
        self,
        method: str = "average",
        ascending: bool = True,
        pct: bool = False,
        numeric_only: bool = False,
        **kwargs,
    ):
        # TODO: SNOW-1063358: Modin upgrade - modin.pandas.window.Rolling
        return self._aggregate(
            method_name="rank",
            method=method,
            ascending=ascending,
            pct=pct,
            numeric_only=numeric_only,
            **kwargs,
        )


@_inherit_docstrings(
    pandas.core.window.expanding.Expanding,
    excluded=[pandas.core.window.expanding.Expanding.__init__],
    modify_doc=doc_replace_dataframe_with_link,
)
class Expanding(metaclass=TelemetryMeta):
    def __init__(
        self,
        dataframe,
        min_periods: int = 1,
        axis: Union[int, str] = 0,
        method: str = "single",
    ) -> None:
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        if min_periods is not None and not isinstance(min_periods, int):
            raise ValueError("min_periods must be an integer")
        if isinstance(min_periods, int) and min_periods < 0:
            raise ValueError("min_periods must be >= 0")

        self._dataframe = dataframe
        self._query_compiler = dataframe._query_compiler
        self.expanding_kwargs = {
            "min_periods": min_periods,
            "axis": axis,
            "method": method,
        }
        self.axis = axis
        if method != "single":
            WarningMessage.ignored_argument(
                operation="Expanding",
                argument="method",
                message="Snowpark pandas API executes on Snowflake. Ignoring engine related arguments to select a different execution engine.",
            )

    def count(
        self,
        numeric_only: bool = False,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_count(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                numeric_only=numeric_only,
            )
        )

    def sum(
        self,
        numeric_only: bool = False,
        engine: Optional[Literal["cython", "numba"]] = None,
        engine_kwargs: Optional[dict[str, bool]] = None,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_sum(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                numeric_only=numeric_only,
                engine=engine,
                engine_kwargs=engine_kwargs,
            )
        )

    def mean(
        self,
        numeric_only: bool = False,
        engine: Optional[Literal["cython", "numba"]] = None,
        engine_kwargs: Optional[dict[str, bool]] = None,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_mean(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                numeric_only=numeric_only,
                engine=engine,
                engine_kwargs=engine_kwargs,
            )
        )

    def median(
        self,
        numeric_only: bool = False,
        engine: Optional[Literal["cython", "numba"]] = None,
        engine_kwargs: Optional[dict[str, bool]] = None,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_median(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                numeric_only=numeric_only,
                engine=engine,
                engine_kwargs=engine_kwargs,
            )
        )

    def var(
        self,
        ddof: int = 1,
        numeric_only: bool = False,
        engine: Optional[Literal["cython", "numba"]] = None,
        engine_kwargs: Optional[dict[str, bool]] = None,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_var(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                ddof=ddof,
                numeric_only=numeric_only,
                engine=engine,
                engine_kwargs=engine_kwargs,
            )
        )

    def std(
        self,
        ddof: int = 1,
        numeric_only: bool = False,
        engine: Optional[Literal["cython", "numba"]] = None,
        engine_kwargs: Optional[dict[str, bool]] = None,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_std(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                ddof=ddof,
                numeric_only=numeric_only,
                engine=engine,
                engine_kwargs=engine_kwargs,
            )
        )

    def min(
        self,
        numeric_only: bool = False,
        engine: Optional[Literal["cython", "numba"]] = None,
        engine_kwargs: Optional[dict[str, bool]] = None,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_min(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                numeric_only=numeric_only,
                engine=engine,
                engine_kwargs=engine_kwargs,
            )
        )

    def max(
        self,
        numeric_only: bool = False,
        engine: Optional[Literal["cython", "numba"]] = None,
        engine_kwargs: Optional[dict[str, bool]] = None,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_max(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                numeric_only=numeric_only,
                engine=engine,
                engine_kwargs=engine_kwargs,
            )
        )

    def corr(
        self,
        other: Optional[SnowparkDataFrame] = None,
        pairwise: Optional[bool] = None,
        ddof: int = 1,
        numeric_only: bool = False,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_corr(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                other=other,
                pairwise=pairwise,
                ddof=ddof,
                numeric_only=numeric_only,
            )
        )

    def cov(
        self,
        other: Optional[SnowparkDataFrame] = None,
        pairwise: Optional[bool] = None,
        ddof: int = 1,
        numeric_only: bool = False,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_cov(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                other=other,
                pairwise=pairwise,
                ddof=ddof,
                numeric_only=numeric_only,
            )
        )

    def skew(
        self,
        numeric_only: bool = False,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_skew(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                numeric_only=numeric_only,
            )
        )

    def kurt(
        self,
        numeric_only: bool = False,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_kurt(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                numeric_only=numeric_only,
            )
        )

    def apply(
        self,
        func: Any,
        raw: bool = False,
        engine: Optional[Literal["cython", "numba"]] = None,
        engine_kwargs: Optional[dict[str, bool]] = None,
        args: Optional[tuple] = None,
        kwargs: Optional[dict] = None,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_apply(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                func=func,
                raw=raw,
                engine=engine,
                engine_kwargs=engine_kwargs,
                args=args,
                kwargs=kwargs,
            )
        )

    def aggregate(
        self,
        func: Any,
        *args,
        **kwargs,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_aggregate(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                func=func,
                *args,
                **kwargs,
            )
        )

    def quantile(
        self,
        quantile: float,
        interpolation: str = "linear",
        numeric_only: bool = False,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_quantile(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                quantile=quantile,
                interpolation=interpolation,
                numeric_only=numeric_only,
            )
        )

    def sem(
        self,
        ddof: int = 1,
        numeric_only: bool = False,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_sem(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                ddof=ddof,
                numeric_only=numeric_only,
            )
        )

    def rank(
        self,
        method: str = "average",
        ascending: bool = True,
        pct: bool = False,
        numeric_only: bool = False,
    ):
        # TODO: SNOW-1063366: Modin upgrade - modin.pandas.window.Expanding
        return self._dataframe.__constructor__(
            query_compiler=self._query_compiler.expanding_rank(
                fold_axis=self.axis,
                expanding_kwargs=self.expanding_kwargs,
                method=method,
                ascending=ascending,
                pct=pct,
                numeric_only=numeric_only,
            )
        )
