"""Fused cross-entropy: same maths as FusedLinearCE, but the per-chunk softmax/grad chain is one compiled kernel set.

On a 288 GB/s card the eager version streams ~5 GB of fp32 logits-sized tensors per chunk; here the bf16 logits are
read twice and the bf16 gradient written once.
"""
import torch


def _chunk(logits, tgt, z_loss: float, inv_n: float):
    lg = logits.float()
    lse = torch.logsumexp(lg, dim=-1)
    tl = lg.gather(1, tgt[:, None]).squeeze(1)
    loss = (lse - tl).sum()
    p = torch.exp(lg - lse[:, None])
    if z_loss:
        loss = loss + z_loss * (lse * lse).sum()
        p = p * (1.0 + 2.0 * z_loss * lse[:, None])
    cols = torch.arange(lg.shape[1], device=lg.device)[None, :]
    p = torch.where(cols == tgt[:, None], p - 1.0, p)
    return loss, (p * inv_n).to(torch.bfloat16)


chunk_compiled = torch.compile(_chunk, dynamic=False)


class FastLinearCE(torch.autograd.Function):
    @staticmethod
    def forward(ctx, h, w, targets, chunk: int, z_loss: float):
        N = h.shape[0]
        wb = w.to(h.dtype)
        grad_h = torch.empty_like(h)
        grad_w = torch.zeros_like(w, dtype=torch.float32)
        total = torch.zeros((), device=h.device, dtype=torch.float32)
        for i in range(0, N, chunk):
            hc, tc = h[i : i + chunk], targets[i : i + chunk]
            loss_i, pb = chunk_compiled(hc @ wb.t(), tc, z_loss, 1.0 / N)
            total += loss_i
            grad_h[i : i + chunk] = pb @ wb
            grad_w.add_((pb.t() @ hc).float())
        ctx.save_for_backward(grad_h, grad_w)
        return total / N

    @staticmethod
    def backward(ctx, g):
        grad_h, grad_w = ctx.saved_tensors
        return grad_h * g.to(grad_h.dtype), grad_w * g, None, None, None


def _chunk_multi(logits, tgt, z_loss: float, inv_n: float):
    """Patch-level loss: one distribution per position scored against all k tokens of the next patch.
    tgt [n, k]; loss = sum_k CE(logits, tgt[:, k]); grad = k*softmax - sum_k onehot."""
    lg = logits.float()
    k = tgt.shape[1]
    lse = torch.logsumexp(lg, dim=-1)
    loss = (k * lse - lg.gather(1, tgt).sum(1)).sum()
    p = k * torch.exp(lg - lse[:, None])
    if z_loss:
        loss = loss + k * z_loss * (lse * lse).sum()
        p = p * (1.0 + 2.0 * z_loss * lse[:, None])
    cols = torch.arange(lg.shape[1], device=lg.device)[None, :]
    for j in range(k):   # unrolled at trace time and fused: no logits-sized one-hot buffer
        p = p - (cols == tgt[:, j : j + 1]).to(p.dtype)
    return loss, (p * inv_n).to(torch.bfloat16)


chunk_multi = torch.compile(_chunk_multi, dynamic=False)


class PatchLinearCE(torch.autograd.Function):
    """FastLinearCE for patch-level training: h [N, d], targets [N, k]; returns the mean per-token CE."""
    @staticmethod
    def forward(ctx, h, w, targets, chunk: int, z_loss: float):
        N, k = targets.shape
        wb = w.to(h.dtype)
        grad_h = torch.empty_like(h)
        grad_w = torch.zeros_like(w, dtype=torch.float32)
        total = torch.zeros((), device=h.device, dtype=torch.float32)
        for i in range(0, N, chunk):
            hc, tc = h[i : i + chunk], targets[i : i + chunk]
            loss_i, pb = chunk_multi(hc @ wb.t(), tc, z_loss, 1.0 / (N * k))
            total += loss_i
            grad_h[i : i + chunk] = pb @ wb
            grad_w.add_((pb.t() @ hc).float())
        ctx.save_for_backward(grad_h, grad_w)
        return total / (N * k)

    @staticmethod
    def backward(ctx, g):
        grad_h, grad_w = ctx.saved_tensors
        return grad_h * g.to(grad_h.dtype), grad_w * g, None, None, None


# Variable assistant counts and CE tails must not specialize each compiled call.
chunk_sft = torch.compile(_chunk, dynamic=True)


class SFTLinearCE(torch.autograd.Function):
    @staticmethod
    def forward(ctx, h, w, targets, chunk: int, z_loss: float):
        N = h.shape[0]
        wb = w.to(h.dtype)
        grad_h = torch.empty_like(h)
        grad_w = torch.zeros_like(w, dtype=torch.float32)
        total = torch.zeros((), device=h.device, dtype=torch.float32)
        for i in range(0, N, chunk):
            hc, tc = h[i : i + chunk], targets[i : i + chunk]
            loss_i, pb = chunk_sft(hc @ wb.t(), tc, z_loss, torch.tensor(1.0 / N, device=h.device))
            total += loss_i
            grad_h[i : i + chunk] = pb @ wb
            grad_w.add_((pb.t() @ hc).float())
        ctx.save_for_backward(grad_h, grad_w)
        return total / N

    @staticmethod
    def backward(ctx, g):
        grad_h, grad_w = ctx.saved_tensors
        return grad_h * g.to(grad_h.dtype), grad_w * g, None, None, None