"""Fused CUDA kernels for memory-bound blocks: SiTU-GLU, the LM loss, token log-probs."""

from __future__ import annotations

import torch
import torch.nn.functional as F
import triton
import triton.language as tl
from triton.language.extra import libdevice

from tiny_gdn.nn_common import (
    linear_cross_entropy_reference,
    linear_token_logprobs_reference,
    situ_glu_reference,
)

BLOCK_SIZE = 1024
VOCAB_BLOCK_SIZE = 4096
LOGITS_CHUNK_BYTES = 1 << 30


@triton.jit
def _situ_glu_forward_kernel(
    gate_up,
    output,
    total,
    width,
    gate_cap,
    up_cap,
    BLOCK: tl.constexpr,
):
    offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
    mask = offsets < total
    row = offsets // width
    gate_offsets = offsets + row * width
    gate = tl.load(gate_up + gate_offsets, mask=mask, other=0.0).to(tl.float32)
    up = tl.load(gate_up + gate_offsets + width, mask=mask, other=0.0).to(tl.float32)
    gated = gate_cap * libdevice.tanh(gate / gate_cap) * tl.sigmoid(gate)
    result = gated * up_cap * libdevice.tanh(up / up_cap)
    tl.store(output + offsets, result.to(output.dtype.element_ty), mask=mask)


@triton.jit
def _situ_glu_backward_kernel(
    gate_up,
    grad_output,
    grad_gate_up,
    total,
    width,
    gate_cap,
    up_cap,
    BLOCK: tl.constexpr,
):
    offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
    mask = offsets < total
    row = offsets // width
    gate_offsets = offsets + row * width
    gate = tl.load(gate_up + gate_offsets, mask=mask, other=0.0).to(tl.float32)
    up = tl.load(gate_up + gate_offsets + width, mask=mask, other=0.0).to(tl.float32)
    grad = tl.load(grad_output + offsets, mask=mask, other=0.0).to(tl.float32)
    gate_tanh = libdevice.tanh(gate / gate_cap)
    up_tanh = libdevice.tanh(up / up_cap)
    sigmoid = tl.sigmoid(gate)
    gated = gate_cap * gate_tanh * sigmoid
    capped_up = up_cap * up_tanh
    gate_slope = sigmoid * ((1.0 - gate_tanh * gate_tanh) + gate_cap * gate_tanh * (1.0 - sigmoid))
    grad_gate = grad * capped_up * gate_slope
    grad_up = grad * gated * (1.0 - up_tanh * up_tanh)
    dtype = grad_gate_up.dtype.element_ty
    tl.store(grad_gate_up + gate_offsets, grad_gate.to(dtype), mask=mask)
    tl.store(grad_gate_up + gate_offsets + width, grad_up.to(dtype), mask=mask)


def _launch_grid(total: int) -> tuple[int]:
    return (triton.cdiv(total, BLOCK_SIZE),)


class _SiTUGLUFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, gate_up: torch.Tensor, gate_cap: float, up_cap: float) -> torch.Tensor:
        gate_up = gate_up.contiguous()
        width = gate_up.shape[-1] // 2
        output = gate_up.new_empty((*gate_up.shape[:-1], width))
        total = output.numel()
        if total:
            _situ_glu_forward_kernel[_launch_grid(total)](
                gate_up, output, total, width, gate_cap, up_cap, BLOCK=BLOCK_SIZE
            )
        ctx.save_for_backward(gate_up)
        ctx.caps = (gate_cap, up_cap)
        return output

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]:
        (gate_up,) = ctx.saved_tensors
        gate_cap, up_cap = ctx.caps
        grad_output = grad_output.contiguous()
        grad_gate_up = torch.empty_like(gate_up)
        width = gate_up.shape[-1] // 2
        total = grad_output.numel()
        if total:
            _situ_glu_backward_kernel[_launch_grid(total)](
                gate_up,
                grad_output,
                grad_gate_up,
                total,
                width,
                gate_cap,
                up_cap,
                BLOCK=BLOCK_SIZE,
            )
        return grad_gate_up, None, None


