# ===----------------------------------------------------------------------=== #
# Nabla 2025
#
# 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.
# ===----------------------------------------------------------------------=== #

"""Linear algebra operations."""

import numpy as np
from max.graph import TensorValue, ops

from ..core.array import Array
from ..utils.shape_utils import get_broadcasted_shape
from .operation import BinaryOperation

# Public API
__all__ = ["matmul"]


class MatMulOp(BinaryOperation):
    """Matrix multiplication operation with batching support."""

    def __init__(self):
        super().__init__("dot_general")

    def forward(self, *args: Array) -> Array:
        """Forward pass for binary operations."""
        if len(args) != 2:
            raise ValueError(f"Binary operation requires 2 arguments, got {len(args)}")

        # Move arrays to best device
        from .operation import move_to_best_device

        args = move_to_best_device(*args)
        arg1, arg2 = args[0], args[1]

        from ..ops.view import broadcast_batch_dims, broadcast_to, reshape

        arg1_has_rank_1 = len(arg1.shape) == 1
        arg2_has_rank_1 = len(arg2.shape) == 1

        if arg1_has_rank_1:
            arg1 = reshape(arg1, (1, arg1.shape[0]))

        if arg2_has_rank_1:
            arg2 = reshape(arg2, (arg2.shape[0], 1))

        self._validate_inputs(arg1, arg2)

        output_shape = self.compute_output_shape(arg1.shape, arg2.shape)
        output_batch_dims = self.compute_output_batch_dims(
            arg1.batch_dims, arg2.batch_dims
        )
        output_dtype = self.compute_output_dtype(arg1, arg2)
        if arg1.traced:
            arg1 = broadcast_to(arg1, output_shape[:-2] + arg1.shape[-2:])
            arg1 = broadcast_batch_dims(arg1, output_batch_dims)
        if arg2.traced:
            arg2 = broadcast_to(arg2, output_shape[:-2] + arg2.shape[-2:])
            arg2 = broadcast_batch_dims(arg2, output_batch_dims)

        res = Array(
            shape=output_shape,
            dtype=output_dtype,
            device=arg1.device,
            materialize=False,
            name=self.name,
            batch_dims=output_batch_dims,
        )

        res.set_maxpr(self.maxpr)
        res.add_arguments(arg1, arg2)
        res.vjp_rule = self.vjp_rule
        res.jvp_rule = self.jvp_rule
        res.custom_kernel_path = self.custom_kernel_path()

        if not res.stage_realization:
            self.eagerxpr([arg1, arg2], res)

        # If the inputs were rank 1, we need to reshape the output
        if arg1_has_rank_1 and arg2_has_rank_1:
            res = reshape(res, output_shape[:-2] + (1, 1))
        elif arg1_has_rank_1:
            res = reshape(res, output_shape[:-2] + (res.shape[1],))
        elif arg2_has_rank_1:
            res = reshape(res, output_shape[:-2] + (res.shape[0],))

        res.creator_op = self
        return res

    def compute_output_shape(self, *input_shapes: tuple) -> tuple:
        """Compute output shape for matrix multiplication with compatible signature."""
        if len(input_shapes) != 2:
            raise ValueError(
                f"Matrix multiplication requires 2 input shapes, got {len(input_shapes)}"
            )
        shape1, shape2 = input_shapes[0], input_shapes[1]

        if shape1[-1] != shape2[-2]:
            raise ValueError(
                f"Shapes {shape1} and {shape2} are not compatible for matrix multiplication"
            )

        return get_broadcasted_shape(
            shape1,
            shape2,
            ignore_axes=[-2, -1],
            replace_ignored_dims=[shape1[-2], shape2[-1]],
        )

    def _validate_inputs(self, arg1: Array, arg2: Array) -> None:
        """Validate matrix multiplication inputs."""
        if not isinstance(arg1, Array) or not isinstance(arg2, Array):
            raise TypeError("Both arguments must be Array instances")
        if arg1.dtype != arg2.dtype:
            raise ValueError(f"Dtypes {arg1.dtype} and {arg2.dtype} are incompatible")
        if arg1.device != arg2.device:
            raise ValueError(
                f"Devices {arg1.device} and {arg2.device} are incompatible"
            )
        if arg1.shape[-1] != arg2.shape[-2]:
            raise ValueError(
                f"Shapes {arg1.shape} and {arg2.shape} are not compatible for matrix multiplication"
            )

    def maxpr(self, args: list[TensorValue], output: Array) -> None:
        x_val, y_val = args[0], args[1]
        x_shape = x_val.shape
        y_shape = y_val.shape
        output_shape = output.batch_dims + output.shape

        if len(output_shape) <= 4:
            output.tensor_value = ops.matmul(args[0], args[1])
        else:
            if x_shape[:-2] != y_shape[:-2]:
                raise ValueError(
                    f"Shapes {x_shape} and {y_shape} are not compatible for matrix multiplication "
                    f"(batch dimensions mismatch: {x_shape[:-2]} vs {y_shape[:-2]})"
                )
            # now we can simpply reshape the args to a rank3 tensor respecitvely and then do a batche dmamtul on this one
            batch_dims_x = [int(dim) for dim in x_shape[:-2]]
            batch_dims_y = [int(dim) for dim in y_shape[:-2]]
            new_shape_x = (
                np.prod(batch_dims_x).item(),
                int(x_shape[-2]),
                int(x_shape[-1]),
            )
            new_shape_y = (
                np.prod(batch_dims_y).item(),
                int(y_shape[-2]),
                int(y_shape[-1]),
            )
            x_val_b = ops.reshape(x_val, new_shape_x)
            y_val_b = ops.reshape(y_val, new_shape_y)
            matmul_result = ops.matmul(x_val_b, y_val_b)
            reshaped_result = ops.reshape(
                matmul_result,
                tuple(args[0].shape[:-2])
                + (matmul_result.shape[-2], matmul_result.shape[-1]),
            )
            output.tensor_value = reshaped_result

    def eagerxpr(self, args: list[Array], output: Array) -> None:
        arg0_numpy = args[0].to_numpy()
        arg1_numpy = args[1].to_numpy()
        np_result = np.matmul(arg0_numpy, arg1_numpy)
        output.impl_(np_result)

    def vjp_rule(
        self, primals: list[Array], cotangent: Array, output: Array
    ) -> list[Array]:
        x, y = primals
        from .view import transpose

        return [matmul(cotangent, transpose(y)), matmul(transpose(x), cotangent)]

    def jvp_rule(
        self, primals: list[Array], tangents: list[Array], output: Array
    ) -> Array:
        x, y = primals
        tx, ty = tangents

        from .binary import add

        return add(matmul(x, ty), matmul(tx, y))


