"""``torch.compile`` backend for the Neuron decode graphs: torch-neuronx's ``neuron`` backend plus fp8 legalization.

trn2 has no ``F8E4M3FN`` (neuronx-cc rejects any graph holding one; issue #24), but the model carries fp8 weights,
fp8 KV records and fp8-quantized activations. This backend makes every fp8 value a ``uint8`` before lowering:

* fp8 graph inputs are passed as zero-copy ``uint8`` views;
* a pass over the Dynamo graph, ahead of AOT, carries every fp8 value as ``uint8``.
  Data movement (views, slices, cat, index, copy, where, ...) is byte-exact as is; ``fp8 -> float`` becomes an
  exact 256-entry table lookup and ``float -> fp8`` an arithmetic round-to-nearest-even encode with torch's
  saturating semantics (|x| > 448 and inf -> 448, NaN -> 0x7F). Negative zero encodes as +0 (same value).

Usage: ``torch.compile(fn, backend="deepinfer_neuron")`` (registered on import).
"""

from __future__ import annotations

import os
import operator
from collections.abc import Callable, Sequence
from typing import Any

import torch
import torch.fx
from torch._dynamo import register_backend
from torch._subclasses.fake_tensor import unset_fake_temporarily

FP8 = torch.float8_e4m3fn
U8 = torch.uint8


def _is_fp8(t: Any) -> bool:
    return isinstance(t, torch.Tensor) and t.dtype == FP8


def _encode_scalar(value: float) -> int:
    with unset_fake_temporarily():
        return int(torch.tensor(float(value)).to(FP8).view(U8).item())


# --------------------------------------------------------------------------- #
# Reference codec (traced into the graph)
# --------------------------------------------------------------------------- #


def _decode(x_u8: torch.Tensor) -> torch.Tensor:
    """Exact e4m3fn bits -> fp32, elementwise (no table, no bitcast): (8 + m) * 2^(e - 10), subnormals m * 2^-9.

    Float arithmetic only (Vector / Scalar engines): integer shifts and masks run on GpSimd, and this decode is in
    every graph that reads an fp8 activation or cache. ``exp2`` is exact on integers on the device (see
    nki_attention._pow2)."""
    b = x_u8.to(torch.float32)
    neg = b >= 128.0
    mag = torch.where(neg, b - 128.0, b)
    e = torch.floor(mag * 0.125)
    m = mag - 8.0 * e
    v = torch.where(e == 0.0, m * 2.0**-9, (m + 8.0) * torch.exp2(e - 10.0))
    v = torch.where(neg, -v, v)
    return torch.where(mag == 127.0, torch.full_like(v, float("nan")), v)


def _decode_as(dtype: torch.dtype) -> Callable[[torch.Tensor], torch.Tensor]:
    def decode(x_u8: torch.Tensor) -> torch.Tensor:
        return _decode(x_u8).to(dtype)

    return decode


def _rne(m: torch.Tensor) -> torch.Tensor:
    f = torch.floor(m)
    d = m - f
    odd = (f - 2.0 * torch.floor(f * 0.5)) == 1.0  # not torch.remainder: it lowers to GpSimd
    return f + ((d > 0.5) | ((d == 0.5) & odd)).to(m.dtype)


def _encode(x: torch.Tensor) -> torch.Tensor:
    """fp32 -> e4m3fn bits (uint8), round-to-nearest-even, saturating like torch's cast (no table, no bitcast)."""
    return encode_codes(x).to(U8)


def encode_codes(x: torch.Tensor) -> torch.Tensor:
    """``_encode`` as int32 codes: byte assembly should stay in int32 inside Neuron graphs and narrow once at the
    store (uint8 concats of narrowed values crash neuronx-cc or corrupt bytes, issue #25)."""
    xf = x.float()
    nan = xf != xf
    neg = xf < 0
    a = torch.clamp(torch.abs(torch.where(nan, torch.zeros_like(xf), xf)), max=448.0)
    # e = floor(log2(a)) and 2^(3 - e) for a in [2^-6, 448], by comparisons: exact, no log/exp on the device.
    e = torch.full_like(a, -6.0)
    scale = torch.full_like(a, 2.0**9)
    for k in range(-5, 9):
        ge = a >= 2.0**k
        e = e + ge.to(a.dtype)
        scale = scale * torch.where(ge, 0.5, 1.0)
    r = _rne(a * scale)  # [8, 16]; 16 carries into the next exponent below
    code_n = torch.clamp((e + 7.0) * 8.0 + (r - 8.0), max=126.0)
    code_s = _rne(a * 512.0)  # subnormals: a / 2^-9, 8 = smallest normal
    code = torch.where(a >= 2.0**-6, code_n, code_s) + neg.to(a.dtype) * 128.0
    code = torch.where(nan, torch.full_like(code, 127.0), code)
    return code.to(torch.int32)


# --------------------------------------------------------------------------- #
# Graph pass (Dynamo / torch level, ahead of AOT: the codecs inserted here are traced with the graph's devices)
# --------------------------------------------------------------------------- #

