"""Make transformers' block-FP8 linears differentiable for LoRA fine-tuning.

The on-disk GLM-5.2 checkpoint stores weights as block-FP8 (e4m3, [128,128]
block scales). transformers' FP8 matmul kernels have NO autograd formula, so
loss.backward() raises:
  "Trying to backward through ...w8a8_block_dynamic_fp8_matmul... no autograd formula".

For LoRA the base weight is FROZEN: we only need grad_INPUT to flow (no weight grad).
We replace the FP8 matmul paths with a differentiable path that, WHEN GRAD IS NEEDED,
dequantizes the FP8 weight on the fly (transient, freed after the op) and runs a normal
F.linear / bmm. When no grad is needed (inference, or grad-checkpoint no_grad fwd) we
keep the fast native FP8 kernel.

Key lever: patch module-level `fp8_linear`. Python resolves module globals at call time,
so FP8Linear.forward AND the eager FP8Experts.linear (both call `fp8_linear(...)`) pick
up the patched version automatically. We also patch the grouped/batched experts-interface
paths in case the model dispatches there.

`_need_grad(input)` gate => only layers downstream of the earliest LoRA adapter run the
(slower) differentiable path; earlier layers keep the fast FP8 kernel. So a LATE-band
LoRA keeps both compute and transient memory bounded.

Dequant mirrors FineGrainedFP8HfQuantizer._dequantize_one:
  W_dq[r,c] = W_fp8[r,c] * scale[r//block_m, c//block_n]
"""
import torch
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

_FP8_DTYPE = torch.float8_e4m3fn

class _FP8LinearFn(torch.autograd.Function):
    """Differentiable block-FP8 linear for a FROZEN base weight.

    Forward: dequant W (bf16, transient) -> F.linear -> FREE Wdq (only the small
    output activation is kept in the graph). Backward: RECOMPUTE Wdq transiently
    and return grad_input = grad_out @ Wdq. No weight grad (base is frozen), so the
    bf16 weight is never stored across the whole forward -> bounded memory.
    The fp8 weight + scale are resident base params (saved by reference, ~free).
    """
    @staticmethod
    def forward(ctx, x, weight_fp8, scale_inv, bias, block_size, out_dtype):
        Wdq = dequant_block_fp8(weight_fp8, scale_inv, out_dtype=out_dtype, block_size=block_size)
        y = F.linear(x, Wdq, bias)
        ctx.save_for_backward(weight_fp8, scale_inv)
        ctx.block_size = block_size
        ctx.out_dtype = out_dtype
        del Wdq
        return y

    @staticmethod
    def backward(ctx, grad_out):
        weight_fp8, scale_inv = ctx.saved_tensors
        grad_x = None
        if ctx.needs_input_grad[0]:
            Wdq = dequant_block_fp8(weight_fp8, scale_inv, out_dtype=ctx.out_dtype,
                                    block_size=ctx.block_size)
            grad_x = grad_out.matmul(Wdq)   # y = x @ W^T  ->  dx = grad_out @ W
            del Wdq
        return grad_x, None, None, None, None, None


def _ckpt_linear(input, weight, scale_inv, bias, block_size, out_dtype):
    """grad_input flows; dequant recomputed in backward (frozen base, no weight grad)."""
    if torch.is_grad_enabled() and isinstance(input, torch.Tensor) and input.requires_grad:
        return _FP8LinearFn.apply(input, weight, scale_inv, bias, block_size, out_dtype)
    Wdq = dequant_block_fp8(weight, scale_inv, out_dtype=out_dtype, block_size=block_size)
    return F.linear(input, Wdq, bias)


def _to_local(t):
    try:
        from torch.distributed.tensor import DTensor
        if isinstance(t, DTensor):
            return t.to_local()
    except Exception:
        pass
    return t


