# 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.
"""Contains SparseCore-specific Pallas abstractions."""

from __future__ import annotations

import collections
from collections.abc import Sequence
import contextlib
import dataclasses
from typing import Any

import jax
from jax._src import core as jax_core
from jax._src import pretty_printer as pp
from jax._src import tree_util
from jax._src.pallas import core as pallas_core
from jax._src.pallas.mosaic import core as tpu_core
from jax._src.pallas.mosaic import tpu_info
from jax._src.state import types as state_types
import jax.numpy as jnp


class SparseCorePushStreamWarning(UserWarning):
  """Warns that destination semaphores are ignored for push streams."""


def get_sparse_core_info() -> tpu_info.SparseCoreInfo:
  """Returns the SparseCore information for the current device.

  Raises:
    RuntimeError: If the current TPU does not have SparseCores.
  """
  sc_info = tpu_info.get_tpu_info().sparse_core
  if sc_info is None:
    raise RuntimeError("The current TPU does not have SparseCores")
  return sc_info


@dataclasses.dataclass(frozen=True, kw_only=True)
class ScalarSubcoreMesh(pallas_core.Mesh):
  axis_name: str
  num_cores: int = dataclasses.field(
      default_factory=lambda: get_sparse_core_info().num_cores
  )

  def __post_init__(self):
    sc_info = get_sparse_core_info()
    if self.num_cores > sc_info.num_cores:
      raise ValueError(
          f"Mesh has {self.num_cores} cores, but the current TPU chip has only"
          f" {sc_info.num_cores} SparseCores"
      )

  @property
  def core_type(self) -> tpu_core.CoreType:
    return tpu_core.CoreType.SC_SCALAR_SUBCORE

  @property
  def default_memory_space(self) -> tpu_core.MemorySpace:
    return tpu_core.MemorySpace.HBM

  @property
  def shape(self):
    return collections.OrderedDict({self.axis_name: self.num_cores})

  @property
  def size(self) -> int:
    return self.num_cores

  @property
  def dimension_semantics(self) -> Sequence[tpu_core.DimensionSemantics]:
    return [tpu_core.GridDimensionSemantics.CORE_PARALLEL]

  def discharges_effect(self, effect):
    del effect  # Unused.
    return False

  def check_is_compatible_with(self, other_mesh):
    if isinstance(other_mesh, ScalarSubcoreMesh):
      raise ValueError("You can't use two different ScalarSubcoreMeshes.")
    elif isinstance(other_mesh, VectorSubcoreMesh):
      if (self.axis_name == other_mesh.core_axis_name
          and self.num_cores == other_mesh.num_cores):
        return True
      raise ValueError(f"{self} should have the same core axis name and number"
                       f" of cores as the VectorSubcoreMesh {other_mesh}.")
    elif isinstance(other_mesh, tpu_core.TensorCoreMesh):
      if self.axis_name == other_mesh.axis_name:
        raise ValueError(
            f"{self} should have a different axis name from the TensorCoreMesh"
            f" {other_mesh}."
        )
      return True
    return super().check_is_compatible_with(other_mesh)

  @property
  def supported_memory_spaces(self) -> Sequence[Any]:
    return [
        tpu_core.MemorySpace.VMEM_SHARED,
        tpu_core.MemorySpace.SMEM,
        tpu_core.MemorySpace.SEMAPHORE,
    ]

  @contextlib.contextmanager
  def tracing_context(self):
    yield