# Global operation instance for efficiency
_matmul_op = MatMulOp()


def matmul(arg0, arg1) -> Array:
    """Matrix multiplication with broadcasting support."""
    from .binary import _ensure_array

    arg0 = _ensure_array(arg0)
    arg1 = _ensure_array(arg1)
    return _matmul_op.forward(arg0, arg1)


# # --- Convolution operations using im2col and col2im ---
# # Global operation instances
# _conv2d_op_cache = {}
# _conv2d_transpose_op_cache = {}


# # --- Helper functions for normalization ---
# def _normalize_tuple(value, n, name):
#     if isinstance(value, int):
#         return (value,) * n
#     elif isinstance(value, tuple | list):
#         if len(value) == n:
#             return tuple(value)
#         else:
#             raise ValueError(
#                 f"{name} must be an int or a tuple of {n} ints, got {value}"
#             )
#     else:
#         raise TypeError(
#             f"{name} must be an int or a tuple, got {type(value)} for {name}"
#         )


# def _normalize_padding_arg(padding_arg, name="padding"):
#     if isinstance(padding_arg, int):  # single int for all sides
#         return ((padding_arg, padding_arg), (padding_arg, padding_arg))
#     if isinstance(padding_arg, tuple | list):
#         if len(padding_arg) == 2:
#             if all(isinstance(x, int) for x in padding_arg):  # (symmetric_H, symmetric_W)
#                 ph, pw = padding_arg
#                 return ((ph, ph), (pw, pw))
#             elif all(isinstance(x, tuple | list) and len(x) == 2 and all(isinstance(y, int) for y in x) for x in padding_arg):
#                 # ((H_top, H_bottom), (W_left, W_right))
#                 return tuple(map(tuple, padding_arg))
#         elif len(padding_arg) == 4 and all(isinstance(x, int) for x in padding_arg):
#             # (H_top, H_bottom, W_left, W_right)
#             pt, pb, pl, pr = padding_arg
#             return ((pt, pb), (pl, pr))
#     raise ValueError(
#         f"{name} format is not recognized. Use int, (ph,pw), (pt,pb,pl,pr), or ((pt,pb),(pl,pr)). Got {padding_arg}"
#     )