def dequant_block_fp8(weight, scale_inv, out_dtype=torch.bfloat16, block_size=None):
    """weight:(out,in) fp8 ; scale_inv:(sr,sc) fp32/ue8m0 -> (out,in) out_dtype.

    Block scales use a FIXED block size (default [128,128] from the checkpoint's
    quant config); the scale grid is ceil(out/bm) x ceil(in/bn), so the final
    block along each dim may be PARTIAL (e.g. out=576 -> 5 blocks, last is 64).
    We expand scales via repeat_interleave at the fixed block size, then slice to
    the weight shape -- this correctly handles partial trailing blocks."""
    W = weight.to(torch.float32)
    out, inp = W.shape[-2], W.shape[-1]
    if scale_inv.dim() == 0 or scale_inv.numel() == 1:
        return (W * scale_inv.to(torch.float32)).to(out_dtype)
    if scale_inv.dtype == torch.uint8:  # ue8m0 packed exponent
        s = (scale_inv.to(torch.float32) - 127.0).exp2()
    else:
        s = scale_inv.to(torch.float32)
    sr, sc = s.shape[-2], s.shape[-1]
    if block_size is not None and len(block_size) == 2:
        bm, bn = int(block_size[0]), int(block_size[1])
    elif out % sr == 0 and inp % sc == 0:
        bm, bn = out // sr, inp // sc
    else:
        # partial trailing block (e.g. 576 -> 5x128); default to 128 block edge
        bm = bn = 128
    s_full = s.repeat_interleave(bm, dim=0).repeat_interleave(bn, dim=1)
    s_full = s_full[:out, :inp]
    return (W * s_full).to(out_dtype)


def _need_grad(x):
    return torch.is_grad_enabled() and isinstance(x, torch.Tensor) and x.requires_grad


def _diff_batched_expert_mm(hidden, weight, scale, expert_ids, num_experts, block_size=None):
    """Differentiable replacement for finegrained_fp8.batched_matmul.
    hidden:(S,in) weight:(E,out,in) scale:(E,sr,sc) -> (S,out). One expert at a time."""
    S = hidden.size(0)
    out_dim = weight.size(1)
    y = hidden.new_zeros(S, out_dim)
    for e in torch.unique(expert_ids).tolist():
        if e >= num_experts:
            continue
        mask = expert_ids == e
        rows = hidden[mask]
        if rows.numel() == 0:
            continue
        y[mask] = _ckpt_linear(rows, weight[e], scale[e], None, block_size, hidden.dtype)
    return y


