# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE

from __future__ import annotations

from abc import abstractmethod

import jax

import awkward as ak
from awkward import _reducers
from awkward._nplikes.array_like import maybe_materialize
from awkward._nplikes.numpy_like import NumpyMetadata
from awkward._nplikes.shape import ShapeItem
from awkward._reducers import Reducer
from awkward._typing import Final, Self, TypeVar

np = NumpyMetadata.instance()


_overloads: dict[type[Reducer], type[JAXReducer]] = {}


R = TypeVar("R", bound=Reducer)


def overloads(cls: type[Reducer]):
    def registrar(new_cls: type[R]) -> type[R]:
        _overloads[cls] = new_cls
        return new_cls

    return registrar


class JAXReducer(Reducer):
    @classmethod
    @abstractmethod
    def from_kernel_reducer(cls, reducer: Reducer) -> Self:
        raise NotImplementedError


# Offsets-pipeline note:
# `apply()` and these helpers no longer take a `parents` argument. Where JAX's
# `segment_*` ops still need a per-element segment-id array we synthesise it
# from `offsets` via `_parents_from_offsets`. For the positional-correction
# helpers, the bin index is simply the output position `k` (because
# argmin/argmax always pick an element that belongs to bin `k`), so we use
# `arange(outlength)` directly and avoid the parents indirection — mirroring
# the simplification already made in the CPU/CUDA reduce_adjust_starts kernels.


def _parents_from_offsets(offsets, outlength):
    """Materialise a per-element segment-id array from `offsets`.

    `offsets` has length `outlength + 1`. The returned array has length
    `offsets[outlength]` and contains, at position `i`, the bin index `k` such
    that `offsets[k] <= i < offsets[k+1]`.
    """
    outlen_i = int(outlength)
    if outlen_i == 0:
        return jax.numpy.zeros(0, dtype=jax.numpy.int64)
    counts = offsets[1 : outlen_i + 1] - offsets[:outlen_i]
    return jax.numpy.repeat(
        jax.numpy.arange(outlen_i, dtype=jax.numpy.int64),
        counts,
        total_repeat_length=int(offsets[outlen_i]),
    )


def awkward_JAXNumpyArray_reduce_adjust_starts_64(toptr, outlength, starts):
    if outlength == 0:
        return toptr

    identity = jax.numpy.iinfo(jax.numpy.int64).max
    valid = toptr[:outlength] != identity
    # parents[toptr[k]] == k by construction (argmin/argmax always picks an
    # element of bin k), so the bin index is just `k`.
    adjustments = starts[: int(outlength)]
    updated = jax.numpy.where(valid, toptr[:outlength] - adjustments, toptr[:outlength])

    return toptr.at[:outlength].set(updated)


def awkward_JAXNumpyArray_reduce_adjust_starts_shifts_64(
    toptr, outlength, starts, shifts
):
    if outlength == 0:
        return toptr

    identity = jax.numpy.iinfo(jax.numpy.int64).max
    valid = toptr[:outlength] != identity
    safe_sub_toptr = jax.numpy.where(valid, toptr[:outlength], 0)
    # parents[toptr[k]] == k, so we look up `starts[k]` directly. `shifts` is
    # per-element of the flat input, so it's still indexed by toptr[k].
    delta = shifts[safe_sub_toptr] - starts[: int(outlength)]
    updated = jax.numpy.where(valid, toptr[:outlength] + delta, toptr[:outlength])

    return toptr.at[:outlength].set(updated)


def apply_positional_corrections(
    reduced: ak.contents.NumpyArray,
    offsets: ak.index.Index,
    starts: ak.index.Index,
    shifts: ak.index.Index | None,
) -> ak._nplikes.ArrayLike:
    if shifts is None:
        assert (
            offsets.nplike is reduced.backend.nplike
            and starts.nplike is reduced.backend.nplike
        )
        return awkward_JAXNumpyArray_reduce_adjust_starts_64(
            reduced.data, reduced.length, starts.data
        )

    else:
        assert (
            offsets.nplike is reduced.backend.nplike
            and starts.nplike is reduced.backend.nplike
            and shifts.nplike is reduced.backend.nplike
        )
        return awkward_JAXNumpyArray_reduce_adjust_starts_shifts_64(
            reduced.data,
            reduced.length,
            starts.data,
            shifts.data,
        )


