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

"""Custom fusible dtypes."""

import abc
import dataclasses
import functools
import itertools as it
from typing import Any, TypeVar
from collections.abc import Callable, Sequence

import jax
from jax._src import api_util
from jax._src import core
from jax._src import custom_derivatives
from jax._src import dtypes
from jax._src import linear_util as lu
from jax._src import source_info_util
from jax._src import state
from jax._src import tree_util
from jax._src import util
from jax._src.interpreters import partial_eval as pe
from jax._src.lax.control_flow import conditionals
from jax._src.pallas import core as pallas_core
from jax._src.pallas import pallas_call
from jax._src.pallas import primitives as pallas_primitives
from jax._src.pallas.fuser import block_spec
from jax._src.pallas.fuser.fusible import fusible_p
from jax._src.state import discharge as state_discharge
from jax._src.state import primitives as state_primitives
from jax._src.util import foreach

map, unsafe_map = util.safe_map, map
zip, unsafe_zip = util.safe_zip, zip

T = TypeVar("T")

_physicalize_rules: dict[core.Primitive, Callable[..., Any]] = {}

pack_dtype_p = core.Primitive("pack_dtype")


@pack_dtype_p.def_abstract_eval
def pack_dtype_abstract_eval(*xs, dtype):
  if dtypes.issubdtype(dtype, FusibleElementDType):
    return dtype.abstract_pack(*xs)
  raise ValueError("Attempted to pack non-fusion dtype: {dtype}")


def pack(*xs, dtype):
  return pack_dtype_p.bind(*xs, dtype=dtype)


unpack_dtype_p = core.Primitive("unpack_dtype")
unpack_dtype_p.multiple_results = True


@unpack_dtype_p.def_abstract_eval
def unpack_dtype_abstract_eval(x):
  if dtypes.issubdtype(x.dtype, FusibleElementDType):
    return x.dtype.abstract_unpack(x)
  elif isinstance(x.dtype, state.AbstractRef):
    raise NotImplementedError()
  raise ValueError("Attempted to unpack non-fusion dtype: {dtype}")


def unpack(x):
  return tuple(unpack_dtype_p.bind(x))


class FusibleElementDType(dtypes.extended):
  """Scalar dtype for fusible dtypes."""


class FusibleTyRules:
  allow_conversion: bool = False


class FusionDType(dtypes.ExtendedDType, util.StrictABC):
  """Base class for fusible extended dtypes."""

  _op_registry = {}
  _rules = FusibleTyRules
  type = FusibleElementDType

  @abc.abstractmethod
  def abstract_unpack(self, x) -> Sequence[Any]:
    raise NotImplementedError()

  @abc.abstractmethod
  def abstract_pack(self, *xs):
    raise NotImplementedError()

  @classmethod
  def register_op(cls, primitive):
    def _register_fn(fn):
      cls._op_registry[primitive] = fn

    return _register_fn

  @classmethod
  def get_op_rule(cls, primitive):
    return cls._op_registry.get(primitive)

  @property
  def name(self):
    return str(self)

  @abc.abstractmethod
  def pull_block_spec_one_step(self, aval_out, *args, **kwargs):
    raise NotImplementedError()

  @abc.abstractmethod
  def unpack_push_block_spec(self, aval_in, *args, **kwargs):
    raise NotImplementedError()

  @abc.abstractmethod
  def unpack_pull_block_spec(self, aval_in, *args, **kwargs):
    raise NotImplementedError()


def physicalize(f):
  """Runs a function that contains fusible extended dtypes."""

  def wrapper(*args, **kwargs):
    if kwargs:
      raise NotImplementedError()
    flattened_args, treedef = jax.tree.flatten(args)
    debug_info = api_util.debug_info("physicalize", f, args, kwargs)
    wrapped_fun, out_tree_thunk = api_util.flatten_fun_nokwargs(
        lu.wrap_init(f, debug_info=debug_info), treedef
    )
    avals = [core.typeof(a) for a in flattened_args]
    jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, avals)
    new_jaxpr = physicalize_closed_jaxpr(core.ClosedJaxpr(jaxpr, consts))
    out_flat = core.eval_jaxpr(
        new_jaxpr.jaxpr, new_jaxpr.consts, *flattened_args
    )
    return tree_util.tree_unflatten(out_tree_thunk(), out_flat)

  return wrapper