_CREATE = {torch.empty, torch.zeros, torch.full, torch.empty_like, torch.zeros_like, torch.full_like, torch.tensor,
           torch.empty_strided, torch.ones, torch.ones_like}
_CREATE_METHODS = {"new_empty", "new_zeros", "new_full", "new_ones", "new_tensor", "new_empty_strided"}
_TO_METHODS = {"float": torch.float32, "bfloat16": torch.bfloat16, "half": torch.float16, "double": torch.float64}
_FILL_ARG = {torch.full: 1, torch.full_like: 1, "new_full": 1, "fill_": 0, "masked_fill": 1, "masked_fill_": 1}


def _ev(node: Any) -> Any:
    return node.meta.get("example_value") if isinstance(node, torch.fx.Node) else node


def _dtype_arg(node: torch.fx.Node, positional: int | None) -> Any:
    if "dtype" in node.kwargs:
        return node.kwargs["dtype"]
    if positional is not None and len(node.args) > positional and isinstance(node.args[positional], torch.dtype):
        return node.args[positional]
    return None


def legalize_fp8(gm: torch.fx.GraphModule) -> bool:
    """Carry every fp8 value of a Dynamo graph as uint8 bytes. Returns True if the graph changed.

    Only dtype-sensitive operations are rewritten (bitcasts, conversions, creation with dtype=fp8, implicit
    conversions in copy_/setitem, fp8 fill values); everything else on fp8 (views, slicing, cat, indexing, where
    between fp8 tensors, ...) moves bytes and runs unchanged on the uint8 carrier.
    """
    graph = gm.graph
    carriers: set[torch.fx.Node] = set()  # nodes whose fp8 value is now a uint8 tensor
    changed = False

    def is_fp8(a: Any) -> bool:
        return a in carriers or _is_fp8(_ev(a))

    for node in list(graph.nodes):
        if node.op not in ("call_function", "call_method"):
            continue
        tgt, args = node.target, node.args
        out_fp8 = _is_fp8(_ev(node)) or (
            isinstance(_ev(node), (tuple, list)) and any(_is_fp8(v) for v in _ev(node))
        )
        with graph.inserting_before(node):
            # Bitcasts: x.view(dtype)
            if node.op == "call_method" and tgt == "view" and len(args) == 2 and isinstance(args[1], torch.dtype):
                src, dst = args
                # Always a fresh node (a same-dtype view is a no-op): a node is a carrier or a plain uint8 tensor,
                # never both, even when the bitcast itself becomes an identity.
                if dst == FP8:
                    repl = graph.call_method("view", (src, U8))
                    _swap(node, repl, carriers, carrier=True)
                elif is_fp8(src):
                    repl = graph.call_method("view", (src, dst))
                    _swap(node, repl, carriers, carrier=False)
                changed = True
                continue
            # Conversions: x.to(dtype) / x.float() / x.type(dtype) ...
            dst = None
            if node.op == "call_method" and tgt in _TO_METHODS:
                dst = _TO_METHODS[tgt]
            elif node.op == "call_method" and tgt in ("to", "type"):
                dst = _dtype_arg(node, 1)
                if dst is None and len(args) > 2 and isinstance(args[2], torch.dtype):
                    dst = args[2]
            if dst is not None:
                src = args[0]
                if is_fp8(src) and dst != FP8:
                    repl = graph.call_function(_decode_as(dst), (src,))
                    _swap(node, repl, carriers, carrier=False)
                elif dst == FP8 and not is_fp8(src):
                    repl = graph.call_function(_encode, (src,))
                    _swap(node, repl, carriers, carrier=True)
                elif is_fp8(src) and dst == FP8:  # fp8 -> fp8 (copy / device move): bytes
                    _set_dtype(node, U8)
                    carriers.add(node)
                else:  # no fp8 involved
                    continue
                changed = True
                continue
            # Creation with dtype=fp8
            if (node.op == "call_function" and tgt in _CREATE) or (node.op == "call_method" and tgt in _CREATE_METHODS):
                if _dtype_arg(node, None) == FP8 or (out_fp8 and "dtype" not in node.kwargs and is_fp8(args[0] if args else None)):
                    if "dtype" in node.kwargs:
                        node.kwargs = {**node.kwargs, "dtype": U8}
                    _encode_fill(node)
                    carriers.add(node)
                    changed = True
                continue
            # Implicit conversions on copy_ / setitem into an fp8 tensor (or out of one).
            if (node.op == "call_method" and tgt == "copy_") or (node.op == "call_function" and tgt is operator.setitem):
                dst_t, src = args[0], args[-1]
                if isinstance(src, torch.fx.Node) and is_fp8(dst_t) != is_fp8(src):
                    conv = graph.call_function(_encode, (src,)) if is_fp8(dst_t) else graph.call_function(_decode_as(_ev(dst_t).dtype), (src,))
                    node.args = (*args[:-1], conv)
                    changed = True
                elif is_fp8(dst_t) and isinstance(src, (int, float)):
                    node.args = (*args[:-1], _encode_scalar(src))
                    changed = True
                if is_fp8(dst_t):
                    carriers.add(node)
                continue
            if node.op == "call_method" and tgt in ("fill_", "masked_fill", "masked_fill_") and is_fp8(args[0]):
                _encode_fill(node)
                carriers.add(node)
                changed = True
                continue
            if out_fp8 or any(is_fp8(a) for a in node.all_input_nodes):
                if out_fp8:
                    carriers.add(node)
                changed = True

    if changed:
        graph.lint()
        gm.recompile()
    return changed


