# Copyright (c) ONNX Project Contributors

# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations

import itertools

import numpy as np

from onnx.reference.op_run import OpRun
from onnx.reference.ops._op_common_indices import _get_index, _get_indices


def _get_pad_shape(
    auto_pad: str,
    input_spatial_shape: tuple[int],
    kernel_spatial_shape: tuple[int],
    strides_spatial: tuple[int],
    output_spatial_shape: tuple[int],
) -> tuple[int]:
    pad_shape = [0] * len(input_spatial_shape)
    if auto_pad in ("SAME_UPPER", "SAME_LOWER"):
        for i in range(len(input_spatial_shape)):
            pad_shape[i] = (
                (output_spatial_shape[i] - 1) * strides_spatial[i]
                + kernel_spatial_shape[i]
                - input_spatial_shape[i]
            )
    elif auto_pad == "VALID":
        pass
    if len(pad_shape) == 0:
        raise RuntimeError(
            f"Unable to compute pad shape, auto_pad={auto_pad!r}, "
            f"input_spatial_shape={input_spatial_shape!r}, "
            f"kernel_spatial_shape={kernel_spatial_shape!r}, "
            f"strides_spatial={strides_spatial!r}."
        )
    return tuple(pad_shape)


def _get_output_shape_no_ceil(
    auto_pad: str,
    input_spatial_shape: tuple[int],
    kernel_spatial_shape: tuple[int],
    strides_spatial: tuple[int],
) -> tuple[int]:
    out_shape = [0] * len(input_spatial_shape)
    if auto_pad in ("SAME_UPPER", "SAME_LOWER"):
        for i in range(len(input_spatial_shape)):
            out_shape[i] = int(
                np.ceil(float(input_spatial_shape[i]) / float(strides_spatial[i]))
            )
    elif auto_pad == "VALID":
        for i in range(len(input_spatial_shape)):
            out_shape[i] = int(
                np.ceil(
                    float(input_spatial_shape[i] - (kernel_spatial_shape[i] - 1))
                    / float(strides_spatial[i])
                )
            )
    return tuple(out_shape)


def _get_output_shape(
    auto_pad: str,
    input_spatial_shape: tuple[int],
    kernel_spatial_shape: tuple[int],
    strides_spatial: tuple[int],
    pad_shape: tuple[int] | None = None,
    ceil_mode: int | None = 0,
) -> tuple[int]:
    if not ceil_mode:
        out_shape = _get_output_shape_no_ceil(
            auto_pad, input_spatial_shape, kernel_spatial_shape, strides_spatial
        )
    else:
        round_fct = np.ceil if ceil_mode else np.floor
        out_shape = [0] * len(input_spatial_shape)
        if auto_pad in ("SAME_UPPER", "SAME_LOWER"):
            for i in range(len(input_spatial_shape)):
                out_shape[i] = int(
                    round_fct(float(input_spatial_shape[i]) / float(strides_spatial[i]))
                )
        elif auto_pad == "VALID":
            if pad_shape is None:
                raise ValueError(  # pragma: no cogitver
                    "pad_shape cannot be None if auto_pad is "
                    "'VALID' and ceil_mode is 1."
                )
            for i in range(len(input_spatial_shape)):
                out_shape[i] = int(
                    round_fct(
                        float(
                            input_spatial_shape[i]
                            + pad_shape[i]
                            - kernel_spatial_shape[i]
                        )
                        / float(strides_spatial[i])
                        + 1
                    )
                )
    if len(out_shape) == 0:
        raise RuntimeError(
            f"Unable to compute output shape, auto_pad={auto_pad!r}, "
            f"input_spatial_shape={input_spatial_shape!r}, "
            f"kernel_spatial_shape={kernel_spatial_shape!r}, "
            f"strides_spatial={strides_spatial!r}, ceil_mode={ceil_mode!r}."
        )
    if min(out_shape) <= 0:
        raise RuntimeError(
            f"output shape cannot be null or negative, out_shape={out_shape!r}, "
            f"auto_pad={auto_pad!r}, input_spatial_shape={input_spatial_shape!r}, "
            f"kernel_spatial_shape={kernel_spatial_shape!r}, "
            f"strides_spatial={strides_spatial!r}, ceil_mode={ceil_mode!r}."
        )
    return tuple(out_shape)


