# Copyright 2026 The JAX Authors.
#
# 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
#
#     https://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 function implementations as hijax primitives."""

from collections.abc import Callable
import functools
import operator
import string
from typing import Any

import numpy as np
import opt_einsum
from opt_einsum.parser import find_output_shape, parse_einsum_input

from jax._src import ad_util
from jax._src import api
from jax._src import core
from jax._src import dtypes
from jax._src import numpy as jnp
from jax._src.hijax import (
    linearize_from_jvp,
    vjp_from_jvp,
    HiPrim,
)
from jax._src.partition_spec import PartitionSpec as P
from jax._src.interpreters import ad
from jax._src.lax import control_flow
from jax._src.lax import lax
from jax._src.numpy import util
from jax._src.numpy.einsum import _einsum as einsum_impl, Unoptimized
from jax._src.typing import Array, ArrayLike, DTypeLike


class SearchSorted(HiPrim):
  """HiJAX primitive for binary search."""
  valid_methods = ("compare_all", "scan", "scan_unrolled", "sort")

  side: str
  dimension: int
  batch_dims: int
  method: str
  out_dtype: np.dtype

  def __init__(
      self,
      sorted_arr_aval: core.ShapedArray,
      query_aval: core.ShapedArray,
      *,
      side: str,
      dimension: int,
      batch_dims: int,
      method: str,
      out_dtype: np.dtype):
    batch_dims = operator.index(batch_dims)
    if batch_dims < 0 or batch_dims >= sorted_arr_aval.ndim:
      raise ValueError(
          f"batch_dims={batch_dims} must be in range [0, {sorted_arr_aval.ndim})"
      )
    dimension = operator.index(dimension)
    if not batch_dims <= dimension < sorted_arr_aval.ndim:
      raise ValueError(
          f"dimension={dimension} must be in range [{batch_dims},"
          f" {sorted_arr_aval.ndim})"
      )
    if sorted_arr_aval.dtype != query_aval.dtype:
      raise ValueError(
          "dtypes of sorted_arr and query must match; got "
          f"{sorted_arr_aval.dtype} and {query_aval.dtype}"
      )
    if side not in ["left", "right"]:
      raise ValueError(
          f"invalid argument side={side!r}, expected 'left' or 'right'"
      )
    if method not in self.valid_methods:
      raise ValueError(
          f"invalid argument {method=}, expected one of {list(self.valid_methods)}"
      )
    if sorted_arr_aval.shape[:batch_dims] != query_aval.shape[:batch_dims]:
      raise ValueError(
          "batch dimension sizes must match; got"
          f" {sorted_arr_aval.shape[:batch_dims]} != {query_aval.shape[:batch_dims]}"
      )
    if not dtypes.issubdtype(out_dtype, np.integer):
      raise ValueError(f"out_dtype should be an integer type; got {out_dtype}")
    # Attempt this here to catch overflow errors early.
    out_dtype.type(sorted_arr_aval.shape[dimension])
    self.in_avals = (sorted_arr_aval, query_aval)
    self.out_aval = core.typeof(api.eval_shape(
      functools.partial(_searchsorted_impl,
        dimension=dimension, batch_dims=batch_dims, side=side,
        dtype=out_dtype, method=method),
        sorted_arr_aval, query_aval))
    self.params = dict(
      side=side,
      dimension=dimension,
      batch_dims=batch_dims,
      method=method,
    )
    super().__init__()

  def expand(self, sorted_arr: ArrayLike, query: ArrayLike) -> Array:  # pyrefly: ignore[bad-override]
    return _searchsorted_impl(
      sorted_arr, query,
      dimension=self.dimension,
      batch_dims=self.batch_dims,
      side=self.side,
      dtype=self.out_aval.dtype,
      method=self.method)

  def batch(
      self,
      axis_data: Any,
      args: tuple[Array, Array],
      dims: tuple[int | None, int | None]
  ) -> tuple[Array, int | None]:
    del axis_data  # unused
    sorted_arr, query = args
    batch_dims = self.batch_dims
    dimension = self.dimension

    if dims[0] is None and dims[1] is None:
      # Neither array is batched
      out_bdim = None
    elif dims[1] is None:
      # Only sorted_arr is batched
      assert dims[0] is not None  # for type checker
      sorted_arr = jnp.moveaxis(sorted_arr, dims[0], batch_dims)
      dimension += 1
      out_bdim = batch_dims
    elif dims[0] is None:
      # Only query is batched
      assert dims[1] is not None  # for type checker
      query = jnp.moveaxis(query, dims[1], batch_dims)
      out_bdim = sorted_arr.ndim - 1
    else:
      # Both are batched
      sorted_arr = jnp.moveaxis(sorted_arr, dims[0], 0)
      query = jnp.moveaxis(query, dims[1], 0)
      batch_dims += 1
      dimension += 1
      out_bdim = 0

    batched_prim = SearchSorted(
      core.typeof(sorted_arr),
      core.typeof(query),
      side=self.side,
      dimension=dimension,
      batch_dims=batch_dims,
      method=self.method,
      out_dtype=self.out_aval.dtype,
    )
    return batched_prim(sorted_arr, query), out_bdim

  def jvp(self, primals: tuple[Array, Array], tangents: Any) -> tuple[Array | np.ndarray, Array | np.ndarray]:
    del tangents  # unused
    primal_out = self(*primals)
    return primal_out, np.empty(primal_out.shape, dtype=dtypes.float0)

  def vjp_fwd(self, nzs_in: Any, *args: Array) -> tuple[Array, None]:
    return (self(*args), None)

  def vjp_bwd_retval(self, res: Any, g: Any):
    return (ad_util.zeros_like_aval(self.in_avals[0]),
            ad_util.zeros_like_aval(self.in_avals[1]))

  lin, linearized = linearize_from_jvp