def segment_argmin(data, segment_ids, num_segments):
    """
    Applies a segmented argmin-style reduction.

    Parameters:
        data: jax.numpy.ndarray — the values to reduce.
        segment_ids: same shape as data — indicates segment groupings.
        num_segments: int — total number of segments.

    Returns:
        jax.numpy.ndarray — indices of min within each segment.
    """
    indices = jax.numpy.arange(data.shape[0])

    # Find the minimum value in each segment
    min_vals = jax.ops.segment_min(data, segment_ids, num_segments=num_segments)

    # Find where the data equals the minimum value in each segment
    is_min = data == min_vals[segment_ids]

    # Mask the indices where data matches the minimum value
    masked_indices = jax.numpy.where(is_min, indices, data.shape[0])

    # Return the index of the first minimum value in each segment
    return jax.ops.segment_min(masked_indices, segment_ids, num_segments=num_segments)


@overloads(_reducers.ArgMin)
class ArgMin(JAXReducer):
    name: Final = "argmin"
    needs_position: Final = True
    preferred_dtype: Final = np.int64

    @classmethod
    def from_kernel_reducer(cls, reducer: Reducer) -> Self:
        assert isinstance(reducer, _reducers.ArgMin)
        return cls()

    @classmethod
    def _return_dtype(cls, given_dtype):
        return np.int64

    def apply(
        self,
        array: ak.contents.NumpyArray,
        offsets: ak.index.Index,
        starts: ak.index.Index,
        shifts: ak.index.Index | None,
        outlength: ShapeItem,
    ) -> ak.contents.NumpyArray:
        assert isinstance(array, ak.contents.NumpyArray)
        parents_data = _parents_from_offsets(offsets.data, outlength)
        result = segment_argmin(*maybe_materialize(array.data, parents_data), outlength)
        result = jax.numpy.asarray(result, dtype=self.preferred_dtype)
        result_array = ak.contents.NumpyArray(result, backend=array.backend)
        corrected_data = apply_positional_corrections(
            result_array, offsets, starts, shifts
        )
        return ak.contents.NumpyArray(corrected_data, backend=array.backend)


def segment_argmax(data, segment_ids, num_segments):
    """
    Applies a segmented argmax-style reduction.

    Parameters:
        data: jax.numpy.ndarray — the values to reduce.
        segment_ids: same shape as data — indicates segment groupings.
        num_segments: int — total number of segments.

    Returns:
        jax.numpy.ndarray — indices of max within each segment.
    """
    indices = jax.numpy.arange(data.shape[0])

    # Find the maximum value in each segment
    max_vals = jax.ops.segment_max(data, segment_ids, num_segments=num_segments)

    # Find where the data equals the maximum value in each segment
    is_max = data == max_vals[segment_ids]

    # Mask the indices where data matches the maximum value
    masked_indices = jax.numpy.where(is_max, indices, data.shape[0])

    # Return the index of the first maximum value in each segment
    return jax.ops.segment_min(masked_indices, segment_ids, num_segments=num_segments)


@overloads(_reducers.ArgMax)
class ArgMax(JAXReducer):
    name: Final = "argmax"
    needs_position: Final = True
    preferred_dtype: Final = np.int64

    @classmethod
    def from_kernel_reducer(cls, reducer: Reducer) -> Self:
        assert isinstance(reducer, _reducers.ArgMax)
        return cls()

    @classmethod
    def _return_dtype(cls, given_dtype):
        return np.int64

    def apply(
        self,
        array: ak.contents.NumpyArray,
        offsets: ak.index.Index,
        starts: ak.index.Index,
        shifts: ak.index.Index | None,
        outlength: ShapeItem,
    ) -> ak.contents.NumpyArray:
        assert isinstance(array, ak.contents.NumpyArray)
        parents_data = _parents_from_offsets(offsets.data, outlength)
        result = segment_argmax(*maybe_materialize(array.data, parents_data), outlength)
        result = jax.numpy.asarray(result, dtype=self.preferred_dtype)
        result_array = ak.contents.NumpyArray(result, backend=array.backend)
        corrected_data = apply_positional_corrections(
            result_array, offsets, starts, shifts
        )
        return ak.contents.NumpyArray(corrected_data, backend=array.backend)