# def flip(x: Array, axis: int | tuple[int, ...]) -> Array:
#     """
#     Reverses the order of elements in an array along the given axes.
#     This is an implementation of np.flip using fundamental slicing.
#     """
#     if isinstance(axis, int):
#         axes_to_flip = (axis,)
#     else:
#         axes_to_flip = axis

#     # Create a list of slice(None) objects, one for each dimension
#     slicer = [slice(None)] * len(x.shape)

#     # For each axis to be flipped, set the corresponding slice to ::-1
#     for ax in axes_to_flip:
#         slicer[ax] = slice(None, None, -1)

#     # Use tuple slicing on the array. The Nabla Array class's __getitem__
#     # must support this to be Python-idiomatic.
#     return x[tuple(slicer)]

# def _conv2d_filter_gradient(
#     x: Array, dy: Array, stride: tuple, dilation: tuple, padding: tuple, groups: int
# ) -> Array:
#     """
#     Computes `grad_W = conv(permute(x), permute(dy))` for a standard conv2d.
#     Returns a filter gradient in HWIO layout.
#     """
#     from ..ops import view

#     # Permute input x (NHWC) to be the data for the new conv: (Cin, H, W, N)
#     x_perm = view.transpose(x, (3, 1, 2, 0))

#     # Permute grad_output dy (NH'W'Cout) to be the filter for the new conv: (H', W', N, Cout)
#     dy_perm = view.transpose(dy, (1, 2, 0, 3))

#     # The new convolution's stride is the original's dilation, and vice versa.
#     # This is a standard identity for this gradient formulation.
#     grad_filter_permuted = conv2d(
#         x_perm, dy_perm, stride=dilation, dilation=stride, padding=padding, groups=groups
#     )

#     # The output is (Cin, kH, kW, Cout). Permute back to standard filter layout.
#     return view.transpose

# class Conv2DOp(BinaryOperation):
#     # ... This class is likely correct, but its VJP depends on the functions below ...
#     # Keep the version from my previous answer. The key fix is in Conv2DTransposeOp's VJP.
#     # ... For completeness, I'll include it with the corrected VJP rule call ...
#     """2D Convolution operation.

#     Data Layout: NHWC (batch, height, width, in_channels)
#     Filter Layout: HWIO (height, width, in_channels/groups, out_channels)
#     """

#     def __init__(self, stride, dilation, padding, groups):
#         super().__init__("conv2d")
#         self.stride = stride
#         self.dilation = dilation
#         self.padding = padding
#         self.groups = groups

#     def compute_output_shape(self, *input_shapes: tuple) -> tuple:
#         input_shape, filter_shape = input_shapes
#         n, h_in, w_in, c_in = input_shape
#         k_h, k_w, f_cin_div_g, f_cout = filter_shape

#         if c_in != f_cin_div_g * self.groups:
#             raise ValueError(
#                 f"Input channels ({c_in}) must match filter's effective input channels "
#                 f"({f_cin_div_g} * {self.groups} groups = {f_cin_div_g * self.groups}). "
#                 f"Input shape: {input_shape}, Filter shape: {filter_shape}"
#             )

#         (pad_h_top, pad_h_bottom), (pad_w_left, pad_w_right) = self.padding
#         dil_h, dil_w = self.dilation
#         s_h, s_w = self.stride

#         h_out = (h_in + pad_h_top + pad_h_bottom - dil_h * (k_h - 1) - 1) // s_h + 1
#         w_out = (w_in + pad_w_left + pad_w_right - dil_w * (k_w - 1) - 1) // s_w + 1
#         c_out = f_cout

#         if h_out <= 0 or w_out <= 0:
#             raise ValueError(f"Computed non-positive output dimensions for Conv2D: {(n, h_out, w_out, c_out)}")

#         return (n, h_out, w_out, c_out)

#     def forward(self, *args: Array) -> Array:
#         # Standard forward pass logic
#         from .operation import move_to_best_device
#         input_arr, filter_arr = move_to_best_device(*args)
#         self._validate_inputs(input_arr, filter_arr)

#         output_shape = self.compute_output_shape(input_arr.shape, filter_arr.shape)
#         res = Array(
#             shape=output_shape, dtype=self.compute_output_dtype(input_arr, filter_arr),
#             device=input_arr.device, materialize=False, name=self.name,
#             batch_dims=input_arr.batch_dims,
#         )
#         res.set_maxpr(self.maxpr)
#         res.add_arguments(input_arr, filter_arr)
#         res.vjp_rule = self.vjp_rule
#         res.jvp_rule = self.jvp_rule
#         if not res.stage_realization:
#             self.eagerxpr([input_arr, filter_arr], res)
#         return res