def situ_glu(gate_up: torch.Tensor, gate_cap: float, up_cap: float) -> torch.Tensor:
    """``softcap(gate) * sigmoid(gate) * softcap(up)`` over ``[..., gate | up]``.

    The CUDA path evaluates in FP32 and rounds once; backward recomputes the
    activations from ``gate_up`` instead of saving intermediates.
    """

    if gate_up.shape[-1] % 2:
        raise ValueError("gate_up must have an even last dimension")
    if not gate_up.is_cuda:
        return situ_glu_reference(gate_up, gate_cap, up_cap)
    return _SiTUGLUFunction.apply(gate_up, gate_cap, up_cap)


@triton.jit
def _cross_entropy_rows_kernel(
    logits,
    targets,
    losses,
    log_partitions,
    valid_count,
    row_stride,
    vocab_size,
    z_loss_coefficient,
    ignore_index,
    COMPUTE_GRAD: tl.constexpr,
    BLOCK: tl.constexpr,
):
    """Per row: logsumexp and NLL; optionally overwrite the row with d(mean loss)/d(logits)."""

    row = tl.program_id(0).to(tl.int64)
    row_logits = logits + row * row_stride
    target = tl.load(targets + row)
    columns = tl.arange(0, BLOCK)
    if target == ignore_index:
        tl.store(losses + row, 0.0)
        tl.store(log_partitions + row, 0.0)
        if COMPUTE_GRAD:
            for start in range(0, vocab_size, BLOCK):
                offsets = start + columns
                zeros = tl.zeros((BLOCK,), dtype=logits.dtype.element_ty)
                tl.store(row_logits + offsets, zeros, mask=offsets < vocab_size)
    else:
        running_max = tl.full((), float("-inf"), tl.float32)
        running_sum = tl.zeros((), dtype=tl.float32)
        for start in range(0, vocab_size, BLOCK):
            offsets = start + columns
            values = tl.load(
                row_logits + offsets, mask=offsets < vocab_size, other=float("-inf")
            ).to(tl.float32)
            new_max = tl.maximum(running_max, tl.max(values, axis=0))
            running_sum = running_sum * tl.exp(running_max - new_max) + tl.sum(
                tl.exp(values - new_max), axis=0
            )
            running_max = new_max
        log_partition = running_max + tl.log(running_sum)
        target_logit = tl.load(row_logits + target).to(tl.float32)
        tl.store(losses + row, log_partition - target_logit)
        tl.store(log_partitions + row, log_partition)
        if COMPUTE_GRAD:
            grad_scale = 1.0 / tl.load(valid_count).to(tl.float32)
            softmax_scale = grad_scale * (1.0 + 2.0 * z_loss_coefficient * log_partition)
            for start in range(0, vocab_size, BLOCK):
                offsets = start + columns
                in_vocab = offsets < vocab_size
                values = tl.load(row_logits + offsets, mask=in_vocab, other=0.0).to(tl.float32)
                grad = tl.exp(values - log_partition) * softmax_scale
                grad = tl.where(offsets == target, grad - grad_scale, grad)
                tl.store(row_logits + offsets, grad.to(logits.dtype.element_ty), mask=in_vocab)