@overloads(_reducers.Count)
class Count(JAXReducer):
    name: Final = "count"
    preferred_dtype: Final = np.int64
    needs_position: Final = False

    @classmethod
    def from_kernel_reducer(cls, reducer: Reducer) -> Self:
        assert isinstance(reducer, _reducers.Count)
        return cls()

    @classmethod
    def _return_dtype(cls, given_dtype):
        return np.int64

    def apply(
        self,
        array: ak.contents.NumpyArray,
        offsets: ak.index.Index,
        starts: ak.index.Index,
        shifts: ak.index.Index | None,
        outlength: ShapeItem,
    ) -> ak.contents.NumpyArray:
        assert isinstance(array, ak.contents.NumpyArray)
        parents_data = _parents_from_offsets(offsets.data, outlength)
        result = jax.numpy.ones_like(
            *maybe_materialize(array.data), dtype=self.preferred_dtype
        )
        result = jax.ops.segment_sum(
            result, *maybe_materialize(parents_data), outlength
        )
        result = jax.numpy.asarray(result, dtype=self.preferred_dtype)

        if np.issubdtype(array.dtype, np.complexfloating):
            return ak.contents.NumpyArray(
                result.view(array.dtype), backend=array.backend
            )
        else:
            return ak.contents.NumpyArray(result, backend=array.backend)


def segment_count_nonzero(data, segment_ids, num_segments):
    """
    Counts the number of non-zero elements in `data` per segment.

    Parameters:
        data: jax.numpy.ndarray — input values to count.
        segment_ids: jax.numpy.ndarray — same shape as data, segment assignment.
        num_segments: int — total number of segments.

    Returns:
        jax.numpy.ndarray — count of non-zero values per segment.
    """
    # Create a binary mask where non-zero entries become 1
    nonzero_mask = jax.numpy.where(data != 0, 1, 0)

    # Sum the mask using segment_sum to count per segment
    return jax.ops.segment_sum(nonzero_mask, segment_ids, num_segments=num_segments)


@overloads(_reducers.CountNonzero)
class CountNonzero(JAXReducer):
    name: Final = "count_nonzero"
    preferred_dtype: Final = np.int64
    needs_position: Final = False

    @classmethod
    def from_kernel_reducer(cls, reducer: Reducer) -> Self:
        assert isinstance(reducer, _reducers.CountNonzero)
        return cls()

    @classmethod
    def _return_dtype(cls, given_dtype):
        return np.int64

    def apply(
        self,
        array: ak.contents.NumpyArray,
        offsets: ak.index.Index,
        starts: ak.index.Index,
        shifts: ak.index.Index | None,
        outlength: ShapeItem,
    ) -> ak.contents.NumpyArray:
        assert isinstance(array, ak.contents.NumpyArray)
        parents_data = _parents_from_offsets(offsets.data, outlength)
        result = segment_count_nonzero(
            *maybe_materialize(array.data, parents_data), outlength
        )
        result = jax.numpy.asarray(result, dtype=self.preferred_dtype)

        return ak.contents.NumpyArray(result, backend=array.backend)