#     def _validate_inputs(self, input_arr: Array, filter_arr: Array) -> None:
#         if len(input_arr.shape) != 4 or len(filter_arr.shape) != 4:
#             raise ValueError("Conv2D requires 4D input and filter tensors.")
#         if input_arr.device != filter_arr.device:
#             raise ValueError(f"Devices {input_arr.device} and {filter_arr.device} are incompatible")

#     def maxpr(self, args: list[TensorValue], output: Array) -> None:
#         input_val, filter_val = args
#         (pt, pb), (pl, pr) = self.padding
#         output.tensor_value = ops.conv2d(
#             x=input_val, filter=filter_val, stride=self.stride,
#             dilation=self.dilation, padding=(pt, pb, pl, pr), groups=self.groups
#         )

#     def eagerxpr(self, args: list[Array], output: Array) -> None:
#         input_arr, filter_arr = args
#         input_torch = torch.from_numpy(np.transpose(input_arr.to_numpy(), (0, 3, 1, 2)))
#         filter_torch = torch.from_numpy(np.transpose(filter_arr.to_numpy(), (3, 2, 0, 1)))
#         (pad_h, _), (pad_w, _) = self.padding
#         result_torch = F.conv2d(
#             input=input_torch, weight=filter_torch, bias=None, stride=self.stride,
#             padding=(pad_h, pad_w), dilation=self.dilation, groups=self.groups
#         )
#         result_nhwc = np.transpose(result_torch.numpy(), (0, 2, 3, 1))
#         output.impl_(result_nhwc)

#     def vjp_rule(self, primals: list[Array], cotangent: Array, output: Array) -> list[Array]:
#         """VJP of Y = conv(X, W)"""
#         input_arr, filter_arr = primals # filter_arr is HWIO

#         # 1. grad_input = conv_transpose(dY, W_flipped_180)
#         flipped_filter = flip(filter_arr, axis=(0, 1))
#         # Filter for conv_transpose must be HWOI. Swap channels of our HWIO filter.
#         filter_for_grad_input = flipped_filter.transpose((0, 1, 3, 2))

#         # Calculate output_padding to restore original input shape
#         h_in, w_in = input_arr.shape[1:3]
#         h_out, w_out = cotangent.shape[1:3]
#         k_h, k_w = filter_arr.shape[0:2]
#         (pt, pb), (pl, pr) = self.padding
#         sh, sw = self.stride
#         dh, dw = self.dilation
#         out_pad_h = h_in - ((h_out - 1) * sh - (pt + pb) + (k_h - 1) * dh + 1)
#         out_pad_w = w_in - ((w_out - 1) * sw - (pl + pr) + (k_w - 1) * dw + 1)

#         grad_input = conv2d_transpose(
#             cotangent, filter_for_grad_input, stride=self.stride, dilation=self.dilation,
#             padding=self.padding, output_padding=(max(0,out_pad_h), max(0,out_pad_w)), groups=self.groups
#         )

#         # 2. grad_filter = conv(permute(X), permute(dY))
#         grad_filter = _conv2d_filter_gradient(
#             input_arr, cotangent, self.stride, self.dilation, self.padding, self.groups
#         )

#         return [grad_input, grad_filter]


#     def jvp_rule(self, primals: list[Array], tangents: list[Array], output: Array) -> Array:
#         input_arr, filter_arr = primals
#         input_tangent, filter_tangent = tangents
#         from .binary import add
#         res1 = conv2d(input_tangent, filter_arr, stride=self.stride, dilation=self.dilation, padding=self.padding, groups=self.groups)
#         res2 = conv2d(input_arr, filter_tangent, stride=self.stride, dilation=self.dilation, padding=self.padding, groups=self.groups)
#         return add(res1, res2)

# class Conv2DTransposeOp(BinaryOperation):
#     # This is the class with the key fixes.
#     # ... (Keep the __init__, compute_output_shape, forward, _validate_inputs, maxpr, eagerxpr from my PREVIOUS answer)
#     # The only change is in the VJP RULE.

#     def __init__(self, stride, dilation, padding, output_padding, groups):
#         super().__init__("conv2d_transpose")
#         self.stride = stride
#         self.dilation = dilation
#         self.padding = padding
#         self.output_padding = output_padding
#         self.groups = groups