class _LinearCrossEntropyFunction(torch.autograd.Function):
    """Chunked ``states @ weight.T`` + cross-entropy + z-loss with gradients precomputed.

    Each logits chunk is consumed in place, so peak extra memory is one chunk
    instead of ``[tokens, vocab]`` logits plus their softmax for backward.
    """

    @staticmethod
    def forward(
        ctx,
        states: torch.Tensor,
        weight: torch.Tensor,
        targets: torch.Tensor,
        z_loss_coefficient: float,
        ignore_index: int,
        compute_grad: bool,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        states = states.contiguous()
        targets = targets.contiguous()
        rows, vocab_size = states.shape[0], weight.shape[0]
        valid = targets.ne(ignore_index)
        valid_count = valid.sum(dtype=torch.float32)
        losses = states.new_empty(rows, dtype=torch.float32)
        log_partitions = states.new_empty(rows, dtype=torch.float32)
        grad_states = torch.empty_like(states) if compute_grad else None
        grad_weight = (
            torch.zeros(weight.shape, dtype=torch.float32, device=weight.device)
            if compute_grad
            else None
        )
        chunk_rows = max(1, LOGITS_CHUNK_BYTES // (vocab_size * states.element_size()))
        for start in range(0, rows, chunk_rows):
            end = min(start + chunk_rows, rows)
            chunk = states[start:end]
            logits = chunk @ weight.t()
            _cross_entropy_rows_kernel[(end - start,)](
                logits,
                targets[start:end],
                losses[start:end],
                log_partitions[start:end],
                valid_count,
                logits.stride(0),
                vocab_size,
                z_loss_coefficient,
                ignore_index,
                COMPUTE_GRAD=compute_grad,
                BLOCK=VOCAB_BLOCK_SIZE,
                num_warps=8,
            )
            if compute_grad:
                torch.mm(logits, weight, out=grad_states[start:end])
                if chunk.dtype == torch.float32:
                    grad_weight.addmm_(logits.t(), chunk)
                else:
                    torch.addmm(
                        grad_weight, logits.t(), chunk, out_dtype=torch.float32, out=grad_weight
                    )
        cross_entropy = losses.sum() / valid_count
        z_loss = log_partitions.square().sum() / valid_count
        objective = cross_entropy + z_loss_coefficient * z_loss
        if compute_grad:
            ctx.save_for_backward(grad_states, grad_weight)
        ctx.weight_dtype = weight.dtype
        ctx.mark_non_differentiable(cross_entropy, z_loss)
        return objective, cross_entropy, z_loss

    @staticmethod
    def backward(
        ctx, grad_objective: torch.Tensor, *_: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, None, None, None, None]:
        grad_states, grad_weight = ctx.saved_tensors
        scale = grad_objective.to(torch.float32)
        grad_states = (grad_states * scale).to(grad_states.dtype)
        grad_weight = (grad_weight * scale).to(ctx.weight_dtype)
        return grad_states, grad_weight, None, None, None, None


class _WeightedLinearCrossEntropyFunction(torch.autograd.Function):
    """``Σ_i w_i (CE_i + z log Z_i²)`` over chunked ``states @ weight.T``, gradients precomputed.

    The row kernel runs with unit scale and each row's gradient is then
    multiplied by its weight, so per-token weights cost one in-place multiply.
    """

    @staticmethod
    def forward(
        ctx,
        states: torch.Tensor,
        weight: torch.Tensor,
        targets: torch.Tensor,
        token_weights: torch.Tensor,
        z_loss_coefficient: float,
        compute_grad: bool,
        excluded_token_id: int | None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        states = states.contiguous()
        targets = targets.contiguous()
        token_weights = token_weights.float().contiguous()
        rows, vocab_size = states.shape[0], weight.shape[0]
        unit = torch.ones((), dtype=torch.float32, device=states.device)
        losses = states.new_empty(rows, dtype=torch.float32)
        log_partitions = states.new_empty(rows, dtype=torch.float32)
        grad_states = torch.empty_like(states) if compute_grad else None
        grad_weight = (
            torch.zeros(weight.shape, dtype=torch.float32, device=weight.device)
            if compute_grad
            else None
        )
        chunk_rows = max(1, LOGITS_CHUNK_BYTES // (vocab_size * states.element_size()))
        for start in range(0, rows, chunk_rows):
            end = min(start + chunk_rows, rows)
            chunk = states[start:end]
            logits = chunk @ weight.t()
            if excluded_token_id is not None:
                logits[:, excluded_token_id] = float("-inf")
            _cross_entropy_rows_kernel[(end - start,)](
                logits,
                targets[start:end],
                losses[start:end],
                log_partitions[start:end],
                unit,
                logits.stride(0),
                vocab_size,
                z_loss_coefficient,
                -100,
                COMPUTE_GRAD=compute_grad,
                BLOCK=VOCAB_BLOCK_SIZE,
                num_warps=8,
            )
            if compute_grad:
                logits.mul_(token_weights[start:end, None].to(logits.dtype))
                torch.mm(logits, weight, out=grad_states[start:end])
                if chunk.dtype == torch.float32:
                    grad_weight.addmm_(logits.t(), chunk)
                else:
                    torch.addmm(
                        grad_weight, logits.t(), chunk, out_dtype=torch.float32, out=grad_weight
                    )
        objective = (token_weights * (losses + z_loss_coefficient * log_partitions.square())).sum()
        if compute_grad:
            ctx.save_for_backward(grad_states, grad_weight)
        ctx.weight_dtype = weight.dtype
        ctx.mark_non_differentiable(losses, log_partitions)
        return objective, losses, log_partitions

    @staticmethod
    def backward(
        ctx, grad_objective: torch.Tensor, *_: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, None, None, None, None, None]:
        grad_states, grad_weight = ctx.saved_tensors
        scale = grad_objective.to(torch.float32)
        grad_states = (grad_states * scale).to(grad_states.dtype)
        grad_weight = (grad_weight * scale).to(ctx.weight_dtype)
        return grad_states, grad_weight, None, None, None, None, None


def weighted_linear_cross_entropy(
    states: torch.Tensor,
    weight: torch.Tensor,
    targets: torch.Tensor,
    token_weights: torch.Tensor,
    *,
    z_loss_coefficient: float = 0.0,
    excluded_token_id: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """``(Σ_i w_i (CE_i + z log Z_i²), CE [rows], log Z [rows])`` without full logits.

    Every row is scored; only the weighted sum carries gradients.
    ``excluded_token_id`` gets probability zero (its logit is -inf).
    """

    if states.ndim != 2 or targets.shape != states.shape[:1] or token_weights.shape != targets.shape:
        raise ValueError("states must be [tokens, hidden], targets and token_weights [tokens]")
    if not states.is_cuda:
        logits = F.linear(states, weight).float()
        if excluded_token_id is not None:
            logits[:, excluded_token_id] = float("-inf")
        log_partitions = torch.logsumexp(logits, dim=-1)
        losses = log_partitions - logits.gather(1, targets[:, None]).squeeze(1)
        weights = token_weights.float()
        objective = (weights * (losses + z_loss_coefficient * log_partitions.square())).sum()
        return objective, losses.detach(), log_partitions.detach()
    compute_grad = torch.is_grad_enabled() and (states.requires_grad or weight.requires_grad)
    return _WeightedLinearCrossEntropyFunction.apply(
        states, weight, targets, token_weights, z_loss_coefficient, compute_grad, excluded_token_id
    )


def linear_cross_entropy(
    states: torch.Tensor,
    weight: torch.Tensor,
    targets: torch.Tensor,
    *,
    z_loss_coefficient: float = 0.0,
    ignore_index: int = -100,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Mean ``cross_entropy(states @ weight.T, targets)`` without full logits.

    Returns ``(objective, cross_entropy, z_loss)``. Only ``objective =
    cross_entropy + z_loss_coefficient * z_loss`` carries gradients; ``z_loss``
    is the mean squared log-partition over non-ignored targets.
    """

    if states.ndim != 2 or targets.shape != states.shape[:1]:
        raise ValueError("states must be [tokens, hidden] and targets [tokens]")
    if not states.is_cuda:
        return linear_cross_entropy_reference(
            states,
            weight,
            targets,
            z_loss_coefficient=z_loss_coefficient,
            ignore_index=ignore_index,
        )
    compute_grad = torch.is_grad_enabled() and (states.requires_grad or weight.requires_grad)
    return _LinearCrossEntropyFunction.apply(
        states, weight, targets, z_loss_coefficient, ignore_index, compute_grad
    )


@triton.jit
def _token_logprob_grad_rows_kernel(
    logits,
    targets,
    log_partitions,
    row_grads,
    row_stride,
    vocab_size,
    ignore_index,
    BLOCK: tl.constexpr,
):
    """Overwrite each row with ``g * (onehot(target) - softmax)``, the gradient of ``g * log p``."""

    row = tl.program_id(0).to(tl.int64)
    row_logits = logits + row * row_stride
    target = tl.load(targets + row)
    columns = tl.arange(0, BLOCK)
    if target == ignore_index:
        for start in range(0, vocab_size, BLOCK):
            offsets = start + columns
            zeros = tl.zeros((BLOCK,), dtype=logits.dtype.element_ty)
            tl.store(row_logits + offsets, zeros, mask=offsets < vocab_size)
    else:
        scale = tl.load(row_grads + row).to(tl.float32)
        log_partition = tl.load(log_partitions + row)
        for start in range(0, vocab_size, BLOCK):
            offsets = start + columns
            in_vocab = offsets < vocab_size
            values = tl.load(row_logits + offsets, mask=in_vocab, other=0.0).to(tl.float32)
            grad = -scale * tl.exp(values - log_partition)
            grad = tl.where(offsets == target, grad + scale, grad)
            tl.store(row_logits + offsets, grad.to(logits.dtype.element_ty), mask=in_vocab)


class _LinearTokenLogprobFunction(torch.autograd.Function):
    """Per-row ``log_softmax(states @ weight.T)[target]`` without keeping logits.

    Preference losses weight every sequence differently, so the gradient
    cannot be precomputed in forward like the mean LM loss. Forward keeps one
    log-partition per row; backward recomputes each logits chunk and turns it
    into the scaled softmax gradient in place.
    """

    @staticmethod
    def forward(
        ctx,
        states: torch.Tensor,
        weight: torch.Tensor,
        targets: torch.Tensor,
        ignore_index: int,
    ) -> torch.Tensor:
        states = states.contiguous()
        targets = targets.contiguous()
        rows, vocab_size = states.shape[0], weight.shape[0]
        negative_logprobs = states.new_empty(rows, dtype=torch.float32)
        log_partitions = states.new_empty(rows, dtype=torch.float32)
        unused_count = states.new_ones((), dtype=torch.float32)
        chunk_rows = max(1, LOGITS_CHUNK_BYTES // (vocab_size * states.element_size()))
        for start in range(0, rows, chunk_rows):
            end = min(start + chunk_rows, rows)
            logits = states[start:end] @ weight.t()
            _cross_entropy_rows_kernel[(end - start,)](
                logits,
                targets[start:end],
                negative_logprobs[start:end],
                log_partitions[start:end],
                unused_count,
                logits.stride(0),
                vocab_size,
                0.0,
                ignore_index,
                COMPUTE_GRAD=False,
                BLOCK=VOCAB_BLOCK_SIZE,
                num_warps=8,
            )
        ctx.save_for_backward(states, weight, targets, log_partitions)
        ctx.ignore_index = ignore_index
        ctx.chunk_rows = chunk_rows
        return negative_logprobs.neg_()

    @staticmethod
    def backward(
        ctx, grad_logprobs: torch.Tensor
    ) -> tuple[torch.Tensor | None, torch.Tensor | None, None, None]:
        states, weight, targets, log_partitions = ctx.saved_tensors
        row_grads = grad_logprobs.to(torch.float32).contiguous()
        rows, vocab_size = states.shape[0], weight.shape[0]
        grad_states = torch.empty_like(states) if ctx.needs_input_grad[0] else None
        grad_weight = (
            torch.zeros(weight.shape, dtype=torch.float32, device=weight.device)
            if ctx.needs_input_grad[1]
            else None
        )
        for start in range(0, rows, ctx.chunk_rows):
            end = min(start + ctx.chunk_rows, rows)
            chunk = states[start:end]
            logits = chunk @ weight.t()
            _token_logprob_grad_rows_kernel[(end - start,)](
                logits,
                targets[start:end],
                log_partitions[start:end],
                row_grads[start:end],
                logits.stride(0),
                vocab_size,
                ctx.ignore_index,
                BLOCK=VOCAB_BLOCK_SIZE,
                num_warps=8,
            )
            if grad_states is not None:
                torch.mm(logits, weight, out=grad_states[start:end])
            if grad_weight is not None:
                if chunk.dtype == torch.float32:
                    grad_weight.addmm_(logits.t(), chunk)
                else:
                    torch.addmm(
                        grad_weight, logits.t(), chunk, out_dtype=torch.float32, out=grad_weight
                    )
        return (
            grad_states,
            None if grad_weight is None else grad_weight.to(weight.dtype),
            None,
            None,
        )


def linear_token_logprobs(
    states: torch.Tensor,
    weight: torch.Tensor,
    targets: torch.Tensor,
    *,
    ignore_index: int = -100,
) -> torch.Tensor:
    """fp32 ``log_softmax(states @ weight.T)[row, target]`` per row without full logits.

    Rows whose target is ``ignore_index`` return 0 and receive no gradient.
    """

    if states.ndim != 2 or targets.shape != states.shape[:1]:
        raise ValueError("states must be [tokens, hidden] and targets [tokens]")
    if not states.is_cuda:
        return linear_token_logprobs_reference(
            states, weight, targets, ignore_index=ignore_index
        )
    return _LinearTokenLogprobFunction.apply(states, weight, targets, ignore_index)