class Nonzero(HiPrim):
  """HiJAX primitive for nonzero."""

  size: int
  axes: tuple[int, ...]
  out_dtype: np.dtype

  def __init__(
      self,
      a_aval: core.ShapedArray,
      *fill_value_avals: core.ShapedArray,
      size: int,
      axes: tuple[int, ...],
      out_dtype: np.dtype):
    if not core.is_symbolic_dim(size):
      size = operator.index(size)
    if fill_value_avals and len(fill_value_avals) != len(axes):
      raise ValueError(f"Expected {len(axes)} fill values, got {len(fill_value_avals)}")
    if any(fv.dtype != out_dtype for fv in fill_value_avals):
      raise ValueError(f"Expected fill values to have dtype {out_dtype}, got {fill_value_avals}")

    self.in_avals = (a_aval, *fill_value_avals)
    out_dtype = dtypes._maybe_canonicalize_explicit_dtype(out_dtype, "nonzero")
    out_shape = self._shape_rule(a_aval, *fill_value_avals, size=size, axes=axes)
    out_sharding = self._sharding_rule(a_aval, *fill_value_avals, axes=axes)
    out_vma = core.standard_vma_rule("nonzero", *self.in_avals)
    out_aval = a_aval.update(
        shape=out_shape,
        dtype=out_dtype,
        weak_type=False,
        sharding=out_sharding,
        manual_axis_type=a_aval.mat.update(varying=out_vma),
    )
    self.out_aval = (out_aval,) * len(axes)

    self.params = dict(
        size=size,
        axes=axes,
        out_dtype=out_dtype,
    )
    super().__init__()

  @staticmethod
  def _shape_rule(
      a_aval: core.ShapedArray,
      *fill_value_avals: core.ShapedArray,
      size: int,
      axes: tuple[int, ...],
  ) -> core.Shape:
    if not core.is_symbolic_dim(size) and size < 0:
      raise ValueError(f"size must be a positive integer; got {size=}")
    if not all(0 <= ax < a_aval.ndim for ax in axes):
      raise ValueError(f"axes out of range for array with {a_aval.ndim} dimensions:  {axes=}")
    if len(axes) != len(set(axes)):
      raise ValueError(f"duplicate axes are not allowed: {axes=}")
    batch_shape = tuple(
        s for i, s in enumerate(a_aval.shape) if i not in axes
    )
    for fv_aval in fill_value_avals:
      try:
        broadcasted = lax.broadcast_shapes(fv_aval.shape, batch_shape)
      except ValueError as e:
        raise ValueError(
            f"fill_value shape {fv_aval.shape} is not broadcast-compatible with "
            f"batch shape {batch_shape}"
        ) from e
      if broadcasted != batch_shape:
        raise ValueError(
            f"fill_value shape {fv_aval.shape} cannot be broadcast to "
            f"batch shape {batch_shape} without expanding it."
        )
    return (*batch_shape, size)

  @staticmethod
  def _sharding_rule(
      a_aval: core.ShapedArray,
      *fill_value_avals: core.ShapedArray,
      axes: tuple[int, ...],
  ):
    if any(a_aval.sharding.spec[i] is not None for i in axes):
      raise ValueError(
          f"Cannot apply nonzero over sharded axes {axes} for array with"
          f" sharding {a_aval.sharding}"
      )
    batch_spec = P(
        *(s for i, s in enumerate(a_aval.sharding.spec) if i not in axes)
    )
    for fv_aval in fill_value_avals:
      nlead = len(batch_spec) - fv_aval.ndim
      if any(
          fv_s is not None and fv_s != b_s
          for fv_s, b_s in zip(fv_aval.sharding.spec, batch_spec[nlead:])
      ):
        raise ValueError(
            f"fill_value sharding {fv_aval.sharding.spec} is not compatible"
            f" with batch sharding {batch_spec}"
        )
    return a_aval.sharding.update(spec=P(*batch_spec, None))

  def expand(self, a: ArrayLike, *fill_value: ArrayLike) -> tuple[Array, ...]:  # pyrefly: ignore[bad-override]
    return _nonzero_impl(a, *fill_value, size=self.size, axes=self.axes, out_dtype=self.out_dtype)

  def batch(
      self,
      axis_data: Any,
      args: tuple[Array, ...],
      dims: tuple[int | None, ...]
  ) -> tuple[tuple[Array, ...], int | None | tuple[int | None, ...]]:
    del axis_data  # unused
    a, *fvs = args
    d_a, *d_fvs = dims

    if d_a is None and all(d is None for d in d_fvs):
      return self(*args), (None,) * len(self.axes)

    # If a is not batched but some fv is, we broadcast a to have the batch dimension.
    if d_a is None:
      B = None
      for fv, d_fv in zip(fvs, d_fvs):
        if d_fv is not None:
          B = fv.shape[d_fv]
          break
      assert B is not None
      # Broadcast a to (B, *a.shape)
      a = lax.broadcast_in_dim(a, (B, *a.shape), tuple(range(1, a.ndim + 1)))
      d_a = 0

    # Move batch dim of a to 0
    if d_a != 0:
      a = jnp.moveaxis(a, d_a, 0)

    # Since batch dim of a is at 0, all original axes are shifted by 1.
    new_axes = tuple(ax + 1 for ax in self.axes)

    # Reshape fvs
    reshaped_fvs = []
    for fv, d_fv in zip(fvs, d_fvs):
      if d_fv is not None:
        if d_fv != 0:
          fv = jnp.moveaxis(fv, d_fv, 0)
        B = fv.shape[0]
        # Reshape to (B, 1, ..., 1)
        # We need a.ndim - len(new_axes) - 1 ones.
        num_ones = a.ndim - len(new_axes) - 1
        shape = (B,) + (1,) * num_ones
        reshaped_fvs.append(lax.reshape(fv, shape))
      else:
        reshaped_fvs.append(fv)

    batched_prim = Nonzero(
        core.typeof(a),
        *(core.typeof(fv) for fv in reshaped_fvs),
        size=self.size,
        axes=new_axes,
        out_dtype=self.out_dtype,
    )

    out_dims = (0,) * len(new_axes)
    return batched_prim(a, *reshaped_fvs), out_dims

  def jvp(self, primals: tuple[Array], tangents: Any) -> tuple[tuple[Array, ...], tuple[Array | np.ndarray, ...]]:
    del tangents  # unused
    primal_out = self(*primals)
    tangents_out = tuple(np.empty(p.shape, dtype=dtypes.float0) for p in primal_out)
    return primal_out, tangents_out

  def vjp_fwd(self, nzs_in: Any, *args: Array) -> tuple[tuple[Array, ...], None]:
    return (self(*args), None)

  def vjp_bwd_retval(self, res: Any, g: Any):
    return tuple(ad_util.zeros_like_aval(aval) for aval in self.in_avals)

  lin, linearized = linearize_from_jvp