def install(verbose=True):
    from transformers.integrations import finegrained_fp8 as fp8

    # ---- core: module-level fp8_linear (dense attn LoRA targets + eager experts) ----
    if not getattr(fp8, "_fp8_linear_diff_patched", False):
        _orig_fp8_linear = fp8.fp8_linear

        def diff_fp8_linear(input, weight, weight_scale_inv, block_size=None,
                            activation_scale=None, output_dtype=None, bias=None):
            if not _need_grad(input):
                return _orig_fp8_linear(input, weight, weight_scale_inv,
                                        block_size=block_size,
                                        activation_scale=activation_scale,
                                        output_dtype=output_dtype, bias=bias)
            return _ckpt_linear(input, weight, weight_scale_inv, bias,
                                block_size, (output_dtype or input.dtype))

        fp8.fp8_linear = diff_fp8_linear
        fp8._fp8_linear_diff_patched = True
        if verbose:
            print("[fp8_diff_patch] patched fp8_linear", flush=True)

    # ---- grouped FP8GroupedLinear (shared/fused grouped) ----
    if hasattr(fp8, "FP8GroupedLinear") and not getattr(fp8.FP8GroupedLinear, "_diff_patched", False):
        _orig_grp = fp8.FP8GroupedLinear.forward

        def grp_forward(self, x):
            if self.weight.element_size() > 1 or not _need_grad(x):
                return _orig_grp(self, x)
            input_shape = x.shape[:-2]
            hidden_dim = x.shape[-1]
            w = _to_local(self.weight)
            s = _to_local(self.weight_scale_inv)
            ng = self.n_groups
            Wdq = dequant_block_fp8(w, s, out_dtype=x.dtype, block_size=self.block_size).view(ng, -1, hidden_dim).transpose(1, 2)
            xg = x.reshape(-1, ng, hidden_dim).transpose(0, 1)
            y = torch.bmm(xg, Wdq).transpose(0, 1).reshape(*input_shape, ng, -1)
            if getattr(self, "has_bias", False):
                y = y + self.bias.view(ng, -1)
            return y

        fp8.FP8GroupedLinear.forward = grp_forward
        fp8.FP8GroupedLinear._diff_patched = True
        if verbose:
            print("[fp8_diff_patch] patched FP8GroupedLinear.forward", flush=True)

    # ---- experts batched_mm dispatch ----
    if not getattr(fp8, "_experts_batched_diff_patched", False):
        _orig_experts = fp8.fp8_batched_mm_experts_forward

        def diff_experts_forward(self, hidden_states, top_k_index, top_k_weights):
            if not _need_grad(hidden_states):
                return _orig_experts(self, hidden_states, top_k_index, top_k_weights)
            num_top_k = top_k_index.size(-1)
            num_tokens = hidden_states.size(0)
            hidden_dim = hidden_states.size(-1)
            selected = hidden_states.repeat_interleave(num_top_k, dim=0)
            sample_weights = top_k_weights.reshape(-1)
            expert_ids = top_k_index.reshape(-1)
            sentinel = (expert_ids >= self.num_experts).unsqueeze(-1)
            w_up = _to_local(self.gate_up_proj if self.has_gate else self.up_proj)
            s_up = _to_local(self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv)
            w_dn = _to_local(self.down_proj)
            s_dn = _to_local(self.down_proj_scale_inv)
            proj = _diff_batched_expert_mm(selected, w_up, s_up, expert_ids, self.num_experts, self.block_size)
            proj = self._apply_gate(proj) if self.has_gate else self.act_fn(proj)
            proj = _diff_batched_expert_mm(proj, w_dn, s_dn, expert_ids, self.num_experts, self.block_size)
            weighted = (proj * sample_weights.to(proj.dtype).unsqueeze(-1)).masked_fill(sentinel, 0.0)
            return weighted.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype)

        fp8.fp8_batched_mm_experts_forward = diff_experts_forward
        for tgt in (getattr(fp8.FP8ExpertsInterface, "_global_mapping", None),
                    getattr(fp8, "ALL_FP8_EXPERTS_FUNCTIONS", None)):
            try:
                if tgt is not None and "batched_mm" in tgt:
                    tgt["batched_mm"] = diff_experts_forward
            except Exception:
                pass
        fp8._experts_batched_diff_patched = True
        if verbose:
            print("[fp8_diff_patch] patched fp8_batched_mm_experts_forward", flush=True)

    # ---- experts grouped_mm dispatch (GLM-5.2 default) ----
    if not getattr(fp8, "_experts_grouped_diff_patched", False):
        _orig_grouped = fp8.fp8_grouped_mm_experts_forward

        def diff_grouped_experts_forward(self, hidden_states, top_k_index, top_k_weights):
            if not _need_grad(hidden_states):
                return _orig_grouped(self, hidden_states, top_k_index, top_k_weights)
            num_top_k = top_k_index.size(-1)
            num_tokens = hidden_states.size(0)
            hidden_dim = hidden_states.size(-1)
            sample_weights = top_k_weights.reshape(-1)            # (S,)
            expert_ids = top_k_index.reshape(-1)                  # (S,)
            # token i of pair p is hidden_states[p // num_top_k]
            sel = hidden_states.repeat_interleave(num_top_k, dim=0)  # (S, H)
            sentinel = (expert_ids >= self.num_experts).unsqueeze(-1)
            w_up = _to_local(self.gate_up_proj if self.has_gate else self.up_proj)
            s_up = _to_local(self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv)
            w_dn = _to_local(self.down_proj)
            s_dn = _to_local(self.down_proj_scale_inv)
            proj = _diff_batched_expert_mm(sel, w_up, s_up, expert_ids, self.num_experts, self.block_size)
            proj = self._apply_gate(proj) if self.has_gate else self.act_fn(proj)
            proj = _diff_batched_expert_mm(proj, w_dn, s_dn, expert_ids, self.num_experts, self.block_size)
            weighted = (proj * sample_weights.to(proj.dtype).unsqueeze(-1)).masked_fill(sentinel, 0.0)
            return weighted.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype)

        fp8.fp8_grouped_mm_experts_forward = diff_grouped_experts_forward
        for tgt in (getattr(fp8.FP8ExpertsInterface, "_global_mapping", None),
                    getattr(fp8, "ALL_FP8_EXPERTS_FUNCTIONS", None)):
            try:
                if tgt is not None and "grouped_mm" in tgt:
                    tgt["grouped_mm"] = diff_grouped_experts_forward
            except Exception:
                pass
        fp8._experts_grouped_diff_patched = True
        if verbose:
            print("[fp8_diff_patch] patched fp8_grouped_mm_experts_forward", flush=True)

    return fp8