@util.weakref_lru_cache
def physicalize_closed_jaxpr(jaxpr: core.ClosedJaxpr) -> core.ClosedJaxpr:
  """Replaces all extended dtypes with physical types in a jaxpr."""
  fun = functools.partial(physicalize_interp, jaxpr.jaxpr, jaxpr.consts)
  in_avals = [_physical_aval(aval) for aval in jaxpr.in_avals]
  flat_avals, treedef = tree_util.tree_flatten(in_avals)
  debug_info = api_util.debug_info("physicalize_closed_jaxpr", fun, (), {})
  wrapped_fun, _ = api_util.flatten_fun_nokwargs(
      lu.wrap_init(fun, debug_info=debug_info.with_unknown_names()), treedef
  )
  new_jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, flat_avals)
  assert len(new_jaxpr.constvars) == len(consts), "Mismatched consts"
  return core.ClosedJaxpr(new_jaxpr, consts)


def _physical_aval(aval):
  if isinstance(aval, core.ShapedArray):
    if isinstance(aval.dtype, FusionDType):
      return aval.dtype.abstract_unpack(aval)
    return core.ShapedArray(aval.shape, aval.dtype)
  if isinstance(aval, state.AbstractRef):
    if _is_fusion_type(aval):
      unpacked = aval.dtype.abstract_unpack(aval.inner_aval)
      return tuple(aval.update(inner_aval=u) for u in unpacked)
    return aval
  return aval


def physicalize_jaxpr(jaxpr: core.Jaxpr) -> core.Jaxpr:
  """Replaces all extended dtypes with physical types in a jaxpr."""

  def _flat_jaxpr_eval(consts, args):
    return physicalize_interp(jaxpr, consts, *args)

  in_avals = [_physical_aval(v.aval) for v in jaxpr.invars]
  const_avals = [_physical_aval(v.aval) for v in jaxpr.constvars]
  flat_avals, treedef = jax.tree.flatten((const_avals, in_avals))
  debug_info = api_util.debug_info(
      "physicalize_jaxpr", _flat_jaxpr_eval, (const_avals, in_avals), {})
  wrapped_fun, _ = api_util.flatten_fun_nokwargs(
      lu.wrap_init(_flat_jaxpr_eval, debug_info=debug_info), treedef
  )
  new_jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, flat_avals)
  assert not consts
  new_jaxpr = pe.convert_invars_to_constvars(
      new_jaxpr, len(tree_util.tree_leaves(const_avals))
  )
  return new_jaxpr


@dataclasses.dataclass
class Context:
  avals_in: Sequence[Any]
  avals_out: Sequence[Any]


def physicalize_interp(
    jaxpr: core.Jaxpr, consts: Sequence[core.Value], *args: core.Value
):
  """Physicalizes a jaxpr by replacing fusible dtypes with physical types."""
  # TODO: Merge into JAX core.
  env: dict[core.Var, Any] = {}

  def read_env(var: core.Atom):
    if isinstance(var, core.Literal):
      return var.val
    return env[var]

  def write_env(var: core.Var, val: Any):
    env[var] = val

  foreach(write_env, jaxpr.constvars, consts)
  assert len(jaxpr.invars) == len(
      args
  ), f"Length mismatch: {jaxpr.invars} != {args}"
  foreach(write_env, jaxpr.invars, args)

  for eqn in jaxpr.eqns:
    invals = list(map(read_env, eqn.invars))
    avals_in = tuple(x.aval for x in eqn.invars)
    name_stack = (
        source_info_util.current_name_stack() + eqn.source_info.name_stack
    )
    with (
        source_info_util.user_context(
            eqn.source_info.traceback, name_stack=name_stack
        ),
        eqn.ctx.manager,
    ):
      # need to check types and then invoke the correct rule.
      ctx = Context(
          avals_in=avals_in, avals_out=[var.aval for var in eqn.outvars]
      )
      custom_rule = _phys_find_rule(eqn.primitive, avals_in)
      if custom_rule:
        outvals = custom_rule(ctx, *invals, **eqn.params)
      else:
        bind_params = eqn.primitive.get_bind_params(eqn.params)
        outvals = eqn.primitive.bind(*invals, **bind_params)

    if eqn.primitive.multiple_results:
      assert len(outvals) == len(eqn.outvars), eqn
      foreach(write_env, eqn.outvars, outvals)
    else:
      write_env(eqn.outvars[0], outvals)

  return map(read_env, jaxpr.outvars)