class Einsum(HiPrim):
  """HiJAX primitive for einstein summation.

  Parameters:
    *operands_avals: Abstract values of the operands.
    subscripts: The subscripts of the einsum. These will be normalized
      via parse_einsum_input before being stored.
    optimize: The optimization strategy of the einsum; defaults to "auto".
    precision: Specify the precision of the computation.
    preferred_element_type: Specify the accumulator and result dtype.
  """
  multiple_results = False

  subscripts: str
  optimize: str | bool | tuple[tuple[int, ...], ...]
  precision: lax.PrecisionLike
  preferred_element_type: DTypeLike | None

  def __init__(
      self,
      *operands_avals: core.ShapedArray,
      subscripts: str,
      optimize: str | bool | tuple[tuple[int, ...], ...] = "auto",
      precision: lax.PrecisionLike = None,
      preferred_element_type: DTypeLike | None = None,
  ):
    # Normalize subscripts
    in_subs, out_sub, _ = parse_einsum_input((subscripts, *operands_avals))
    self.subscripts = f"{in_subs}->{out_sub}"
    self.optimize = optimize
    self.precision = precision
    self.preferred_element_type = preferred_element_type
    self.in_avals = operands_avals

    inputs = in_subs.split(',')
    shapes = [a.shape for a in operands_avals]
    out_shape = find_output_shape(inputs, shapes, out_sub)
    if preferred_element_type is None:
      out_dtype = dtypes.result_type(*(a.dtype for a in operands_avals))
    else:
      out_dtype = dtypes.check_and_canonicalize_user_dtype(
          preferred_element_type, 'einsum'
      )
    self.out_aval = core.ShapedArray(out_shape, out_dtype)

    self.params = dict(
        subscripts=self.subscripts,
        optimize=optimize,
        precision=precision,
        preferred_element_type=preferred_element_type,
    )
    super().__init__()

  def expand(self, *operands: ArrayLike) -> Array:
    path_type = (
        'optimal' if self.optimize is True else
        Unoptimized() if self.optimize is False else
        self.optimize
    )
    contract_path_args = (self.subscripts, *operands)
    operands_out, contractions = opt_einsum.contract_path(  # pyrefly: ignore[no-matching-overload]
        *contract_path_args, einsum_call=True, use_blas=True, optimize=path_type)
    contractions = tuple((a, frozenset(b), c) for a, b, c, *_ in contractions)
    return einsum_impl(
        list(operands_out), contractions, precision=self.precision,
        preferred_element_type=self.preferred_element_type)

  def batch(
      self,
      axis_data: Any,
      args: tuple[Array, ...],
      dims: tuple[int | None, ...]
  ) -> tuple[Array, int | None]:
    del axis_data  # unused
    if all(d is None for d in dims):
      return self(*args), None
    input_subs, output_sub = self.subscripts.split('->')
    input_subs_list = input_subs.split(',')

    used_chars = set(input_subs) | set(output_sub) | {',', '-', '>'}
    available_chars = [c for c in string.ascii_letters if c not in used_chars]
    if not available_chars:
      raise ValueError("Run out of characters for batch dimension in einsum.")
    batch_char = available_chars[0]

    new_input_subs_list = []
    for i, dim in enumerate(dims):
      sub = input_subs_list[i]
      if dim is not None:
        sub_list = list(sub)
        sub_list.insert(dim, batch_char)
        new_sub = "".join(sub_list)
      else:
        new_sub = sub
      new_input_subs_list.append(new_sub)

    new_input_subs = ",".join(new_input_subs_list)
    new_output_sub = batch_char + output_sub
    new_subscripts = f"{new_input_subs}->{new_output_sub}"

    batched_prim = Einsum(
        *(core.typeof(arg) for arg in args),
        subscripts=new_subscripts,
        optimize=self.optimize,
        precision=self.precision,
        preferred_element_type=self.preferred_element_type,
    )
    return batched_prim(*args), 0

  def jvp(self, primals: tuple[Array, ...], tangents: Any) -> tuple[Array, Array]:
    primal_out = self(*primals)
    tangent_outs = []
    for i, t in enumerate(tangents):
      if not isinstance(t, ad_util.Zero):
        tangent_outs.append(self(*primals[:i], t, *primals[i+1:]))
    if not tangent_outs:
      return primal_out, ad_util.zeros_like_aval(self.out_aval)
    return primal_out, functools.reduce(lax.add, tangent_outs)

  def transpose(self, out_ct, *maybe_accums):
    in_subs_list = self.subscripts.split('->')[0].split(',')
    out_sub = self.subscripts.split('->')[1]

    accums = [acc for acc in maybe_accums if isinstance(acc, ad.GradAccum)]
    if len(accums) > 1:
      raise NotImplementedError("Transpose of Einsum with multiple linear inputs is not supported.")

    for i, accum in enumerate(maybe_accums):
      if isinstance(accum, ad.GradAccum):
        if isinstance(out_ct, ad_util.Zero):
          accum.accum(ad_util.zeros_like_aval(self.in_avals[i]))
          continue

        orig_sub = in_subs_list[i]
        orig_aval = self.in_avals[i]

        all_ct_input_chars = set(out_sub)
        for k, sub in enumerate(in_subs_list):
          if k != i:
            all_ct_input_chars.update(sub)

        missing_chars = sorted(
            set(orig_sub) - all_ct_input_chars, key=orig_sub.index
        )

        if missing_chars:
          missing_sizes = tuple(
              orig_aval.shape[orig_sub.index(c)] for c in missing_chars
          )
          expanded_out_ct = out_ct[(...,) + (None,) * len(missing_chars)]
          curr_out_ct = jnp.broadcast_to(
              expanded_out_ct, out_ct.shape + missing_sizes
          )
          curr_out_sub = out_sub + "".join(missing_chars)
        else:
          curr_out_ct = out_ct
          curr_out_sub = out_sub

        in_subs = ",".join([*in_subs_list[:i], curr_out_sub, *in_subs_list[i + 1:]])
        ct_subscripts = f"{in_subs}->{orig_sub}"
        ct_operands = (*maybe_accums[:i], curr_out_ct, *maybe_accums[i + 1:])
        ct_val = einsum(
            ct_subscripts,
            *ct_operands,
            optimize=self.optimize,
            precision=self.precision,
        )
        accum.accum(ct_val)

  lin, linearized = linearize_from_jvp
  vjp_fwd, vjp_bwd_retval = vjp_from_jvp