def _swap(node: torch.fx.Node, repl: torch.fx.Node, carriers: set, carrier: bool) -> None:
    node.replace_all_uses_with(repl)
    if carrier:
        carriers.add(repl)
    node.graph.erase_node(node)


def _set_dtype(node: torch.fx.Node, dtype: torch.dtype) -> None:
    if "dtype" in node.kwargs:
        node.kwargs = {**node.kwargs, "dtype": dtype}
        return
    node.args = tuple(dtype if a == FP8 else a for a in node.args)


def _encode_fill(node: torch.fx.Node) -> None:
    key = node.target
    if key not in _FILL_ARG:
        if key in (torch.ones, torch.ones_like) or key == "new_ones":
            raise NotImplementedError(f"fp8 legalization: {key} with fp8")
        return
    i = _FILL_ARG[key] + (1 if node.op == "call_method" else 0)
    if i < len(node.args) and isinstance(node.args[i], (int, float)):
        a = list(node.args)
        a[i] = _encode_scalar(a[i])
        node.args = tuple(a)
    elif "fill_value" in node.kwargs and isinstance(node.kwargs["fill_value"], (int, float)):
        node.kwargs = {**node.kwargs, "fill_value": _encode_scalar(node.kwargs["fill_value"])}
    elif "value" in node.kwargs and isinstance(node.kwargs["value"], (int, float)):
        node.kwargs = {**node.kwargs, "value": _encode_scalar(node.kwargs["value"])}


# --------------------------------------------------------------------------- #
# Backend
# --------------------------------------------------------------------------- #


def _fp8_inputs_as_bytes(gm: torch.fx.GraphModule, example_inputs: Sequence[Any]) -> tuple[list[Any], list[int]]:
    placeholders = [n for n in gm.graph.nodes if n.op == "placeholder"]
    positions = [i for i, x in enumerate(example_inputs) if _is_fp8(x)]
    if not positions:
        return list(example_inputs), []
    first_body = next(n for n in gm.graph.nodes if n.op != "placeholder")
    for i in positions:
        ph = placeholders[i]
        with gm.graph.inserting_before(first_body):
            as_fp8 = gm.graph.call_method("view", (ph, FP8))
        ph.replace_all_uses_with(as_fp8, delete_user_cb=lambda user: user is not as_fp8)
        if "example_value" in ph.meta:
            fake = ph.meta["example_value"]
            as_fp8.meta["example_value"] = fake
            with fake.fake_mode:  # same fake mode, so the static shape (and its source) carry over
                ph.meta["example_value"] = fake.view(U8)
    gm.recompile()
    # Fresh (non-view) uint8 fakes: a view of the fp8 input reaches the lowering with symbolic sizes.
    fake_mode = next(ph.meta["example_value"].fake_mode for ph in placeholders if "example_value" in ph.meta)

    def as_bytes(x: torch.Tensor) -> torch.Tensor:
        with fake_mode:
            return torch.empty_strided(x.shape, x.stride(), dtype=U8, device=x.device)

    inputs = [as_bytes(x) if i in positions else x for i, x in enumerate(example_inputs)]
    return inputs, positions


def _print_aliased_inputs(gm: torch.fx.GraphModule, example_inputs: Sequence[Any]) -> None:
    """Debug: graph inputs sharing a storage (aot_autograd rejects mutated aliases of different dtypes)."""
    from torch.multiprocessing.reductions import StorageWeakRef

    names = [n.target for n in gm.graph.nodes if n.op == "placeholder"]
    groups: dict[Any, list[str]] = {}
    for name, x in zip(names, example_inputs):
        if isinstance(x, torch.Tensor):
            key = StorageWeakRef(x.untyped_storage())
            groups.setdefault(key, []).append(f"{name}:{x.dtype}{tuple(x.shape)}@{x.storage_offset()}")
    for members in groups.values():
        if len(members) > 1:
            print("ALIASED INPUTS:", members, flush=True)


@register_backend
def deepinfer_neuron(gm: torch.fx.GraphModule, example_inputs: Sequence[Any], **kwargs: Any) -> Callable:
    from torch_neuronx.neuron_dynamo_backend.backend import neuron_backend

    if os.environ.get("DEEPINFER_NEURON_ALIAS_DEBUG"):
        _print_aliased_inputs(gm, example_inputs)
    inputs, positions = _fp8_inputs_as_bytes(gm, example_inputs)
    if legalize_fp8(gm) and os.environ.get("DEEPINFER_FP8_LEGALIZE_DEBUG"):
        print(gm.graph, flush=True)
    compiled = neuron_backend(gm, inputs, **kwargs)
    if not positions:
        return compiled
    pos = set(positions)

    def run(*args: Any) -> Any:
        return compiled(*[a.view(U8) if i in pos else a for i, a in enumerate(args)])

    return run