"""
torch.compile inductor bug: closure-based mask + SDPA fails on recompilation (CPU only).

Same mask logic written inline works fine. Only the closure composition
pattern triggers a C++ codegen error (undeclared variable) on recompilation.
"""
import torch
import torch.nn.functional as F


# ---- Closure-based mask construction (triggers bug) ----

def causal_fn(b, h, q, kv):
    return kv <= q

def padding_fn(padding_mask):
    def inner(b, h, q, kv):
        return padding_mask[b, kv]
    return inner

def and_masks(f1, f2):
    def combined(b, h, q, kv):
        return f1(b, h, q, kv) & f2(b, h, q, kv)
    return combined

def make_mask_closure(batch_size, q_len, kv_len, q_offset, padding_mask):
    fn = and_masks(causal_fn, padding_fn(padding_mask))
    b = torch.arange(batch_size)[:, None, None, None]
    h = torch.arange(1)[None, :, None, None]
    q = (torch.arange(q_len) + q_offset)[None, None, :, None]
    kv = torch.arange(kv_len)[None, None, None, :]
    return fn(b, h, q, kv).expand(batch_size, 1, q_len, kv_len)


# ---- Inline mask construction (no closures, works fine) ----

def make_mask_inline(batch_size, q_len, kv_len, q_offset, padding_mask):
    b = torch.arange(batch_size)[:, None, None, None]
    kv = torch.arange(kv_len)[None, None, None, :]
    q = (torch.arange(q_len) + q_offset)[None, None, :, None]
    return (padding_mask[b, kv] & (kv <= q)).expand(batch_size, 1, q_len, kv_len)


# ---- Minimal model: mask + SDPA ----

class Model(torch.nn.Module):
    def __init__(self, use_closure=True):
        super().__init__()
        self.proj = torch.nn.Linear(32, 32, bias=False)
        self.use_closure = use_closure

    def forward(self, x, past_k, past_v, padding_mask, q_offset):
        B, S, _ = x.shape
        q = self.proj(x).view(B, S, 4, 8).transpose(1, 2)
        k = x.view(B, S, 4, 8).transpose(1, 2)
        v = k.clone()
        if past_k is not None:
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)

        if self.use_closure:
            mask = make_mask_closure(B, S, k.shape[2], q_offset, padding_mask)
        else:
            mask = make_mask_inline(B, S, k.shape[2], q_offset, padding_mask)

        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        return out.transpose(1, 2).reshape(B, S, 32), k, v


# ---- Test: prefill then decode (shape changes → recompilation) ----

def test(label, use_closure):
    print(f"\n--- {label} ---")
    torch._dynamo.reset()
    model = Model(use_closure=use_closure).eval()
    compiled = torch.compile(model)

    pad_mask = torch.ones(1, 8, dtype=torch.bool)
    pad_mask[0, :2] = False  # 2 padding tokens

    # Prefill: seq_len=8
    x = torch.randn(1, 8, 32)
    with torch.no_grad():
        out, pk, pv = compiled(x, None, None, pad_mask, q_offset=0)
    print(f"  Prefill (kv=8): OK")

    # Decode steps: seq_len=1, kv grows each step → recompilation
    for step in range(3):
        pad_mask = torch.cat([pad_mask, torch.ones(1, 1, dtype=torch.bool)], dim=1)
        x = torch.randn(1, 1, 32)
        try:
            with torch.no_grad():
                out, pk, pv = compiled(x, pk, pv, pad_mask, q_offset=pk.shape[2])
            print(f"  Decode step {step+1} (kv={pk.shape[2]}): OK")
        except Exception as e:
            import traceback
            print(f"  Decode step {step+1}: FAILED — {type(e).__name__}")
            traceback.print_exc()
            return


if __name__ == "__main__":
    print(f"torch {torch.__version__}")
    test("Closure-based mask (BUG)", use_closure=True)
    test("Inline mask (OK)", use_closure=False)

---

--- Closure-based mask (BUG) ---
  Prefill (kv=8): OK
  Decode step 1: FAILED — InductorError
  ...
  CppCompileError: C++ compile error
  ...
  error: 'tmp2' was not declared in this scope
  ...

--- Inline mask (OK) ---
  Prefill (kv=8): OK
  Decode step 1 (kv=9): OK
  Decode step 2 (kv=10): OK
  Decode step 3 (kv=11): OK