def _is_fusion_type(aval: core.AbstractValue):
  """Returns whether an aval is an array containing fusion types."""
  return (
      isinstance(aval, (core.ShapedArray, state.AbstractRef))
      and hasattr(aval, 'dtype')
      and isinstance(aval.dtype, FusionDType)
  )


def _phys_find_rule(primitive, avals: Sequence[core.AbstractValue]):
  """Finds the physicalization rule for a primitive."""
  if primitive in _physicalize_rules:
    return _physicalize_rules[primitive]

  # pyrefly: ignore[missing-attribute]
  fusion_types = {aval.dtype for aval in avals if _is_fusion_type(aval)}
  if len(fusion_types) == 0:
    return None
  elif len(fusion_types) > 1:
    raise ValueError(f"Multiple fusion types for primitive: {fusion_types}")
  fusion_type = fusion_types.pop()
  if primitive not in fusion_type._op_registry:
    raise ValueError(
        f"No implementation found for primitive {primitive} "
        f"for custom type {fusion_type}"
    )
  return fusion_type.get_op_rule(primitive)


def _assert_no_fusion_types(avals: Sequence[core.AbstractValue]):
  if any(_is_fusion_type(aval) for aval in avals):
    raise NotImplementedError(f"Fusion type found in avals: {avals}")


def _pallas_call_physicalize_rule(
    ctx: Context, *args, jaxpr, grid_mapping: pallas_core.GridMapping, **kwargs
):
  _assert_no_fusion_types(ctx.avals_in)
  _assert_no_fusion_types(ctx.avals_out)
  with grid_mapping.trace_env():
    new_jaxpr = physicalize_closed_jaxpr(core.ClosedJaxpr(jaxpr, ()))
  if diff := len(new_jaxpr.jaxpr.invars) - len(jaxpr.invars):
    num_scratch_avals = len(grid_mapping.scratch_avals) + diff
    new_scratch_avals = tuple(v.aval for v in
                              new_jaxpr.jaxpr.invars[-num_scratch_avals:])
    grid_mapping = grid_mapping.replace(
        scratch_avals=new_scratch_avals
    )
  return pallas_call.pallas_call_p.bind(
      *args, jaxpr=new_jaxpr.jaxpr, grid_mapping=grid_mapping, **kwargs
  )


_physicalize_rules[pallas_call.pallas_call_p] = _pallas_call_physicalize_rule


def _cond_physicalize_rule(ctx: Context, *args, branches, **kwargs):
  _assert_no_fusion_types(ctx.avals_out)
  physicalized_branches = tuple(
      physicalize_closed_jaxpr(branch) for branch in branches
  )
  flat_args = jax.tree.leaves(args)
  return conditionals.cond_p.bind(
      *flat_args, branches=physicalized_branches, **kwargs
  )


_physicalize_rules[conditionals.cond_p] = _cond_physicalize_rule


@lu.transformation2
def _physicalize_transform(f, *args):
  vals, zeros = args[::2], args[1::2]
  assert len(vals) == len(zeros)
  wrapper = lambda *inner_vals: f(
      *it.chain.from_iterable(zip(inner_vals, zeros))
  )
  return physicalize(wrapper)(*vals)


@lu.transformation2
def _physicalize_transform_bwd(f, const_avals, *args):
  return [custom_derivatives.Zero(a) for a in const_avals] + list(
      physicalize(f)(*args)
  )


def _custom_vjp_call_physicalize_rule(
    ctx: Context, *args, call_jaxpr, num_consts, fwd_jaxpr_thunk, bwd, **kwargs
):
  _assert_no_fusion_types(ctx.avals_out)
  new_jaxpr = physicalize_closed_jaxpr(call_jaxpr)
  fun = lu.wrap_init(core.jaxpr_as_fun(new_jaxpr),
                     debug_info=call_jaxpr.jaxpr.debug_info)
  fwd = custom_derivatives.lift_fwd(num_consts, fwd_jaxpr_thunk)
  fwd_physicalized = _physicalize_transform(fwd)
  const_avals, _ = util.split_list(new_jaxpr.in_avals, [num_consts])
  bwd_physicalized = _physicalize_transform_bwd(bwd, const_avals)
  kwargs['subfuns'] = (fun, fwd_physicalized, bwd_physicalized)
  return custom_derivatives.custom_vjp_call_p.bind(*args, **kwargs)