@dataclasses.dataclass(frozen=True, kw_only=True)
class VectorSubcoreMesh(pallas_core.Mesh):
  core_axis_name: str
  subcore_axis_name: str
  num_cores: int = dataclasses.field(
      default_factory=lambda: get_sparse_core_info().num_cores
  )
  num_subcores: int = dataclasses.field(
      default_factory=lambda: get_sparse_core_info().num_subcores
  )

  def __post_init__(self):
    sc_info = get_sparse_core_info()
    if self.num_cores > sc_info.num_cores:
      raise ValueError(
          f"Mesh has {self.num_cores} cores, but the current TPU chip has only"
          f" {sc_info.num_cores} SparseCores"
      )
    if self.num_subcores > sc_info.num_subcores:
      raise ValueError(
          f"Mesh has {self.num_subcores} subcores, but the current TPU chip has"
          f" only {sc_info.num_subcores} subcores"
      )

  @property
  def core_type(self) -> tpu_core.CoreType:
    return tpu_core.CoreType.SC_VECTOR_SUBCORE

  @property
  def default_memory_space(self) -> tpu_core.MemorySpace:
    return tpu_core.MemorySpace.HBM

  @property
  def shape(self):
    return collections.OrderedDict({
        self.core_axis_name: self.num_cores,
        self.subcore_axis_name: self.num_subcores,
    })

  @property
  def size(self) -> int:
    return self.num_cores * self.num_subcores

  @property
  def dimension_semantics(self) -> Sequence[tpu_core.DimensionSemantics]:
    return [
        tpu_core.GridDimensionSemantics.CORE_PARALLEL,
        tpu_core.GridDimensionSemantics.SUBCORE_PARALLEL,
    ]

  def discharges_effect(self, effect):
    del effect  # Unused.
    return False

  def check_is_compatible_with(self, other_mesh):
    if isinstance(other_mesh, VectorSubcoreMesh):
      raise ValueError("You can't use two different VectorSubcoreMeshes.")
    elif isinstance(other_mesh, ScalarSubcoreMesh):
      if (other_mesh.axis_name == self.core_axis_name
          and other_mesh.num_cores == self.num_cores):
        return True
      raise ValueError(f"{self} should have the same core axis name and number"
                       f" of cores as the ScalarSubcoreMesh {other_mesh}.")
    elif isinstance(other_mesh, tpu_core.TensorCoreMesh):
      if self.core_axis_name == other_mesh.axis_name:
        raise ValueError(
            f"{self} should have a different core axis name from the"
            f" TensorCoreMesh {other_mesh}."
        )
      if self.subcore_axis_name == other_mesh.axis_name:
        raise ValueError(
            f"{self} should have a different subcore axis name from the"
            f" TensorCoreMesh {other_mesh}."
        )
      return True
    return super().check_is_compatible_with(other_mesh)

  @property
  def supported_memory_spaces(self) -> Sequence[Any]:
    return [
        tpu_core.MemorySpace.VMEM,
        tpu_core.MemorySpace.VMEM_SHARED,
        tpu_core.MemorySpace.SMEM,
        tpu_core.MemorySpace.SEMAPHORE,
    ]

  @contextlib.contextmanager
  def tracing_context(self):
    yield


def supported_shapes(dtype: jax.typing.DTypeLike) -> Sequence[tuple[int, ...]]:
  """Returns all supported array shapes for the given dtype on SparseCore."""
  sc_info = get_sparse_core_info()
  num_lanes = sc_info.num_lanes
  itemsize = jnp.dtype(dtype).itemsize
  if itemsize > 4:
    raise ValueError(f"Unsupported dtype: {dtype}")
  packing_factor = 4 // itemsize
  if packing_factor == 1:
    return [(num_lanes,)]
  return [(num_lanes * packing_factor,), (packing_factor, num_lanes)]


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class Indices:
  """Indices for a gather or a scatter on SparseCore.

  Attributes:
    values: The values of the indices. Can be an array or a ref.
    ignored_value: If not None, the indices with this value will be ignored.
  """

  values: Any
  ignored_value: int | None = jax.tree.static(default=None)

  def pretty_print(
      self, context: jax_core.JaxprPpContext, *, print_dtype: bool = True
  ) -> str:
    values = self.values.pretty_print(context, print_dtype=print_dtype)
    if self.ignored_value is None:
      return values
    return f"{values}~{self.ignored_value}"


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True, slots=True)
class SharedMemRefSliceTransform(state_types.Transform):
  """Slices a VMEM_SHARED ref along its minor dimension into subcore VMEM."""

  def transform_type(self, x: jax_core.AbstractValue) -> jax_core.AbstractValue:
    match x:
      case state_types.AbstractRef():
        memory_space = x.memory_space
        if isinstance(memory_space, pallas_core.CoreMemorySpace):
          memory_space = memory_space.memory_space
        if memory_space != tpu_core.MemorySpace.VMEM_SHARED:
          raise ValueError(
              "slice_subcore_vmem_from_shared requires a VMEM_SHARED"
              f" reference, got memory space: {x.memory_space}"
          )
        return x.update(
            inner_aval=self.transform_type(x.inner_aval),
            memory_space=tpu_core.MemorySpace.VMEM,
        )
      case jax_core.ShapedArray():
        if not x.shape:
          raise ValueError(
              "Cannot slice subcore VMEM from a 0D VMEM_SHARED reference."
          )
        num_subcores = get_sparse_core_info().num_subcores
        if x.shape[-1] % num_subcores != 0:
          raise ValueError(
              f"VMEM_SHARED last dimension ({x.shape[-1]}) must be divisible "
              "by the total number of physical subcores, including inactive "
              f"ones ({num_subcores})."
          )
        new_shape = (*x.shape[:-1], x.shape[-1] // num_subcores)
        return x.update(shape=new_shape)
      case _:
        raise TypeError(f"Cannot slice subcore VMEM from {x}")

  def pretty_print(self, context: jax_core.JaxprPpContext) -> pp.Doc:
    del context
    return pp.text("{shared_memref_slice}")