@overloads(_reducers.Sum)
class Sum(JAXReducer):
    name: Final = "sum"
    preferred_dtype: Final = np.float64
    needs_position: Final = False

    def __init__(self, dtype=None):
        # Forced accumulator dtype (mirrors awkward._reducers.Sum); honoured so
        # ak.mean/ak.var's float64 accumulation isn't silently dropped on jax.
        self._dtype = None if dtype is None else np.dtype(dtype)

    @classmethod
    def from_kernel_reducer(cls, reducer: Reducer) -> Self:
        assert isinstance(reducer, _reducers.Sum)
        return cls(reducer._dtype)

    def axis_none_reducer(self) -> AxisNoneSum:
        return AxisNoneSum(self._dtype)

    def apply(
        self,
        array: ak.contents.NumpyArray,
        offsets: ak.index.Index,
        starts: ak.index.Index,
        shifts: ak.index.Index | None,
        outlength: ShapeItem,
    ) -> ak.contents.NumpyArray:
        assert isinstance(array, ak.contents.NumpyArray)
        if array.dtype.kind == "M":
            raise TypeError(f"cannot compute the sum (ak.sum) of {array.dtype!r}")

        if self._dtype is not None and array.dtype.kind not in "cmM":
            input_array = array.data.astype(self._dtype)
        elif array.dtype.kind == "b":
            input_array = array.data.astype(np.int64)
        else:
            input_array = array.data
        parents_data = _parents_from_offsets(offsets.data, outlength)
        result = jax.ops.segment_sum(
            *maybe_materialize(input_array, parents_data), outlength
        )

        if array.dtype.kind == "m":
            return ak.contents.NumpyArray(
                array.backend.nplike.asarray(result, dtype=array.dtype)
            )
        elif np.issubdtype(array.dtype, np.complexfloating):
            return ak.contents.NumpyArray(result.view(array.dtype))
        else:
            return ak.contents.NumpyArray(result, backend=array.backend)


@overloads(_reducers.SumOfSquares)
class SumOfSquares(JAXReducer):
    name: Final = "sumofsquares"
    preferred_dtype: Final = np.float64
    needs_position: Final = False

    @classmethod
    def from_kernel_reducer(cls, reducer: Reducer) -> Self:
        assert isinstance(reducer, _reducers.SumOfSquares)
        return cls()

    def apply(
        self,
        array: ak.contents.NumpyArray,
        offsets: ak.index.Index,
        starts: ak.index.Index,
        shifts: ak.index.Index | None,
        outlength: ShapeItem,
    ) -> ak.contents.NumpyArray:
        assert isinstance(array, ak.contents.NumpyArray)
        if array.dtype.kind == "c":
            raise TypeError(
                f"cannot compute the sum-of-squares (ak.var/ak.std) of {array.dtype!r}"
            )
        # Square in float64, then segment_sum (differentiable, no x*x overflow).
        data = array.data.astype(np.float64)
        squared = data * data
        parents_data = _parents_from_offsets(offsets.data, outlength)
        result = jax.ops.segment_sum(
            *maybe_materialize(squared, parents_data), outlength
        )
        return ak.contents.NumpyArray(result, backend=array.backend)


@overloads(_reducers.SumOfPowers)
class SumOfPowers(JAXReducer):
    name: Final = "sumofpowers"
    preferred_dtype: Final = np.float64
    needs_position: Final = False

    def __init__(self, n=2):
        self._n = int(n)

    @classmethod
    def from_kernel_reducer(cls, reducer: Reducer) -> Self:
        assert isinstance(reducer, _reducers.SumOfPowers)
        return cls(reducer._n)

    def apply(
        self,
        array: ak.contents.NumpyArray,
        offsets: ak.index.Index,
        starts: ak.index.Index,
        shifts: ak.index.Index | None,
        outlength: ShapeItem,
    ) -> ak.contents.NumpyArray:
        assert isinstance(array, ak.contents.NumpyArray)
        if array.dtype.kind == "c":
            raise TypeError(
                f"cannot compute the sum-of-powers (ak.moment) of {array.dtype!r}"
            )
        # x**n in float64, then segment_sum (differentiable, no x**n overflow).
        powered = array.data.astype(np.float64) ** self._n
        parents_data = _parents_from_offsets(offsets.data, outlength)
        result = jax.ops.segment_sum(
            *maybe_materialize(powered, parents_data), outlength
        )
        return ak.contents.NumpyArray(result, backend=array.backend)