_physicalize_rules[custom_derivatives.custom_vjp_call_p] = _custom_vjp_call_physicalize_rule


def _run_state_rule(ctx: Context, *args, jaxpr, which_linear, is_initialized):
  _assert_no_fusion_types(ctx.avals_in)
  _assert_no_fusion_types(ctx.avals_out)
  jaxpr = physicalize_jaxpr(jaxpr)
  return state_discharge.run_state_p.bind(
      *args,
      jaxpr=jaxpr,
      which_linear=which_linear,
      is_initialized=is_initialized,
  )


_physicalize_rules[state_discharge.run_state_p] = _run_state_rule


def _core_map_rule(ctx: Context, *args, jaxpr, **params):
  _assert_no_fusion_types(ctx.avals_in)
  _assert_no_fusion_types(ctx.avals_out)
  assert not jaxpr.invars
  with core.extend_axis_env_nd(params["mesh"].shape.items()):
    jaxpr = physicalize_jaxpr(jaxpr)
  return pallas_core.core_map_p.bind(*args, jaxpr=jaxpr, **params)


_physicalize_rules[pallas_core.core_map_p] = _core_map_rule


def _run_scoped_rule(ctx: Context, *args, jaxpr, **params):
  _assert_no_fusion_types(ctx.avals_out)
  jaxpr = physicalize_jaxpr(jaxpr)
  flat_args = tree_util.tree_leaves(args)
  assert len(flat_args) == len(
      jaxpr.constvars
  ), f"Length mismatch: {len(flat_args)=} != {len(jaxpr.constvars)=}"
  return pallas_primitives.run_scoped_p.bind(*flat_args, jaxpr=jaxpr, **params)


_physicalize_rules[pallas_primitives.run_scoped_p] = _run_scoped_rule


def _scan_rule(ctx: Context, *args, jaxpr, **params):
  _assert_no_fusion_types(ctx.avals_in)
  _assert_no_fusion_types(ctx.avals_out)
  jaxpr = physicalize_closed_jaxpr(jaxpr)
  return jax.lax.scan_p.bind(*args, jaxpr=jaxpr, **params)


_physicalize_rules[jax.lax.scan_p] = _scan_rule


def _while_rule(
    ctx: Context, *args, body_jaxpr, cond_jaxpr, body_nconsts,
    cond_nconsts, **params
):
  _assert_no_fusion_types(ctx.avals_out)
  cond_avals = [v.aval for v in cond_jaxpr.jaxpr.invars]
  _, cond_in_avals = util.split_list(cond_avals, [cond_nconsts])
  _assert_no_fusion_types(cond_in_avals)
  new_cond_jaxpr = physicalize_closed_jaxpr(cond_jaxpr)
  new_num_cond_consts = (
      cond_nconsts
      + len(new_cond_jaxpr.jaxpr.invars)
      - len(cond_jaxpr.jaxpr.invars)
  )

  body_avals = [v.aval for v in body_jaxpr.jaxpr.invars]
  _, body_in_avals = util.split_list(body_avals, [body_nconsts])
  _assert_no_fusion_types(body_in_avals)
  new_body_jaxpr = physicalize_closed_jaxpr(body_jaxpr)
  new_num_body_consts = (
      body_nconsts
      + len(new_body_jaxpr.jaxpr.invars)
      - len(body_jaxpr.jaxpr.invars)
  )
  flat_args = tree_util.tree_leaves(args)
  cond_consts, body_consts, flat_args = util.split_list(
      flat_args, [new_num_cond_consts, new_num_body_consts]
  )
  assert len(flat_args) + len(body_consts) == len(
      new_body_jaxpr.jaxpr.invars), (
      f"Length mismatch: {len(flat_args) + len(body_consts)} !="
      f" {len(new_body_jaxpr.jaxpr.invars)=}"
  )
  assert len(flat_args) + len(cond_consts) == len(
      new_cond_jaxpr.jaxpr.invars), (
      f"Length mismatch: {len(flat_args) + len(cond_consts)} !="
      f" {len(new_cond_jaxpr.jaxpr.invars)=}"
  )
  return jax.lax.while_p.bind(
      *(cond_consts + body_consts + flat_args),
      body_jaxpr=new_body_jaxpr,
      cond_jaxpr=new_cond_jaxpr,
      body_nconsts=new_num_body_consts,
      cond_nconsts=new_num_cond_consts,
      **params,
  )


