from __future__ import annotations

import collections
import re
import warnings
from copy import copy
from typing import TYPE_CHECKING
from typing import Any
from typing import Iterator

from narwhals._expression_parsing import evaluate_output_names_and_aliases
from narwhals._expression_parsing import is_simple_aggregation
from narwhals._pandas_like.utils import horizontal_concat
from narwhals._pandas_like.utils import native_series_from_iterable
from narwhals._pandas_like.utils import select_columns_by_name
from narwhals._pandas_like.utils import set_columns
from narwhals.utils import Implementation
from narwhals.utils import find_stacklevel

if TYPE_CHECKING:
    from typing_extensions import Self

    from narwhals._pandas_like.dataframe import PandasLikeDataFrame
    from narwhals._pandas_like.expr import PandasLikeExpr

POLARS_TO_PANDAS_AGGREGATIONS = {
    "sum": "sum",
    "mean": "mean",
    "median": "median",
    "max": "max",
    "min": "min",
    "std": "std",
    "var": "var",
    "len": "size",
    "n_unique": "nunique",
    "count": "count",
}


class PandasLikeGroupBy:
    def __init__(
        self: Self, df: PandasLikeDataFrame, keys: list[str], *, drop_null_keys: bool
    ) -> None:
        self._df = df
        self._keys = keys
        # Drop index to avoid potential collisions:
        # https://github.com/narwhals-dev/narwhals/issues/1907.
        native_frame = df._native_frame.reset_index(drop=True)
        if (
            self._df._implementation is Implementation.PANDAS
            and self._df._backend_version < (1, 1)
        ):  # pragma: no cover
            if (
                not drop_null_keys
                and self._df.simple_select(*self._keys)._native_frame.isna().any().any()
            ):
                msg = "Grouping by null values is not supported in pandas < 1.1.0"
                raise NotImplementedError(msg)
            self._grouped = native_frame.groupby(
                list(self._keys),
                sort=False,
                as_index=True,
                observed=True,
            )
        else:
            self._grouped = native_frame.groupby(
                list(self._keys),
                sort=False,
                as_index=True,
                dropna=drop_null_keys,
                observed=True,
            )

    def agg(self: Self, *exprs: PandasLikeExpr) -> PandasLikeDataFrame:  # noqa: PLR0915
        implementation = self._df._implementation
        backend_version = self._df._backend_version

        new_names: list[str] = copy(self._keys)
        for expr in exprs:
            _, aliases = evaluate_output_names_and_aliases(expr, self._df, self._keys)
            new_names.extend(aliases)

        all_aggs_are_simple = True
        for expr in exprs:
            if not (
                is_simple_aggregation(expr)
                and re.sub(r"(\w+->)", "", expr._function_name)
                in POLARS_TO_PANDAS_AGGREGATIONS
            ):
                all_aggs_are_simple = False

        # dict of {output_name: root_name} that we count n_unique on
        # We need to do this separately from the rest so that we
        # can pass the `dropna` kwargs.
        nunique_aggs: dict[str, str] = {}
        simple_aggs: dict[str, list[str]] = collections.defaultdict(list)

        # ddof to (output_names, aliases) mapping
        std_aggs: dict[int, tuple[list[str], list[str]]] = collections.defaultdict(
            lambda: ([], [])
        )
        var_aggs: dict[int, tuple[list[str], list[str]]] = collections.defaultdict(
            lambda: ([], [])
        )

        expected_old_names: list[str] = []
        simple_agg_new_names: list[str] = []

        if all_aggs_are_simple:
            for expr in exprs:
                output_names, aliases = evaluate_output_names_and_aliases(
                    expr, self._df, self._keys
                )
                if expr._depth == 0:
                    # e.g. agg(nw.len()) # noqa: ERA001
                    function_name = POLARS_TO_PANDAS_AGGREGATIONS.get(
                        expr._function_name, expr._function_name
                    )
                    for alias in aliases:
                        expected_old_names.append(f"{self._keys[0]}_{function_name}")
                        simple_aggs[self._keys[0]].append(function_name)
                        simple_agg_new_names.append(alias)
                    continue

                # e.g. agg(nw.mean('a')) # noqa: ERA001
                function_name = re.sub(r"(\w+->)", "", expr._function_name)
                function_name = POLARS_TO_PANDAS_AGGREGATIONS.get(
                    function_name, function_name
                )

                is_n_unique = function_name == "nunique"
                is_std = function_name == "std"
                is_var = function_name == "var"
                for output_name, alias in zip(output_names, aliases):
                    if is_n_unique:
                        nunique_aggs[alias] = output_name
                    elif is_std and (ddof := expr._kwargs["ddof"]) != 1:
                        std_aggs[ddof][0].append(output_name)
                        std_aggs[ddof][1].append(alias)
                    elif is_var and (ddof := expr._kwargs["ddof"]) != 1:
                        var_aggs[ddof][0].append(output_name)
                        var_aggs[ddof][1].append(alias)
                    else:
                        expected_old_names.append(f"{output_name}_{function_name}")
                        simple_aggs[output_name].append(function_name)
                        simple_agg_new_names.append(alias)

            result_aggs = []

            if simple_aggs:
                result_simple_aggs = self._grouped.agg(simple_aggs)
                result_simple_aggs.columns = [
                    f"{a}_{b}" for a, b in result_simple_aggs.columns
                ]
                if not (
                    set(result_simple_aggs.columns) == set(expected_old_names)
                    and len(result_simple_aggs.columns) == len(expected_old_names)
                ):  # pragma: no cover
                    msg = (
                        f"Safety assertion failed, expected {expected_old_names} "
                        f"got {result_simple_aggs.columns}, "
                        "please report a bug at https://github.com/narwhals-dev/narwhals/issues"
                    )
                    raise AssertionError(msg)

                # Rename columns, being very careful
                expected_old_names_indices: dict[str, list[int]] = (
                    collections.defaultdict(list)
                )
                for idx, item in enumerate(expected_old_names):
                    expected_old_names_indices[item].append(idx)
                index_map: list[int] = [
                    expected_old_names_indices[item].pop(0)
                    for item in result_simple_aggs.columns
                ]
                result_simple_aggs.columns = [simple_agg_new_names[i] for i in index_map]
                result_aggs.append(result_simple_aggs)

            if nunique_aggs:
                result_nunique_aggs = self._grouped[list(nunique_aggs.values())].nunique(
                    dropna=False
                )
                result_nunique_aggs.columns = list(nunique_aggs.keys())

                result_aggs.append(result_nunique_aggs)

            if std_aggs:
                result_aggs.extend(
                    [
                        set_columns(
                            self._grouped[std_output_names].std(ddof=ddof),
                            columns=std_aliases,
                            implementation=implementation,
                            backend_version=backend_version,
                        )
                        for ddof, (std_output_names, std_aliases) in std_aggs.items()
                    ]
                )
            if var_aggs:
                result_aggs.extend(
                    [
                        set_columns(
                            self._grouped[var_output_names].var(ddof=ddof),
                            columns=var_aliases,
                            implementation=implementation,
                            backend_version=backend_version,
                        )
                        for ddof, (var_output_names, var_aliases) in var_aggs.items()
                    ]
                )

            if result_aggs:
                output_names_counter = collections.Counter(
                    c for frame in result_aggs for c in frame
                )
                if any(v > 1 for v in output_names_counter.values()):
                    msg = ""
                    for key, value in output_names_counter.items():
                        if value > 1:
                            msg += f"\n- '{key}' {value} times"
                        else:  # pragma: no cover
                            pass
                    msg = f"Expected unique output names, got:{msg}"
                    raise ValueError(msg)
                result = horizontal_concat(
                    dfs=result_aggs,
                    implementation=implementation,
                    backend_version=backend_version,
                )
            else:
                # No aggregation provided
                result = self._df.__native_namespace__().DataFrame(
                    list(self._grouped.groups.keys()), columns=self._keys
                )
            # Keep inplace=True to avoid making a redundant copy.
            # This may need updating, depending on https://github.com/pandas-dev/pandas/pull/51466/files
            result.reset_index(inplace=True)  # noqa: PD002
            return self._df._from_native_frame(
                select_columns_by_name(result, new_names, backend_version, implementation)
            )

        if self._df._native_frame.empty:
            # Don't even attempt this, it's way too inconsistent across pandas versions.
            msg = (
                "No results for group-by aggregation.\n\n"
                "Hint: you were probably trying to apply a non-elementary aggregation with a "
                "pandas-like API.\n"
                "Please rewrite your query such that group-by aggregations "
                "are elementary. For example, instead of:\n\n"
                "    df.group_by('a').agg(nw.col('b').round(2).mean())\n\n"
                "use:\n\n"
                "    df.with_columns(nw.col('b').round(2)).group_by('a').agg(nw.col('b').mean())\n\n"
            )
            raise ValueError(msg)

        warnings.warn(
            "Found complex group-by expression, which can't be expressed efficiently with the "
            "pandas API. If you can, please rewrite your query such that group-by aggregations "
            "are simple (e.g. mean, std, min, max, ...). \n\n"
            "Please see: "
            "https://narwhals-dev.github.io/narwhals/pandas_like_concepts/improve_group_by_operation/",
            UserWarning,
            stacklevel=find_stacklevel(),
        )

        def func(df: Any) -> Any:
            out_group = []
            out_names = []
            for expr in exprs:
                results_keys = expr(self._df._from_native_frame(df))
                for result_keys in results_keys:
                    out_group.append(result_keys._native_series.iloc[0])
                    out_names.append(result_keys.name)
            return native_series_from_iterable(
                out_group,
                index=out_names,
                name="",
                implementation=implementation,
            )

        if implementation is Implementation.PANDAS and backend_version >= (2, 2):
            result_complex = self._grouped.apply(func, include_groups=False)
        else:  # pragma: no cover
            result_complex = self._grouped.apply(func)

        # Keep inplace=True to avoid making a redundant copy.
        # This may need updating, depending on https://github.com/pandas-dev/pandas/pull/51466/files
        result_complex.reset_index(inplace=True)  # noqa: PD002

        return self._df._from_native_frame(
            select_columns_by_name(
                result_complex, new_names, backend_version, implementation
            )
        )

    def __iter__(self: Self) -> Iterator[tuple[Any, PandasLikeDataFrame]]:
        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                message=".*a length 1 tuple will be returned",
                category=FutureWarning,
            )
            for key, group in self._grouped:
                yield (key, self._df._from_native_frame(group))