#     def compute_output_shape(self, *input_shapes: tuple) -> tuple:
#         input_shape, filter_shape = input_shapes
#         n, h_in, w_in, c_in = input_shape
#         k_h, k_w, f_cout, f_cin_div_g = filter_shape
#         if c_in != f_cin_div_g * self.groups:
#              raise ValueError(
#                 f"Input channels ({c_in}) must match filter's effective input channels "
#                 f"({f_cin_div_g} * {self.groups} groups = {f_cin_div_g * self.groups}). "
#                 f"This is the 'I' in HWOI. "
#                 f"Input shape: {input_shape}, Filter shape: {filter_shape}"
#             )
#         (pad_h_top, pad_h_bottom), (pad_w_left, pad_w_right) = self.padding
#         out_pad_h, out_pad_w = self.output_padding
#         dil_h, dil_w = self.dilation
#         s_h, s_w = self.stride
#         h_out = (h_in - 1) * s_h - (pad_h_top + pad_h_bottom) + dil_h * (k_h - 1) + 1 + out_pad_h
#         w_out = (w_in - 1) * s_w - (pad_w_left + pad_w_right) + dil_w * (k_w - 1) + 1 + out_pad_w
#         c_out = f_cout
#         if h_out <= 0 or w_out <= 0:
#             raise ValueError(f"Computed non-positive output dimensions for Conv2DTranspose: {(n, h_out, w_out, c_out)}")
#         return (n, h_out, w_out, c_out)

#     def forward(self, *args: Array) -> Array:
#         from .operation import move_to_best_device
#         input_arr, filter_arr = move_to_best_device(*args)
#         self._validate_inputs(input_arr, filter_arr)
#         output_shape = self.compute_output_shape(input_arr.shape, filter_arr.shape)
#         res = Array(
#             shape=output_shape, dtype=self.compute_output_dtype(input_arr, filter_arr),
#             device=input_arr.device, materialize=False, name=self.name,
#             batch_dims=input_arr.batch_dims,
#         )
#         res.set_maxpr(self.maxpr)
#         res.add_arguments(input_arr, filter_arr)
#         res.vjp_rule = self.vjp_rule
#         res.jvp_rule = self.jvp_rule
#         if not res.stage_realization:
#             self.eagerxpr([input_arr, filter_arr], res)
#         return res

#     def _validate_inputs(self, input_arr: Array, filter_arr: Array) -> None:
#         if len(input_arr.shape) != 4 or len(filter_arr.shape) != 4:
#             raise ValueError("Conv2DTranspose requires 4D input and filter tensors.")
#         if input_arr.device != filter_arr.device:
#             raise ValueError(f"Devices {input_arr.device} and {filter_arr.device} are incompatible")

#     def maxpr(self, args: list[TensorValue], output: Array) -> None:
#         input_val, filter_val = args
#         (pt, pb), (pl, pr) = self.padding
#         if self.groups > 1:
#             from ..ops.view import split, concatenate
#             input_chunks = split(input_val, self.groups, axis=3)
#             filter_chunks = split(filter_val, self.groups, axis=3)
#             output_chunks = []
#             for i in range(self.groups):
#                 chunk_out = ops.conv2d_transpose(
#                     input_chunks[i], filter_chunks[i], stride=self.stride,
#                     dilation=self.dilation, padding=(pt, pb, pl, pr),
#                     output_paddings=self.output_padding
#                 )
#                 output_chunks.append(chunk_out)
#             output.tensor_value = concatenate(output_chunks, axis=3)
#         else:
#             output.tensor_value = ops.conv2d_transpose(
#                 input_val, filter_val, stride=self.stride, dilation=self.dilation,
#                 padding=(pt, pb, pl, pr), output_paddings=self.output_padding
#             )

#     def eagerxpr(self, args: list[Array], output: Array) -> None:
#         input_arr, filter_arr = args
#         input_torch = torch.from_numpy(np.transpose(input_arr.to_numpy(), (0, 3, 1, 2)))
#         filter_torch = torch.from_numpy(np.transpose(filter_arr.to_numpy(), (3, 2, 0, 1)))
#         (pad_h, _), (pad_w, _) = self.padding
#         result_torch = F.conv_transpose2d(
#             input=input_torch, weight=filter_torch, bias=None, stride=self.stride,
#             padding=(pad_h, pad_w), output_padding=self.output_padding,
#             groups=self.groups, dilation=self.dilation
#         )
#         result_nhwc = np.transpose(result_torch.numpy(), (0, 2, 3, 1))
#         output.impl_(result_nhwc)

#     def vjp_rule(self, primals: list[Array], cotangent: Array, output: Array) -> list[Array]:
#         """VJP of Y = conv_transpose(X, W)"""
#         input_arr, filter_arr = primals # filter_arr is HWOI

