# Copyright 2025 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.

import abc
from collections.abc import Sequence
import dataclasses
import math
import threading
from typing import Any, Literal

from jax import lax
from jax._src import core as jax_core
from jax._src import source_info_util
from jax._src.pallas import primitives
from jax._src.state import indexing
from jax._src.state import types as state_types
from jax._src.util import safe_map
import jax.numpy as jnp
import numpy as np


def get_uninitialized_value(
    dtype, uninitialized_memory: Literal["nan", "zero"]
):
  if uninitialized_memory == "nan":
    if jnp.issubdtype(dtype, jnp.floating):
      return np.nan
    elif jnp.issubdtype(dtype, jnp.integer):
      return jnp.iinfo(dtype).max
    elif jnp.issubdtype(dtype, jnp.bool):
      return True
  if uninitialized_memory == "zero":
    return 0
  raise NotImplementedError(uninitialized_memory + " + " + str(dtype))


def get_uninitialized_array(
    shape, dtype, uninitialized_memory: Literal["nan", "zero"]
):
  return jnp.full(
      shape,
      get_uninitialized_value(dtype, uninitialized_memory),
      dtype,
  )


def pad_to_block_dimension(
    value, block_shape, uninitialized_memory: Literal["nan", "zero"]
):
  """Pads values so the shape evenly divides into block dimensions.

  For example, if values has a shape of (33, 2, 5) with a block_shape of
  (32, 2, 4), this function will pad the value of shape to (64, 2, 8).

  Args:
    value: Array to be padded.
    block_shape: Block shapes to use for padding. If None, no padding will be
      performed.
    uninitialized_memory: The value to use for padding.

  Returns:
    A padded array.
  """
  padded_shape = tuple(
      ((v - 1) // b + 1) * b for v, b in zip(value.shape, block_shape)
  )
  if padded_shape != value.shape:
    pad_width = tuple((0, a - b) for a, b in zip(padded_shape, value.shape))
    pad_value = get_uninitialized_array((), value.dtype, uninitialized_memory)
    value = jnp.pad(value, pad_width, constant_values=pad_value)
  return value


@dataclasses.dataclass(frozen=True, kw_only=True)
class LoggingInfo(abc.ABC):
  """Holds information for logging."""

  source_info: source_info_util.SourceInfo | None = None
  device_id: int

  @abc.abstractmethod
  def get_location_str(self) -> str:
    """Returns a string representation of the location (device/core/thread)."""
    raise NotImplementedError()

  def get_source_info_str(self) -> str:
    if self.source_info is None:
      return "None"
    return source_info_util.summarize(self.source_info)

  def format(self, message: str, line_prefix: str | None = None) -> str:
    """Formats a message for logging, across multiple lines.

    Each lines is prefixed with location information (device/core/thread) and
    with
    the given `line_prefix`.

    Args:
      message: The message to format.
      line_prefix: The prefix to add to each line, or None.

    Returns:
      The formatted message. Each line is prefixed with location information
      (device/core/thread) and with the given `line_prefix`. A final lines with
      a summary of `self.source_info`, if not `None`, is appended.
    """

    if line_prefix is None:
      line_prefix = f"{self.get_location_str()}: "
    else:
      line_prefix = f"{self.get_location_str()}: {line_prefix.strip()} --"

    lines = [f"{line_prefix} {l.strip()}" for l in message.split("\n")]
    if self.source_info:
      lines.append(f"{line_prefix} {self.get_source_info_str()}.")

    return "\n".join(lines)


@dataclasses.dataclass(frozen=True, kw_only=True)
class TPULoggingInfo(LoggingInfo):
  """Logging info for TPU interpret mode."""

  local_core_id: int

  def get_location_str(self) -> str:
    return f"Device {self.device_id}, core {self.local_core_id}"


class Counter:
  """A simple counter that is thread-safe."""

  def __init__(self, initial_value: int):
    self.value = initial_value
    self.lock = threading.Lock()

  def get_next(self):
    with self.lock:
      result = self.value
      self.value += 1
    return result


# TODO(sharadmv): De-dup this w/ the impl in primitives.py.
def _device_id_dict_to_mesh(device_id_dict, axis_sizes, axis_indices):
  physical_axis_dict = {}
  axis_names = axis_sizes.keys()
  for axis, idx in device_id_dict.items():
    if isinstance(axis, tuple) and len(axis) == 1:
      axis = axis[0]
    if isinstance(axis, tuple) and any(a in axis_names for a in axis):
      if not all(a in axis_names for a in axis):
        raise NotImplementedError(
            f"{axis} mixes JAX mesh and Pallas mesh grid axes"
        )
      axes_dimensions = [axis_sizes[name] for name in axis]
      for axis_index, axis_name in enumerate(axis):
        axis_size = axis_sizes[axis_name]
        inner_mesh_size = math.prod(axes_dimensions[axis_index + 1 :])
        minor_divisor = inner_mesh_size

        # Fast path for power of 2s
        if inner_mesh_size & (inner_mesh_size - 1) == 0:
          shift_len = (inner_mesh_size & -inner_mesh_size).bit_length() - 1
          partial_device_idx = idx >> shift_len
        else:
          partial_device_idx = idx // minor_divisor

        if axis_size & (axis_size - 1) == 0:
          device_idx = partial_device_idx & (axis_size - 1)
        else:
          device_idx = partial_device_idx % axis_size
        physical_axis_dict[axis_name] = device_idx
    else:
      physical_axis_dict[axis] = idx
  device_id = []
  for axis in axis_names:
    if axis in physical_axis_dict:
      device_id.append(physical_axis_dict[axis])
    else:
      device_id.append(axis_indices[axis])
  non_mesh_axes = {
      k: v for k, v in physical_axis_dict.items() if k not in axis_names
  }
  return tuple(device_id), non_mesh_axes


def device_coords_to_logical_id(
    device_coords, axis_sizes, axis_indices, *, allow_non_mesh_axes: bool = False
):
  non_mesh_axes = {}
  if isinstance(device_coords, dict):
    device_coords, non_mesh_axes = _device_id_dict_to_mesh(
        device_coords, axis_sizes, axis_indices
    )
    if non_mesh_axes and not allow_non_mesh_axes:
      raise NotImplementedError(non_mesh_axes)
  if not isinstance(device_coords, tuple):
    device_coords = (device_coords,)
  assert len(device_coords) == len(axis_sizes)
  sizes = list(axis_sizes.values())
  ret = 0
  for i in range(len(device_coords)):
    ret += device_coords[i] * math.prod(sizes[i + 1 :])
  if allow_non_mesh_axes:
    return ret, non_mesh_axes
  return ret


def _device_id_to_logical(
    device_id,
    device_id_type,
    axis_sizes,
    axis_indices,
    *,
    core_axis_names: Sequence[jax_core.AxisName] = (),
):
  if device_id is None:
    return None, None
  if device_id_type == primitives.DeviceIdType.MESH:
    logical_id, non_mesh_axes = device_coords_to_logical_id(
        device_id, axis_sizes, axis_indices, allow_non_mesh_axes=True
    )
    core_index = None
    if non_mesh_axes:
      if (
          set(non_mesh_axes.keys()) <= set(core_axis_names)
          and len(non_mesh_axes) == 1
      ):
        (core_index,) = non_mesh_axes.values()
      else:
        raise ValueError(f"Unrecognized axes in device_id: {non_mesh_axes}")
    return logical_id, core_index
  elif device_id_type == primitives.DeviceIdType.LOGICAL:
    return device_id, None
  else:
    raise ValueError(f"Unsupported device ID type: {device_id_type}")


def is_int(dtype):
  return jnp.issubdtype(dtype, jnp.integer)


def is_float(dtype):
  return jnp.issubdtype(dtype, jnp.floating)


@dataclasses.dataclass(frozen=True)
class Placeholder:
  """Placeholder for use in `JaxprEnv` below instead of storing a concrete value."""

  shape: tuple[int, ...]
  dtype: jnp.dtype


class JaxprEnv:
  """An environment for interpreting jaxprs, mapping variables to values."""

  def __init__(
      self,
      *,
      vars: Sequence[jax_core.Var] | None = None,
      values: Sequence[Any] | None = None,
      sentinel_for_floating_point_values: Any = None,
  ):
    self._sentinel_for_floating_point_values = (
        sentinel_for_floating_point_values
    )
    self._env: dict[jax_core.Var, Any] = {}

    if vars is None and values is None:
      return

    vars = vars or []
    values = values or []
    self.write_many(vars, values)

  def read(self, var):
    if isinstance(var, jax_core.Literal):
      result = var.val
    else:
      result = self._env[var]
    if isinstance(result, Placeholder):
      result = lax.full(
          result.shape, self._sentinel_for_floating_point_values, result.dtype
      )
    return result

  def read_many(self, vars):
    return safe_map(self.read, vars)

  def write(self, var, value):
    if self._sentinel_for_floating_point_values and is_float(value.dtype):
      value = Placeholder(value.shape, value.dtype)
    self._env[var] = value

  def write_many(self, vars, values):
    safe_map(self.write, vars, values)


def _transform_slice_or_index(slice_or_idx):
  if isinstance(slice_or_idx, int):
    return slice_or_idx
  else:
    start = int(slice_or_idx.start)
    size = int(slice_or_idx.size)
    stride = int(slice_or_idx.stride)
    return slice(start, start + size * stride, stride)


def _compose_slice_or_index(slice_or_idx1, slice_or_idx2):
  ret = []
  i = 0
  j = 0
  while True:
    if i == len(slice_or_idx1):
      ret.extend(slice_or_idx2[j:])
      return tuple(ret)
    elif j == len(slice_or_idx2):
      ret.extend(slice_or_idx1[i:])
      return tuple(ret)
    elif isinstance(slice_or_idx1[i], int):
      ret.append(slice_or_idx1[i])
      i += 1
    elif isinstance(slice_or_idx2[j], int):
      ret.append(
          slice_or_idx1[i].start + slice_or_idx2[j] * slice_or_idx1[i].step
      )
      i += 1
      j += 1
    else:
      ret.append(
          slice(
              slice_or_idx1[i].start
              + slice_or_idx2[j].start * slice_or_idx1[i].step,
              slice_or_idx1[i].start
              + slice_or_idx2[j].stop * slice_or_idx1[i].step,
              slice_or_idx1[i].step * slice_or_idx2[j].step,
          )
      )
      i += 1
      j += 1


@dataclasses.dataclass(frozen=True)
class Access:
  """What a ref access touches.

  Use as `buffer[access.range].transpose(access.permutation)`. The range has one
  entry per buffer axis and is relative to the buffer's logical shape.
  """

  range: tuple[slice | int, ...]
  permutation: tuple[int, ...] | None = None

  @property
  def inverse_permutation(self) -> tuple[int, ...] | None:
    if self.permutation is None:
      return None
    inverse = [0] * len(self.permutation)
    for position, axis in enumerate(self.permutation):
      inverse[axis] = position
    return tuple(inverse)

  @property
  def shape(self) -> tuple[int, ...]:
    """The shape of the value read or written assuming no OOB access."""
    requested = []
    for r in self.range:
      if isinstance(r, int):
        continue  # An integer index drops its axis.
      step = 1 if r.step is None else r.step
      requested.append(max(0, (r.stop - r.start + step - 1) // step))
    if self.permutation is not None:
      requested = [requested[axis] for axis in self.permutation]
    return tuple(requested)


def _compose_axis(outer: slice, inner):
  """Composes `inner`, stated in `outer`'s coordinates, with `outer`."""
  if not isinstance(inner, indexing.Slice):
    inner = int(inner)
  return _compose_slice_or_index(
      (outer,), (_transform_slice_or_index(inner),)
  )[0]


def to_access(transforms, shape: Sequence[int]) -> Access:
  """Normalizes a ref's transforms into one buffer range plus one permutation.

  Indexing a transposed view is the same as indexing the buffer with the
  indices permuted and transposing the result, so however indexers and
  transposes are interleaved they collapse into this pair.
  """
  # For each buffer axis, what is selected from it.
  selection: list[slice | int] = [slice(0, dim, 1) for dim in shape]
  # The buffer axes still visible.
  visible = list(range(len(shape)))

  for transform in transforms:
    if isinstance(transform, state_types.TransposeTransform):
      permutation = tuple(int(p) for p in transform.permutation)
      visible = [visible[p] for p in permutation]
      continue
    dropped = []
    for position, index in enumerate(transform.indices):
      axis = visible[position]
      outer = selection[axis]
      assert isinstance(outer, slice)
      selection[axis] = _compose_axis(outer, index)
      if not isinstance(index, indexing.Slice):
        # `int` case, since advanced indexing will be rejected in _compose_axis
        dropped.append(axis)
    visible = [axis for axis in visible if axis not in dropped]

  # `buffer[range]` hands back the surviving axes in *buffer* order; the access
  # wants them in `visible` order.
  surviving = sorted(visible)
  permutation = tuple(surviving.index(axis) for axis in visible)
  if list(permutation) == sorted(permutation):
    permutation = None  # Already in buffer order.
  return Access(tuple(selection), permutation)


def to_range(transforms) -> tuple[slice | int, ...]:
  ret = ()
  for transform in transforms:
    # For now, assume only NDIndexer transforms.
    ret = _compose_slice_or_index(
        ret, tuple(_transform_slice_or_index(i) for i in transform.indices)
    )
  return ret


def is_range_out_of_bounds_for_shape(
    rnge: tuple[slice | int, ...], shape: tuple[int, ...]
) -> bool:
  """Returns whether `rnge` is at least partially out of bounds for `shape`."""
  for d, r in zip(shape, rnge, strict=True):
    if isinstance(r, int):
      assert 0 <= r
      if r >= d:
        return True
    elif isinstance(r, slice):
      assert r.start is not None and 0 <= r.start
      assert r.stop is not None and 0 <= r.stop

      if r.step is None:
        if r.stop > d:
          return True
      else:
        assert 0 <= r.step
        num_elements_in_slice = (r.stop - r.start + r.step - 1) // r.step
        if num_elements_in_slice > 0:
          last_index = r.start + (num_elements_in_slice - 1) * r.step
          if last_index >= d:
            return True
    else:
      raise ValueError(f"Unsupported range type: {type(r)}.")
  return False


def clip_range_to_shape(
    rnge: tuple[slice | int, ...], shape: tuple[int, ...]
) -> tuple[slice | int, ...] | None:
  """Clips `slice`s in `rnge` to the `shape`. Returns None if `rnge` is entirely out of bounds."""
  result: list[slice | int] = []
  for r, l in zip(rnge, shape, strict=True):
    if isinstance(r, int):
      if r >= l:
        return None
      result.append(r)
    elif isinstance(r, slice):
      if r.start >= l:
        return None
      result.append(slice(r.start, min(r.stop, l), r.step))
    else:
      raise ValueError(f"Unsupported range type: {type(r)}.")
  return tuple(result)


def get_next_indices(grid, indices):
  next_indices = []
  carry = True
  for dim_size, index in reversed(list(zip(grid, indices))):
    i = jnp.where(carry, index + 1, index)
    carry = dim_size == i
    next_indices.append(jnp.where(carry, 0, i))
  return tuple(reversed(next_indices))


def get_indices(grid, loop_index):
  indices = []
  for dim_size in reversed(grid):
    i = loop_index % dim_size
    loop_index = loop_index // dim_size
    indices.append(i)
  return tuple(reversed(indices))