@overloads(_reducers.AxisNoneSum)
class AxisNoneSum(JAXReducer):
    name: Final = "sum"
    preferred_dtype: Final = np.float64
    needs_position: Final = False

    def __init__(self, dtype=None):
        self._dtype = None if dtype is None else np.dtype(dtype)

    @classmethod
    def from_kernel_reducer(cls, reducer: Reducer) -> Self:
        assert isinstance(reducer, _reducers.AxisNoneSum)
        return cls(reducer._dtype)

    def apply(
        self,
        array: ak.contents.NumpyArray,
        offsets: ak.index.Index,
        starts: ak.index.Index,
        shifts: ak.index.Index | None,
        outlength: ShapeItem,
    ) -> ak.contents.NumpyArray:
        assert isinstance(array, ak.contents.NumpyArray)
        if array.dtype.kind == "M":
            raise TypeError(f"cannot compute the sum (ak.sum) of {array.dtype!r}")

        data = maybe_materialize(array.data)[0]

        if self._dtype is not None and array.dtype.kind not in "cmM":
            data = data.astype(self._dtype)
        elif array.dtype.kind == "b":
            data = data.astype(np.int64)

        result_scalar = jax.numpy.sum(data)
        result = jax.numpy.reshape(result_scalar, (1,))

        if array.dtype.kind == "m":
            return ak.contents.NumpyArray(
                array.backend.nplike.asarray(result, dtype=array.dtype),
                backend=array.backend,
            )
        elif np.issubdtype(array.dtype, np.complexfloating):
            return ak.contents.NumpyArray(
                array.backend.nplike.asarray(result, dtype=array.dtype),
                backend=array.backend,
            )
        else:
            return ak.contents.NumpyArray(result, backend=array.backend)


def segment_prod_with_negatives(data, segment_ids, num_segments):
    """
    Computes the product of elements in each segment, handling negatives and booleans.
    Parameters:
        data: jax.numpy.ndarray — input values to reduce.
        segment_ids: jax.numpy.ndarray — same shape as data, segment assignment.
        num_segments: int — total number of segments.
    Returns:
        jax.numpy.ndarray — product of values per segment.
    """
    # Handle boolean arrays
    if data.dtype == jax.numpy.bool_:
        return jax.ops.segment_min(
            data.astype(jax.numpy.int32), segment_ids, num_segments
        )

    # For numeric arrays, handle negative values and zeros
    # Track zeros to set product to zero if any segment has zeros
    is_zero = data == 0
    has_zeros = (
        jax.ops.segment_sum(is_zero.astype(jax.numpy.int32), segment_ids, num_segments)
        > 0
    )

    # Track signs to determine final sign of product
    is_negative = data < 0
    neg_count = jax.ops.segment_sum(
        is_negative.astype(jax.numpy.int32), segment_ids, num_segments
    )
    sign_products = 1 - 2 * (
        neg_count % 2
    )  # +1 for even negatives, -1 for odd negatives

    # Calculate product of absolute values in log space
    log_abs = jax.numpy.log(jax.numpy.where(is_zero, 1.0, jax.numpy.abs(data)))
    log_products = jax.ops.segment_sum(
        jax.numpy.where(is_zero, 0.0, log_abs), segment_ids, num_segments
    )
    abs_products = jax.numpy.exp(log_products)

    # Apply zeros and signs
    product = jax.numpy.where(has_zeros, 0.0, sign_products * abs_products)
    # floating point accuracy doesn't let us directly cast to integers
    if np.issubdtype(data.dtype, np.integer):
        result = jax.numpy.round(product).astype(data.dtype)
    else:
        result = product
    return result


@overloads(_reducers.Prod)
class Prod(JAXReducer):
    name: Final = "prod"
    preferred_dtype: Final = np.float64
    needs_position: Final = False

    @classmethod
    def from_kernel_reducer(cls, reducer: Reducer) -> Self:
        assert isinstance(reducer, _reducers.Prod)
        return cls()

    def apply(
        self,
        array: ak.contents.NumpyArray,
        offsets: ak.index.Index,
        starts: ak.index.Index,
        shifts: ak.index.Index | None,
        outlength: ShapeItem,
    ) -> ak.contents.NumpyArray:
        assert isinstance(array, ak.contents.NumpyArray)
        # See issue https://github.com/google/jax/issues/9296
        parents_data = _parents_from_offsets(offsets.data, outlength)
        result = segment_prod_with_negatives(
            *maybe_materialize(array.data, parents_data), outlength
        )

        if np.issubdtype(array.dtype, np.complexfloating):
            return ak.contents.NumpyArray(
                result.view(array.dtype), backend=array.backend
            )
        else:
            return ak.contents.NumpyArray(result, backend=array.backend)