#         # 1. grad_input = conv(dY, W_flipped_180)
#         flipped_filter = flip(filter_arr, axis=(0, 1))
#         # Filter for conv2d must be HWIO. Swap channels of our HWOI filter.
#         filter_for_grad_input = flipped_filter.transpose((0, 1, 3, 2))

#         grad_input = conv2d(
#             cotangent, filter_for_grad_input, stride=self.stride,
#             dilation=self.dilation, padding=self.padding, groups=self.groups
#         )

#         # 2. grad_filter = conv(permute(dY), permute(X))
#         # Note the swapped arguments compared to the conv2d VJP.
#         grad_filter_HWIO = _conv2d_filter_gradient(
#             cotangent, input_arr, self.stride, self.dilation, self.padding, self.groups
#         )

#         # The helper returns HWIO. The gradient must match the primal filter's HWOI layout.
#         grad_filter = grad_filter_HWIO.transpose((0, 1, 3, 2))

#         return [grad_input, grad_filter]

#     def jvp_rule(self, primals: list[Array], tangents: list[Array], output: Array) -> Array:
#         input_arr, filter_arr = primals
#         input_tangent, filter_tangent = tangents
#         from .binary import add
#         res1 = conv2d_transpose(
#             input_tangent, filter_arr, stride=self.stride, dilation=self.dilation,
#             padding=self.padding, output_padding=self.output_padding, groups=self.groups)
#         res2 = conv2d_transpose(
#             input_arr, filter_tangent, stride=self.stride, dilation=self.dilation,
#             padding=self.padding, output_padding=self.output_padding, groups=self.groups)
#         return add(res1, res2)

# def conv2d(
#     input_arr: Array, filter_arr: Array, stride=(1, 1),
#     dilation=(1, 1), padding=0, groups=1
# ) -> Array:
#     """Applies a 2D convolution."""
#     norm_stride = _normalize_tuple(stride, 2, "stride")
#     norm_dilation = _normalize_tuple(dilation, 2, "dilation")
#     norm_padding = _normalize_padding_arg(padding, "padding")

#     cache_key = (norm_stride, norm_dilation, norm_padding, groups)
#     if cache_key not in _conv2d_op_cache:
#         _conv2d_op_cache[cache_key] = Conv2DOp(norm_stride, norm_dilation, norm_padding, groups)
#     op = _conv2d_op_cache[cache_key]
#     return op.forward(input_arr, filter_arr)


# def conv2d_transpose(
#     input_arr: Array, filter_arr: Array, stride=(1, 1),
#     dilation=(1, 1), padding=0, output_padding=0, groups=1
# ) -> Array:
#     """Applies a 2D transposed convolution."""
#     norm_stride = _normalize_tuple(stride, 2, "stride")
#     norm_dilation = _normalize_tuple(dilation, 2, "dilation")
#     norm_padding = _normalize_padding_arg(padding, "padding")
#     norm_output_padding = _normalize_tuple(output_padding, 2, "output_padding")

#     cache_key = (norm_stride, norm_dilation, norm_padding, norm_output_padding, groups)
#     if cache_key not in _conv2d_transpose_op_cache:
#         _conv2d_transpose_op_cache[cache_key] = Conv2DTransposeOp(
#             norm_stride, norm_dilation, norm_padding, norm_output_padding, groups
#         )
#     op = _conv2d_transpose_op_cache[cache_key]
#     return op.forward(input_arr, filter_arr)


# # ===----------------------------------------------------------------------=== #
# # Nabla 2025
# #
# # 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.
# # ===----------------------------------------------------------------------=== #

# """Numpy-based convolution utilities for eager execution."""

# from typing import Union

# import numpy as np


# def im2col(
#     input_data: np.ndarray,
#     filter_h: int,
#     filter_w: int,
#     stride: Union[int, tuple[int, int]] = 1,
#     dilation: Union[int, tuple[int, int]] = 1,
#     pad: Union[int, tuple[int, int]] = 0,
# ) -> np.ndarray:
#     """
#     Convert input data to column matrix for convolution.

#     Parameters:
#     -----------
#     input_data : ndarray
#         Input data with shape (N, C, H, W)
#     filter_h : int
#         Filter height
#     filter_w : int
#         Filter width
#     stride : int or tuple
#         Stride for convolution
#     dilation : int or tuple
#         Dilation for convolution
#     pad : int or tuple
#         Padding for input

