#
#  Copyright (2024) Hewlett Packard Enterprise Development LP
#
#  Licensed 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.
#

from typing import Optional

from numpy.typing import NDArray

from .types import CamOp, CamVariant, is_float_type


def validate_args(
    variant: CamVariant,
    op: CamOp,
    inputs: NDArray,
    cam: NDArray,
    reduction_values: Optional[NDArray] = None,
) -> None:
    """
    Ensures that the arguments for a CAM operation are sane and compatible
    by inspecting their dtypes and shapes.
    Raises appropriate errors in case a problem is found.
    """

    # dtype checking

    if inputs.dtype != cam.dtype:
        raise TypeError(
            f"inputs and cam dtypes are not equal: ({inputs.dtype} vs {cam.dtype})"
        )

    if op.is_reduction:
        if reduction_values is None:
            raise TypeError("operation is reduction but no reduction values were given")

        if not is_float_type(reduction_values.dtype):
            raise TypeError("reductions only support float data types")

    # shape checking

    if len(inputs.shape) < 2:
        raise ValueError("inputs are one-dimensional")

    if len(cam.shape) < 2:
        raise ValueError("cam is one-dimensional")

    cam_cell_width = variant.cam_cell_width
    if cam.shape[-1] % cam_cell_width != 0:
        raise ValueError(
            "amount of CAM columns is not divisible by CAM cell encoding width: "
            f"{cam_cell_width} (columns found: {cam.shape[-1]})"
        )

    input_cell_width = variant.input_cell_width
    if inputs.shape[-1] % input_cell_width != 0:
        raise ValueError(
            "amount of input columns is not divisible by input cell encoding width: "
            f"{input_cell_width} (columns found: {inputs.shape[-1]})"
        )

    if inputs.shape[-1] // input_cell_width != cam.shape[-1] // cam_cell_width:
        raise ValueError(
            "input and cam column sizes do not match "
            f"a ratio of {input_cell_width}:{cam_cell_width}"
        )

    if op.is_reduction:
        assert reduction_values is not None
        if len(reduction_values.shape) > 1:
            raise ValueError("reduction values must be one-dimensional")
        if len(inputs.shape) > 2 or len(cam.shape) > 2:
            raise ValueError("reduction operations do not support broadcasting")
        if reduction_values.shape != cam.shape[:-1]:
            raise ValueError("CAM shape does not match reduction value shape")

    inputs_outer_shape, cam_outer_shape = inputs.shape[:-2], cam.shape[:-2]

    for i in range(1, min(len(inputs_outer_shape), len(cam_outer_shape)) + 1):
        if inputs_outer_shape[-i] != cam_outer_shape[-i]:
            raise ValueError(
                "the outer shapes of input and cam are incompatible: "
                f"{inputs_outer_shape} vs {cam_outer_shape}"
            )