@overloads(_reducers.Any)
class Any(JAXReducer):
    name: Final = "any"
    preferred_dtype: Final = np.bool_
    needs_position: Final = False

    @classmethod
    def from_kernel_reducer(cls, reducer: Reducer) -> Self:
        assert isinstance(reducer, _reducers.Any)
        return cls()

    @classmethod
    def _return_dtype(cls, given_dtype):
        return np.bool_

    @staticmethod
    def _max_initial(initial, type):
        if initial is None:
            if type in (
                np.int8,
                np.int16,
                np.int32,
                np.int64,
                np.uint8,
                np.uint16,
                np.uint32,
                np.uint64,
            ):
                return np.iinfo(type).min
            else:
                return -np.inf

        return initial

    def apply(
        self,
        array: ak.contents.NumpyArray,
        offsets: ak.index.Index,
        starts: ak.index.Index,
        shifts: ak.index.Index | None,
        outlength: ShapeItem,
    ) -> ak.contents.NumpyArray:
        assert isinstance(array, ak.contents.NumpyArray)
        parents_data = _parents_from_offsets(offsets.data, outlength)
        result = jax.ops.segment_max(
            *maybe_materialize(array.data, parents_data), outlength
        )
        if array.dtype is not np.dtype(bool):
            result = result.at[result == 0].set(self._max_initial(None, array.dtype))
            result = result > self._max_initial(None, array.dtype)
        result = jax.numpy.asarray(result, dtype=bool)

        return ak.contents.NumpyArray(result, backend=array.backend)


@overloads(_reducers.All)
class All(JAXReducer):
    name: Final = "all"
    preferred_dtype: Final = np.bool_
    needs_position: Final = False

    @classmethod
    def from_kernel_reducer(cls, reducer: Reducer) -> Self:
        assert isinstance(reducer, _reducers.All)
        return cls()

    @classmethod
    def _return_dtype(cls, given_dtype):
        return np.bool_

    def apply(
        self,
        array: ak.contents.NumpyArray,
        offsets: ak.index.Index,
        starts: ak.index.Index,
        shifts: ak.index.Index | None,
        outlength: ShapeItem,
    ) -> ak.contents.NumpyArray:
        assert isinstance(array, ak.contents.NumpyArray)
        parents_data = _parents_from_offsets(offsets.data, outlength)
        result = jax.ops.segment_min(
            *maybe_materialize(array.data, parents_data), outlength
        )
        result = jax.numpy.asarray(result, dtype=bool)

        return ak.contents.NumpyArray(result, backend=array.backend)


@overloads(_reducers.Min)
class Min(JAXReducer):
    name: Final = "min"
    preferred_dtype: Final = np.float64
    needs_position: Final = False

    def __init__(self, initial):
        self._initial = initial

    @property
    def initial(self):
        return self._initial

    @classmethod
    def from_kernel_reducer(cls, reducer: Reducer) -> Self:
        assert isinstance(reducer, _reducers.Min)
        return cls(reducer.initial)

    def axis_none_reducer(self) -> AxisNoneMin:
        return AxisNoneMin(self._initial)

    @staticmethod
    def _min_initial(initial, type):
        if initial is None:
            if type in (
                np.int8,
                np.int16,
                np.int32,
                np.int64,
                np.uint8,
                np.uint16,
                np.uint32,
                np.uint64,
            ):
                return np.iinfo(type).max
            else:
                return np.inf

        return initial

    def apply(
        self,
        array: ak.contents.NumpyArray,
        offsets: ak.index.Index,
        starts: ak.index.Index,
        shifts: ak.index.Index | None,
        outlength: ShapeItem,
    ) -> ak.contents.NumpyArray:
        assert isinstance(array, ak.contents.NumpyArray)

        parents_data = _parents_from_offsets(offsets.data, outlength)
        result = jax.ops.segment_min(
            *maybe_materialize(array.data, parents_data), outlength
        )
        result = jax.numpy.minimum(result, self._min_initial(self.initial, array.dtype))

        if np.issubdtype(array.dtype, np.complexfloating):
            return ak.contents.NumpyArray(
                array.backend.nplike.asarray(
                    result.view(array.dtype), dtype=array.dtype
                ),
                backend=array.backend,
            )
        else:
            return ak.contents.NumpyArray(result, backend=array.backend)