_physicalize_rules[jax.lax.while_p] = _while_rule


def _pack_rule(_, *args, dtype):
  del dtype
  return args


_physicalize_rules[pack_dtype_p] = _pack_rule


def _unpack_rule(_, arg):
  return arg


_physicalize_rules[unpack_dtype_p] = _unpack_rule


def _swap_rule(ctx: Context, ref, val, *args, tree):
  ref_aval, *_ = ctx.avals_in
  if not _is_fusion_type(ref_aval):
    return state_primitives.swap_p.bind(ref, val, *args, tree=tree)
  return ref_aval.dtype.swap(ref, val, *args, tree=tree)


_physicalize_rules[state_primitives.swap_p] = _swap_rule


def _get_rule(ctx: Context, ref, *args, tree):
  ref_aval, *_ = ctx.avals_in
  if not _is_fusion_type(ref_aval):
    return state_primitives.get_p.bind(ref, *args, tree=tree)
  return ref_aval.dtype.get(ref, *args, tree=tree)


_physicalize_rules[state_primitives.get_p] = _get_rule


@block_spec.register_eval_rule(pack_dtype_p)
def _pack_dtype_eval_rule(ctx: block_spec.KernelEvalContext, *args, dtype):
  return dtype.pack_eval_rule(ctx, *args)


@block_spec.register_pull_block_spec_rule(pack_dtype_p)
def _pack_dtype_pull_rule(
    ctx: block_spec.PullRuleContext,
    block_spec: pallas_core.BlockSpec,
    *,
    dtype: FusionDType,
):
  aval_out = ctx.avals_out[0]
  return dtype.pull_block_spec_one_step(aval_out, block_spec)


@block_spec.register_push_block_spec_rule(unpack_dtype_p)
def _unpack_dtype_push_rule(
    ctx: block_spec.PushRuleContext,
    block_spec: pallas_core.BlockSpec,
):
  aval_in = ctx.avals_in[0]
  assert isinstance(aval_in, core.ShapedArray)
  assert isinstance(aval_in.dtype, FusionDType), aval_in.dtype
  return aval_in.dtype.unpack_push_block_spec(aval_in, block_spec)


@block_spec.register_pull_block_spec_rule(unpack_dtype_p)
def _unpack_dtype_pull_rule(
    ctx: block_spec.PushRuleContext,
    block_specs: pallas_core.BlockSpec,
):
  aval_in = ctx.avals_in[0]
  assert isinstance(aval_in, core.ShapedArray)
  assert isinstance(aval_in.dtype, FusionDType), aval_in.dtype
  return aval_in.dtype.unpack_pull_block_spec(aval_in, *block_specs)  # pyrefly: ignore[not-iterable]


@block_spec.register_eval_rule(unpack_dtype_p)
def _unpack_dtype_eval_rule(ctx: block_spec.KernelEvalContext, *args):
  assert ctx.avals_in is not None
  aval_in = ctx.avals_in[0]
  assert isinstance(aval_in, core.ShapedArray)
  assert isinstance(aval_in.dtype, FusionDType), aval_in.dtype
  return aval_in.dtype.unpack_eval_rule(ctx, *args)  # pyrefly: ignore[missing-attribute]


def _fusible_physicalize_rule(
    _, *consts_and_args, jaxpr, num_consts, in_tree, out_tree, func
):
  consts, _ = util.split_list(consts_and_args, [num_consts])
  new_jaxpr = physicalize_closed_jaxpr(core.ClosedJaxpr(jaxpr, consts))
  return fusible_p.bind(
      *consts_and_args,
      jaxpr=new_jaxpr.jaxpr,
      num_consts=num_consts,
      in_tree=in_tree,
      out_tree=out_tree,
      func=func,
  )


_physicalize_rules[fusible_p] = _fusible_physicalize_rule