#     Returns:
#     --------
#     col : ndarray
#         Column matrix with shape (N, C, filter_h, filter_w, out_h, out_w)
#     """
#     n, c, h, w = input_data.shape

#     # Handle stride and dilation as tuples
#     if isinstance(stride, int):
#         stride_h, stride_w = stride, stride
#     else:
#         stride_h, stride_w = stride

#     if isinstance(dilation, int):
#         dilation_h, dilation_w = dilation, dilation
#     else:
#         dilation_h, dilation_w = dilation

#     if isinstance(pad, int):
#         pad_h, pad_w = pad, pad
#     else:
#         pad_h, pad_w = pad

#     out_h = (h + 2 * pad_h - dilation_h * (filter_h - 1) - 1) // stride_h + 1
#     out_w = (w + 2 * pad_w - dilation_w * (filter_w - 1) - 1) // stride_w + 1

#     img = np.pad(
#         input_data, ((0, 0), (0, 0), (pad_h, pad_h), (pad_w, pad_w)), mode="constant"
#     )
#     col = np.ndarray((n, c, filter_h, filter_w, out_h, out_w), dtype=input_data.dtype)

#     for j in range(filter_h):
#         j_lim = j * dilation_h + stride_h * out_h
#         for i in range(filter_w):
#             i_lim = i * dilation_w + stride_w * out_w
#             col[:, :, j, i, :, :] = img[
#                 :,
#                 :,
#                 j * dilation_h : j_lim : stride_h,
#                 i * dilation_w : i_lim : stride_w,
#             ]

#     return col


# def col2im(
#     col: np.ndarray,
#     input_shape: tuple[int, int, int, int],
#     filter_h: int,
#     filter_w: int,
#     stride: Union[int, tuple[int, int]] = 1,
#     dilation: Union[int, tuple[int, int]] = 1,
#     pad: Union[int, tuple[int, int]] = 0,
# ) -> np.ndarray:
#     """
#     Convert column matrix back to input data shape.

#     Parameters:
#     -----------
#     col : ndarray
#         Column matrix with shape (N, C, filter_h, filter_w, out_h, out_w)
#     input_shape : tuple
#         Original input shape (N, C, H, W)
#     filter_h : int
#         Filter height
#     filter_w : int
#         Filter width
#     stride : int or tuple
#         Stride for convolution
#     dilation : int or tuple
#         Dilation for convolution
#     pad : int or tuple
#         Padding for input

#     Returns:
#     --------
#     img : ndarray
#         Reconstructed input data with shape (N, C, H, W)
#     """
#     n, c, h, w = input_shape

#     # Handle stride and dilation as tuples
#     if isinstance(stride, int):
#         stride_h, stride_w = stride, stride
#     else:
#         stride_h, stride_w = stride

#     if isinstance(dilation, int):
#         dilation_h, dilation_w = dilation, dilation
#     else:
#         dilation_h, dilation_w = dilation

#     if isinstance(pad, int):
#         pad_h, pad_w = pad, pad
#     else:
#         pad_h, pad_w = pad

#     out_h = (h + 2 * pad_h - dilation_h * (filter_h - 1) - 1) // stride_h + 1
#     out_w = (w + 2 * pad_w - dilation_w * (filter_w - 1) - 1) // stride_w + 1

#     img = np.zeros(
#         (n, c, h + 2 * pad_h + stride_h - 1, w + 2 * pad_w + stride_w - 1),
#         dtype=col.dtype,
#     )

#     for j in range(filter_h):
#         j_lim = j * dilation_h + stride_h * out_h
#         for i in range(filter_w):
#             i_lim = i * dilation_w + stride_w * out_w
#             img[
#                 :,
#                 :,
#                 j * dilation_h : j_lim : stride_h,
#                 i * dilation_w : i_lim : stride_w,
#             ] += col[:, :, j, i, :, :]

#     return img[:, :, pad_h : h + pad_h, pad_w : w + pad_w]


# def conv2d(input_data, filters, dilation=(1, 1), stride=(1, 1), padding=(0, 0)):
#     """
#     2D convolution using im2col method.

#     Parameters:
#     -----------
#     input_data : ndarray
#         Input data with shape (N, C_in, H, W)
#     filters : ndarray
#         Filters with shape (C_out, C_in, filter_h, filter_w)
#     dilation : tuple
#         Dilation factors (dilation_h, dilation_w)
#     stride : tuple
#         Stride values (stride_h, stride_w)
#     padding : tuple
#         Padding values (pad_h, pad_w)

