"""Multi-block Kahan-compensated softmax via Triton with CUDA graphs.

3-phase approach with CUDA graph capture for minimal launch overhead.
Phase 1: local max + exp-sum per chunk
Phase 2: Kahan-compensated reduction
Phase 3: normalize

The CUDA graph is captured once per (batch, vocab) and replayed.
"""
import torch
import torch.nn as nn
import triton
import triton.language as tl

OP_TYPE = "softmax"
SUPPORTED_PRECISIONS = ["fp32"]
HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]

_graph_cache: dict = {}


@triton.jit
def _phase1(
    X_ptr, PartialMax_ptr, PartialSum_ptr,
    batch, vocab,
    stride_x, stride_pm, stride_ps,
    BLOCK: tl.constexpr,
):
    pid = tl.program_id(0)
    nblk = tl.cdiv(vocab, BLOCK)
    row = pid // nblk
    blk = pid % nblk
    if row >= batch:
        return

    cols = blk * BLOCK + tl.arange(0, BLOCK)
    mask = cols < vocab
    x = tl.load(X_ptr + row * stride_x + cols, mask=mask, other=-float("inf"))

    m = tl.max(x, axis=0)
    e = tl.exp(x - m)
    s = tl.sum(e, axis=0)

    tl.store(PartialMax_ptr + row * stride_pm + blk, m)
    tl.store(PartialSum_ptr + row * stride_ps + blk, s)


@triton.jit
def _phase2(
    PartialMax_ptr, PartialSum_ptr,
    GlobalMax_ptr, GlobalInvSum_ptr,
    batch, nblk,
    stride_pm, stride_ps,
    BLOCK_R: tl.constexpr,
):
    row = tl.program_id(0)
    if row >= batch:
        return

    m = -float("inf")
    for off in range(0, nblk, BLOCK_R):
        ids = off + tl.arange(0, BLOCK_R)
        mask = ids < nblk
        pm = tl.load(PartialMax_ptr + row * stride_pm + ids,
                     mask=mask, other=-float("inf"))
        m = tl.maximum(m, tl.max(pm, axis=0))

    s = 0.0
    c = 0.0
    for off in range(0, nblk, BLOCK_R):
        ids = off + tl.arange(0, BLOCK_R)
        mask = ids < nblk
        pm = tl.load(PartialMax_ptr + row * stride_pm + ids,
                     mask=mask, other=-float("inf"))
        ps = tl.load(PartialSum_ptr + row * stride_ps + ids,
                     mask=mask, other=0.0)
        adj = ps * tl.exp(pm - m)
        bs = tl.sum(adj, axis=0)
        y = bs - c
        t = s + y
        c = (t - s) - y
        s = t

    tl.store(GlobalMax_ptr + row, m)
    tl.store(GlobalInvSum_ptr + row, 1.0 / s)


@triton.jit
def _phase3(
    X_ptr, Y_ptr,
    GlobalMax_ptr, GlobalInvSum_ptr,
    batch, vocab,
    stride_x, stride_y,
    BLOCK: tl.constexpr,
):
    pid = tl.program_id(0)
    nblk = tl.cdiv(vocab, BLOCK)
    row = pid // nblk
    blk = pid % nblk
    if row >= batch:
        return

    gm = tl.load(GlobalMax_ptr + row)
    gi = tl.load(GlobalInvSum_ptr + row)

    cols = blk * BLOCK + tl.arange(0, BLOCK)
    mask = cols < vocab
    x = tl.load(X_ptr + row * stride_x + cols, mask=mask, other=-float("inf"))
    out = tl.exp(x - gm) * gi
    tl.store(Y_ptr + row * stride_y + cols, out, mask=mask)


class Model(nn.Module):
    def __init__(self, batch: int, vocab: int):
        super().__init__()
        self.batch = batch
        self.vocab = vocab

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, V = x.shape
        device = x.device
        BLOCK = 1024
        nblk = (V + BLOCK - 1) // BLOCK
        total = B * nblk

        key = (device, B, nblk)
        graph_info = _graph_cache.get(key)

        if graph_info is None:
            # First call: allocate buffers and capture graph
            pm = torch.empty(B, nblk, device=device, dtype=torch.float32)
            ps = torch.empty(B, nblk, device=device, dtype=torch.float32)
            gm = torch.empty(B, device=device, dtype=torch.float32)
            gi = torch.empty(B, device=device, dtype=torch.float32)
            y = torch.empty_like(x)

            # Warmup individual kernels
            s = torch.cuda.Stream()
            s.wait_stream(torch.cuda.current_stream())
            with torch.cuda.stream(s):
                for _ in range(3):
                    _phase1[(total,)](
                        x, pm, ps,
                        B, V,
                        x.stride(0), nblk, nblk,
                        BLOCK=BLOCK, num_warps=8,
                    )
                    _phase2[(B,)](
                        pm, ps, gm, gi,
                        B, nblk,
                        nblk, nblk,
                        BLOCK_R=1024, num_warps=4,
                    )
                    _phase3[(total,)](
                        x, y, gm, gi,
                        B, V,
                        x.stride(0), y.stride(0),
                        BLOCK=BLOCK, num_warps=8,
                    )
            torch.cuda.current_stream().wait_stream(s)

            # Capture graph
            g = torch.cuda.CUDAGraph()
            with torch.cuda.graph(g):
                _phase1[(total,)](
                    x, pm, ps,
                    B, V,
                    x.stride(0), nblk, nblk,
                    BLOCK=BLOCK, num_warps=8,
                )
                _phase2[(B,)](
                    pm, ps, gm, gi,
                    B, nblk,
                    nblk, nblk,
                    BLOCK_R=1024, num_warps=4,
                )
                _phase3[(total,)](
                    x, y, gm, gi,
                    B, V,
                    x.stride(0), y.stride(0),
                    BLOCK=BLOCK, num_warps=8,
                )

            _graph_cache[key] = (g, x, y, pm, ps, gm, gi)
            return y

        g, x_cap, y, pm, ps, gm, gi = graph_info
        # Copy input to graph's input buffer
        x_cap.copy_(x)
        g.replay()
        return y


BATCH = 8
VOCAB = 32768


def get_inputs():
    x = torch.randn(BATCH, VOCAB, dtype=torch.float32) * 4.0
    return [x]


def get_init_inputs():
    return [BATCH, VOCAB]