def _pool(
    padded: np.ndarray,
    x_shape: tuple[int],
    kernel_shape: tuple[int],
    strides_shape: tuple[int],
    out_shape: tuple[int],
    pad_shape: tuple[int],
    pooling_type: str,
    count_include_pad: int | None = 0,
    ceil_mode: int | None = 0,
    indices: bool = False,
    pads: np.ndarray | None = None,
) -> np.ndarray:
    if pooling_type == "AVG":
        fpool = np.average
    elif pooling_type == "MAX":
        fpool = np.max
    else:
        raise NotImplementedError(
            f"Pooling type {pooling_type!r} does not support. Should be AVG, MAX."
        )
    spatial_size = len(x_shape) - 2
    y = np.zeros([x_shape[0], x_shape[1], *list(out_shape)])
    if indices:
        z = np.full(y.shape, fill_value=-1, dtype=np.int64)
    round_fct = np.ceil if ceil_mode else np.floor

    def loop_range():
        return [
            range(
                int(
                    round_fct(
                        float(x_shape[i + 2] + pad_shape[i] - kernel_shape[i])
                        / float(strides_shape[i])
                        + 1
                    )
                )
            )
            for i in range(spatial_size)
        ]

    for shape in itertools.product(range(x_shape[0]), range(x_shape[1]), *loop_range()):
        window = padded[shape[0], shape[1]]
        listi = [
            range(
                strides_shape[i] * shape[i + 2],
                strides_shape[i] * shape[i + 2] + kernel_shape[i],
            )
            for i in range(spatial_size)
        ]
        listi2 = list(itertools.product(*listi))
        values = []
        for i in listi2:
            try:
                values.append(window[i])
            except IndexError:  # noqa: PERF203
                continue
        window_vals = np.array(values)

        if count_include_pad == 1 and pooling_type == "AVG":
            y[shape] = fpool(window_vals)
        else:
            no_nan = window_vals[np.where(~np.isnan(window_vals))]
            y[shape] = fpool(no_nan)
            if indices:
                try:
                    window_vals_min = np.nan_to_num(window_vals, nan=no_nan.min())
                except TypeError:
                    # argument nan was introduced in numpy 1.17
                    window_vals_min = window_vals.copy()
                    window_vals_min[np.isnan(window_vals_min)] = no_nan.min()
                arg = np.argmax(window_vals_min)
                coordinates = _get_indices(arg, out_shape)
                delta = shape[2:] - pads[:, 0]
                coordinates += delta
                new_arg = _get_index(coordinates, x_shape[2:])
                z[shape] = new_arg
    if indices:
        return y.astype(padded.dtype), z
    return y.astype(padded.dtype)


class CommonPool(OpRun):
    def _run(
        self,
        pooling_type,
        count_include_pad,
        x,
        auto_pad=None,
        ceil_mode=None,
        dilations=None,
        kernel_shape=None,
        pads=None,
        storage_order=None,  # noqa: ARG002
        strides=None,
    ):
        if pooling_type == "MAX" and dilations is None:
            dilations = [1 for s in kernel_shape]
        if pads is None:
            pads = [0 for s in kernel_shape] * 2

        if strides is None or len(strides) == 0:
            strides = [1] * (len(x.shape) - 2)
        kernel_shape = list(kernel_shape)
        auto_pad = "VALID" if auto_pad == "NOTSET" else auto_pad

        if pads is None or len(pads) == 0:
            pad_shape = [0] * (len(x.shape) - 2)
            x_shape = x.shape[2:]
            padded = x
        elif len(pads) == 4:
            pad_top, pad_bottom, pad_left, pad_right = pads
            pad_shape = [pad_top + pad_bottom, pad_left + pad_right]
            x_shape = np.array(x.shape[2:]) + np.array(pad_shape)
            const = np.nan if count_include_pad == 0 else 0
            padded = np.pad(
                x,
                ((0, 0), (0, 0), (pad_top, pad_bottom), (pad_left, pad_right)),
                mode="constant",
                constant_values=const,
            )
        else:
            pad_shape = pads
            x_shape = x.shape[2:]
            padded = x

        if auto_pad in ("SAME_LOWER", "SAME_UPPER"):
            const = np.nan if count_include_pad == 0 else 0
            out_shape = _get_output_shape(
                auto_pad,
                x_shape,
                kernel_shape,
                strides,
                pad_shape,
                ceil_mode,
            )
            pad_shape = _get_pad_shape(
                auto_pad, x_shape, kernel_shape, strides, out_shape
            )
            if auto_pad == "SAME_LOWER":
                pad_bottom = pad_shape[0] // 2
                pad_top = pad_shape[0] - pad_bottom
                pad_right = pad_shape[1] // 2
                pad_left = pad_shape[1] - pad_right
            else:
                pad_top = pad_shape[0] // 2
                pad_bottom = pad_shape[0] - pad_top
                pad_left = pad_shape[1] // 2
                pad_right = pad_shape[1] - pad_left
            padded = np.pad(
                padded,
                ((0, 0), (0, 0), (pad_top, pad_bottom), (pad_left, pad_right)),
                mode="constant",
                constant_values=const,
            )
        else:
            out_shape = _get_output_shape(
                auto_pad,
                x_shape,
                kernel_shape,
                strides,
                pad_shape,
                ceil_mode,
            )

        n_dims = len(pads) // 2
        new_pads = np.array([(pads[i], pads[i + n_dims]) for i in range(n_dims)])
        res = _pool(
            padded,
            x.shape,
            kernel_shape,
            strides,
            out_shape,
            pad_shape,
            pooling_type,
            count_include_pad=count_include_pad,
            ceil_mode=ceil_mode,
            indices=len(self.output) > 1,
            pads=new_pads,
        )
        if isinstance(res, tuple):
            return res
        return (res,)