@overloads(_reducers.AxisNoneMin)
class AxisNoneMin(JAXReducer):
    name: Final = "min"
    preferred_dtype: Final = np.float64
    needs_position: Final = False

    def __init__(self, initial):
        self._initial = initial

    @property
    def initial(self):
        return self._initial

    @classmethod
    def from_kernel_reducer(cls, reducer: Reducer) -> Self:
        assert isinstance(reducer, _reducers.AxisNoneMin)
        return cls(reducer.initial)

    def apply(
        self,
        array: ak.contents.NumpyArray,
        offsets: ak.index.Index,
        starts: ak.index.Index,
        shifts: ak.index.Index | None,
        outlength: ShapeItem,
    ) -> ak.contents.NumpyArray:
        assert isinstance(array, ak.contents.NumpyArray)

        data = maybe_materialize(array.data)[0]

        result_scalar = jax.numpy.min(
            data, initial=Min._min_initial(self.initial, array.dtype)
        )
        result = jax.numpy.reshape(result_scalar, (1,))
        return ak.contents.NumpyArray(result, backend=array.backend)


@overloads(_reducers.Max)
class Max(JAXReducer):
    name: Final = "max"
    preferred_dtype: Final = np.float64
    needs_position: Final = False

    def __init__(self, initial):
        self._initial = initial

    @property
    def initial(self):
        return self._initial

    @classmethod
    def from_kernel_reducer(cls, reducer: Reducer) -> Self:
        assert isinstance(reducer, _reducers.Max)
        return cls(reducer.initial)

    def axis_none_reducer(self) -> AxisNoneMax:
        return AxisNoneMax(self._initial)

    @staticmethod
    def _max_initial(initial, type):
        if initial is None:
            if type in (
                np.int8,
                np.int16,
                np.int32,
                np.int64,
                np.uint8,
                np.uint16,
                np.uint32,
                np.uint64,
            ):
                return np.iinfo(type).min
            else:
                return -np.inf

        return initial

    def apply(
        self,
        array: ak.contents.NumpyArray,
        offsets: ak.index.Index,
        starts: ak.index.Index,
        shifts: ak.index.Index | None,
        outlength: ShapeItem,
    ) -> ak.contents.NumpyArray:
        assert isinstance(array, ak.contents.NumpyArray)

        parents_data = _parents_from_offsets(offsets.data, outlength)
        result = jax.ops.segment_max(
            *maybe_materialize(array.data, parents_data), outlength
        )
        result = jax.numpy.maximum(result, self._max_initial(self.initial, array.dtype))

        if np.issubdtype(array.dtype, np.complexfloating):
            return ak.contents.NumpyArray(
                array.backend.nplike.asarray(
                    result.view(array.dtype), dtype=array.dtype
                ),
                backend=array.backend,
            )
        else:
            return ak.contents.NumpyArray(result, backend=array.backend)


@overloads(_reducers.AxisNoneMax)
class AxisNoneMax(JAXReducer):
    name: Final = "max"
    preferred_dtype: Final = np.float64
    needs_position: Final = False

    def __init__(self, initial):
        self._initial = initial

    @property
    def initial(self):
        return self._initial

    @classmethod
    def from_kernel_reducer(cls, reducer: Reducer) -> Self:
        assert isinstance(reducer, _reducers.AxisNoneMax)
        return cls(reducer.initial)

    def apply(
        self,
        array: ak.contents.NumpyArray,
        offsets: ak.index.Index,
        starts: ak.index.Index,
        shifts: ak.index.Index | None,
        outlength: ShapeItem,
    ) -> ak.contents.NumpyArray:
        assert isinstance(array, ak.contents.NumpyArray)

        data = maybe_materialize(array.data)[0]

        result_scalar = jax.numpy.max(
            data, initial=Max._max_initial(self.initial, array.dtype)
        )
        result = jax.numpy.reshape(result_scalar, (1,))
        return ak.contents.NumpyArray(result, backend=array.backend)


def get_jax_reducer(reducer: Reducer) -> Reducer:
    return _overloads[type(reducer)].from_kernel_reducer(reducer)