def _nonzero_impl(a: ArrayLike, *fill_value: ArrayLike, size: int, axes: tuple[int, ...], out_dtype: np.dtype) -> tuple[Array, ...]:
  """Main implementation of nonzero primitive."""
  a = jnp.asarray(a)
  out_dtype = dtypes._maybe_canonicalize_explicit_dtype(out_dtype, "nonzero")
  axes = tuple(sorted(axes))

  if not axes:
    return ()

  batch_axes = [ax for ax in range(a.ndim) if ax not in axes]
  batch_shape = tuple(a.shape[ax] for ax in batch_axes)
  a_aval = core.typeof(a)
  if a.size == 0 or size == 0:
    out_spec = P(*(a_aval.sharding.spec[ax] for ax in batch_axes), None)
    out_sharding = a_aval.sharding.update(spec=out_spec)
    vma = tuple(
        core.standard_vma_rule(
            "nonzero", a_aval, *(core.typeof(fv) for fv in fill_value)
        )
    )
    return tuple(
        core.pvary(
            jnp.empty((*batch_shape, size), dtype=out_dtype, out_sharding=out_sharding),
            vma,
        )
        for _ in axes
    )

  transposed_a = jnp.transpose(a, (*batch_axes, *axes))
  sub_shape = transposed_a.shape[len(batch_axes):]
  strides = np.cumprod(sub_shape[::-1])[::-1] // sub_shape
  strides = tuple(strides.tolist())

  flattened_a = transposed_a.reshape(*batch_shape, -1)
  mask = flattened_a if flattened_a.dtype == bool else (flattened_a != 0)
  cs_mask = jnp.cumsum(mask, axis=-1)

  def _bincount(c):
    c_aval = core.typeof(c)
    z = jnp.zeros(
        size, dtype=c.dtype, out_sharding=c_aval.sharding.update(spec=P())
    )
    z = core.pvary(z, tuple(c_aval.mat.varying))
    return z.at[c].add(1, mode='drop')

  bincount_fn = _bincount
  for _ in range(len(batch_shape)):
    bincount_fn = api.vmap(bincount_fn)
  bincount = bincount_fn(cs_mask)
  flat_indices = jnp.cumsum(bincount, axis=-1)

  out = [lax.convert_element_type((flat_indices // stride) % sz, out_dtype)
         for stride, sz in zip(strides, sub_shape)]
  counts = mask.sum(axis=-1, keepdims=True)
  fill_mask = lax.expand_dims(jnp.arange(size), range(counts.ndim - 1)) >= counts
  if fill_value:
    return tuple(jnp.where(fill_mask, lax.expand_dims(fv, [np.ndim(fv)]), entry)
                 for fv, entry in zip(fill_value, out))
  else:
    return tuple(jnp.where(fill_mask, 0, entry) for entry in out)


def _searchsorted_impl(sorted_arr: ArrayLike, query: ArrayLike, *, dimension: int,
                       batch_dims: int, side: str, dtype: np.dtype, method: str):
  """Main implementation of searchsorted primitive."""
  sorted_arr = jnp.moveaxis(sorted_arr, dimension, -1)
  dtype = dtypes._maybe_canonicalize_explicit_dtype(dtype, "searchsorted")

  if method == "scan":
    impl: Callable[..., Array] = functools.partial(_searchsorted_scan_impl, unrolled=False)
  elif method == "scan_unrolled":
    impl = functools.partial(_searchsorted_scan_impl, unrolled=True)
  elif method == "compare_all":
    impl = _searchsorted_compare_all_impl
  elif method == "sort":
    impl = _searchsorted_sort_impl
  else:
    raise ValueError(f"Unsupported method: {method}")

  fun = functools.partial(impl, side=side, dtype=dtype)
  for _ in range(sorted_arr.ndim - batch_dims - 1):
    fun = api.vmap(fun, in_axes=(0, None))
  for _ in range(batch_dims):
    fun = api.vmap(fun, in_axes=0)
  return fun(sorted_arr, query)


@api.jit(static_argnames=["side", "dtype", "unrolled"])
def _searchsorted_scan_impl(
    sorted_arr: Array, query: Array, side: str, dtype: np.dtype, unrolled: bool
) -> Array:
  """Scan-based implementation of searchsorted."""
  assert sorted_arr.ndim == 1
  assert side in ["left", "right"]
  (n,) = sorted_arr.shape
  if sorted_arr.size == 0:
    return lax.full(query.shape, fill_value=0, dtype=dtype)
  if query.ndim > 0:
    return api.vmap(functools.partial(_searchsorted_scan_impl, side=side,
                                      dtype=dtype, unrolled=unrolled),
                    in_axes=(None, 0))(sorted_arr, query)

  op = lax._sort_le_comparator if side == "left" else lax._sort_lt_comparator
  unsigned_dtype = np.uint64 if dtypes.iinfo(dtype).bits == 64 else np.uint32
  def body_fun(state, _):
    low, high = state
    mid = low + (high - low) // 2  # use this form to avoid overflow
    go_left = op(query, sorted_arr[mid])
    return (lax.select(go_left, low, mid), lax.select(go_left, mid, high)), ()
  n_levels = int(np.ceil(np.log2(n + 1)))
  sa_aval = core.typeof(sorted_arr)
  vma = tuple(sa_aval.mat.varying)
  sa_scalar_sh = sa_aval.sharding.update(spec=P())
  init = (jnp.array(0, dtype=unsigned_dtype, out_sharding=sa_scalar_sh),
          jnp.array(n, dtype=unsigned_dtype, out_sharding=sa_scalar_sh))
  init = tuple(core.pvary(i, vma) for i in init)
  carry, _ = control_flow.scan(body_fun, init, (), length=n_levels,
                               unroll=n_levels if unrolled else 1)
  return carry[1].astype(dtype)


@api.jit(static_argnames=["side", "dtype"])
def _searchsorted_sort_impl(
    sorted_arr: Array, query: Array, side: str, dtype: np.dtype
) -> Array:
  """Cosorting-based implementation of searchsorted."""
  assert sorted_arr.ndim == 1
  assert side in ["left", "right"]
  working_dtype = (
      np.int32
      if sorted_arr.size + query.size <= np.iinfo(np.int32).max
      else np.int64
  )

  def _rank(x):
    idx = lax.iota(working_dtype, len(x))
    return lax.full_like(idx, 0).at[lax.sort_key_val(x, idx)[1]].set(idx)

  if side == "left":
    index = _rank(lax.concatenate([query.ravel(), sorted_arr], 0))[: query.size]
  else:
    index = _rank(lax.concatenate([sorted_arr, query.ravel()], 0))[
        sorted_arr.size :
    ]
  return lax.reshape(
      lax.sub(index, _rank(query.ravel())), np.shape(query)
  ).astype(dtype)


@api.jit(static_argnames=["side", "dtype"])
def _searchsorted_compare_all_impl(
    sorted_arr: Array, query: Array, side: str, dtype: np.dtype
) -> Array:
  assert sorted_arr.ndim == 1
  assert side in ["left", "right"]
  op = lax._sort_lt_comparator if side == "left" else lax._sort_le_comparator
  comparisons = api.vmap(op, in_axes=(0, None))(sorted_arr, query)
  return comparisons.sum(dtype=dtype, axis=0)


def searchsorted(
    sorted_arr: ArrayLike,
    query: ArrayLike,
    /,
    *,
    side: str = "left",
    dimension: int = 0,
    batch_dims: int = 0,
    method: str = "scan",
    dtype: DTypeLike = "int32",
):
  """Find indices of query values within a sorted array.

  This is a batch-aware implementation of :func:`numpy.searchsorted` built on a
  HiJAX primitive. It adds the `batch_dims` and `dimension` argument, which make
  the API closed under batching.

  Args:
    sorted_arr: N-dimensional array, which is assumed to be sorted in increasing
      order along ``dimension``.
    query: N-dimensional array of query values.
    side: 'left' (default) or 'right'. If 'left', find the index of the first
      suitable location. If 'right', find the index of the last.
    dimension: positive integer specifying the dimension of ``sorted_arr`` along
      which to insert query values. Defaults to the first dimension.
    batch_dims: integer specifying the number of leading dimensions of
      ``sorted_arr`` and ``query`` to treat as shared batch dimensions.
      Defaults to zero.
    method: string specifying the search method: one of 'scan' (default),
      'compare_all', or 'sort'. 'scan' uses a scan-based binary search implementation,
      'compare_all' directly compares all elements in `sorted_arr` to `query`, and
      'sort' uses a cosorting-based implementation.

  Returns:
    An array specifying the insertion locations of `query` into `sorted_arr`.
  """
  sorted_arr, query = core.auto_insert_reshard(sorted_arr, query)
  out_dtype = dtypes._maybe_canonicalize_explicit_dtype(np.dtype(dtype), "searchsorted")
  prim = SearchSorted(
    core.typeof(sorted_arr),
    core.typeof(query),
    side=side,
    dimension=dimension,
    batch_dims=batch_dims,
    method=method,
    out_dtype=out_dtype,
  )
  return prim(sorted_arr, query)


def nonzero(
    a: ArrayLike,
    /,
    *,
    size: int,
    fill_value: ArrayLike | tuple[ArrayLike, ...] | None = None,
    axes: tuple[int, ...] | None = None,
    dtype: DTypeLike = 'int32',
) -> tuple[Array, ...]:
  """Return indices of nonzero elements.

  This is a batch-aware implementation of :func:`numpy.nonzero` built on a
  HiJAX primitive.

  Args:
    a: N-dimensional array.
    size: static integer specifying the number of nonzero entries to return.
    fill_value: optional padding value when ``size`` is specified. Defaults to 0.
    axes: optional tuple of integers specifying the axes to compute the result over.
      Defaults to None (all axes).
    dtype: optional datatype for the returned indices. Defaults to int32.

  Returns:
    Tuple of length ``len(axes)`` containing the indices of each nonzero value.
  """
  out_dtype = dtypes._maybe_canonicalize_explicit_dtype(np.dtype(dtype), "nonzero")
  axes = util.canonicalize_axis_tuple(axes, np.ndim(a))

  if fill_value is not None:
    if isinstance(fill_value, tuple):
      if len(fill_value) != len(axes):
        raise ValueError(f"fill_value tuple must have length equal to number of axes ({len(axes)}); got {len(fill_value)}")
      fill_value_tup = fill_value
    else:
      fill_value_tup = (fill_value,) * len(axes)
    fill_value_tup = tuple(jnp.asarray(fv, dtype=out_dtype) for fv in fill_value_tup)
    for fv in fill_value_tup:
      if fv.ndim != 0:
        raise ValueError(f"fill_value must be a scalar or tuple of scalars; got {fill_value}")
  else:
    fill_value_tup = ()

  a, *fill_value_tup = core.auto_insert_reshard(a, *fill_value_tup)
  prim = Nonzero(
    core.typeof(a),
    *[core.typeof(fv) for fv in fill_value_tup],
    size=size,
    axes=axes,
    out_dtype=out_dtype,
  )
  return prim(a, *fill_value_tup)


def einsum(
    subscripts: str,
    /,
    *operands: ArrayLike,
    optimize: str | bool | tuple[tuple[int, ...], ...] = "auto",
    precision: lax.PrecisionLike = None,
    preferred_element_type: DTypeLike | None = None,
) -> Array:
  """Einstein summation built on a HiJAX primitive.

  Args:
    subscripts: string containing axes names separated by commas.
    *operands: sequence of one or more arrays corresponding to the subscripts.
    optimize: specify how to optimize the order of computation.
    precision: Specify the precision of the computation.
    preferred_element_type: Specify the accumulator and result dtype.

  Returns:
    array containing the result of the einstein summation.
  """
  operands = core.auto_insert_reshard(*operands)
  prim = Einsum(
      *(core.typeof(op) for op in operands),
      subscripts=subscripts,
      optimize=tuple(optimize) if isinstance(optimize, list) else optimize,
      precision=precision,
      preferred_element_type=preferred_element_type,
  )
  return prim(*operands)