#     Returns:
#     --------
#     output : ndarray
#         Convolution output with shape (N, C_out, out_h, out_w)
#     """
#     n, c_in, h, w = input_data.shape
#     c_out, c_in_f, filter_h, filter_w = filters.shape

#     assert c_in == c_in_f, f"Input channels {c_in} != filter input channels {c_in_f}"

#     # Calculate output dimensions
#     pad_h, pad_w = padding
#     stride_h, stride_w = stride
#     dilation_h, dilation_w = dilation

#     out_h = (h + 2 * pad_h - dilation_h * (filter_h - 1) - 1) // stride_h + 1
#     out_w = (w + 2 * pad_w - dilation_w * (filter_w - 1) - 1) // stride_w + 1

#     # Convert input to column matrix
#     col = im2col(input_data, filter_h, filter_w, stride, dilation, padding)
#     col = col.transpose(0, 4, 5, 1, 2, 3).reshape(n * out_h * out_w, -1)

#     # Reshape filters
#     w_col = filters.reshape(c_out, -1)

#     # Perform convolution via matrix multiplication
#     out = np.dot(col, w_col.T)
#     out = out.reshape(n, out_h, out_w, c_out).transpose(0, 3, 1, 2)

#     return out


# def transposed_conv2d(
#     input_data,
#     filters,
#     dilation=(1, 1),
#     stride=(1, 1),
#     padding=(0, 0),
#     output_padding=(0, 0),
# ):
#     """
#     2D transposed convolution using JAX-compatible algorithm.

#     JAX's conv_transpose implementation:
#     1. Upsample input by inserting (stride-1) zeros between elements
#     2. Apply regular convolution with effective padding

#     For transposed convolution, the effective padding is:
#     effective_pad = kernel_size - 1 - original_pad

#     Parameters:
#     -----------
#     input_data : ndarray
#         Input data with shape (N, C_in, H, W)
#     filters : ndarray
#         Filters with shape (C_out, C_in, filter_h, filter_w)
#     dilation : tuple
#         Dilation factors (dilation_h, dilation_w)
#     stride : tuple
#         Stride values (stride_h, stride_w)
#     padding : tuple
#         Original padding values (pad_h, pad_w) from the forward convolution
#     output_padding : tuple
#         Output padding values (out_pad_h, out_pad_w) - not used in JAX-compatible mode

#     Returns:
#     --------
#     output : ndarray
#         Transposed convolution output
#     """
#     n, c_in, h, w = input_data.shape
#     c_out, c_in_f, filter_h, filter_w = filters.shape

#     assert c_in == c_in_f, f"Input channels {c_in} != filter input channels {c_in_f}"

#     pad_h, pad_w = padding
#     stride_h, stride_w = stride
#     dilation_h, dilation_w = dilation

#     # Step 1: Upsample input by inserting (stride-1) zeros between elements
#     if stride_h > 1 or stride_w > 1:
#         # Calculate upsampled dimensions
#         upsampled_h = h + (h - 1) * (stride_h - 1)
#         upsampled_w = w + (w - 1) * (stride_w - 1)

#         # Create upsampled array filled with zeros
#         upsampled = np.zeros(
#             (n, c_in, upsampled_h, upsampled_w), dtype=input_data.dtype
#         )

#         # Insert original values at strided positions
#         upsampled[:, :, ::stride_h, ::stride_w] = input_data
#     else:
#         # No upsampling needed for stride=1
#         upsampled = input_data

#     # Step 2: Calculate effective padding for transposed convolution
#     # For transposed conv, if original conv had padding P and kernel size K,
#     # the effective padding for the underlying regular conv is (K-1-P)
#     effective_pad_h = filter_h - 1 - pad_h
#     effective_pad_w = filter_w - 1 - pad_w

#     # Step 3: Apply regular convolution with effective padding
#     # Use stride=1 since upsampling already handled the stride effect
#     result = conv2d(
#         upsampled,
#         filters,
#         dilation=dilation,
#         stride=(1, 1),
#         padding=(effective_pad_h, effective_pad_w),
#     )

#     # Step 4: Apply output_padding if specified
#     # Output padding adds zeros to the right and bottom of the output
#     out_pad_h, out_pad_w = output_padding
#     if out_pad_h > 0 or out_pad_w > 0:
#         n, c_out, h_out, w_out = result.shape
#         padded_result = np.zeros(
#             (n, c_out, h_out + out_pad_h, w_out + out_pad_w), dtype=result.dtype
#         )
#         padded_result[:, :, :h_out, :w_out] = result
#         result = padded_result

#     return result
