"""CubicV11 long-context architecture and 32k training-memory benchmark.

This is the optimization gate before distillation.  It replaces V7's dense
N x N cosine top-k matrix with fused block-sparse local attention, adds causal
global block summaries, RoPE, GQA, detached depth memory, SwiGLU,
tied embeddings and per-block activation checkpointing.

Run:
    python cubic_v11_long_context_32k.py

The script executes a real forward + backward + Muon step at 32,768 tokens and
falls back to 16,384 only if the full training step cannot fit.
"""

from __future__ import annotations

import gc
import importlib.util
import math
import sys
import time
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

try:
    from torch.nn.attention.flex_attention import create_block_mask, flex_attention
except ImportError as exc:
    raise RuntimeError("CubicV11 long context requires PyTorch FlexAttention") from exc


ROOT = Path(__file__).parent
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
VOCAB_SIZE = 8_192
DIM = 896
DEPTH = 14
QUERY_HEADS = 14
KV_HEADS = 7
HEAD_DIM = DIM // QUERY_HEADS
DEPTH_RANK = 224
DEPTH_HEADS = 7
LOCAL_WINDOW = 4_096
SUMMARY_BLOCK = 256
ROPE_BASE = 500_000.0
TARGET_LENGTHS = (32_768, 16_384)
SEED = 20260720


def load_base():
    source = ROOT / "cubic_v5_muon_benchmark.py"
    spec = importlib.util.spec_from_file_location("v11_muon_base", source)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


# Compile only the sparse attention operator. Compiling the entire 14-layer
# checkpointed graph is slower and substantially more fragile at 32k.
compiled_flex_attention = torch.compile(flex_attention, mode="default", dynamic=False)


class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim: int, base: float):
        super().__init__()
        inv = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
        self.register_buffer("inv_freq", inv, persistent=False)

    def forward(self, seq_len: int, device, dtype):
        positions = torch.arange(seq_len, device=device, dtype=torch.float32)
        angles = torch.outer(positions, self.inv_freq.to(device))
        cos = torch.repeat_interleave(angles.cos(), 2, dim=-1).to(dtype)
        sin = torch.repeat_interleave(angles.sin(), 2, dim=-1).to(dtype)
        return cos.view(1, 1, seq_len, -1), sin.view(1, 1, seq_len, -1)


def rotate_half(x):
    even = x[..., 0::2]
    odd = x[..., 1::2]
    return torch.stack((-odd, even), dim=-1).flatten(-2)


def apply_rope(x, cos, sin):
    return x * cos + rotate_half(x) * sin


class SwiGLU(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        hidden = 2_432  # close to 8/3 * DIM and divisible by 128
        self.gate_up = nn.Linear(dim, 2 * hidden, bias=False)
        self.down = nn.Linear(hidden, dim, bias=False)

    def forward(self, x):
        gate, value = self.gate_up(x).chunk(2, dim=-1)
        return self.down(F.silu(gate) * value)


class CompressedDepthMemory(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm = nn.RMSNorm(DIM)
        self.kv = nn.Linear(DIM, 2 * DEPTH_RANK, bias=False)

    def forward(self, x):
        # The residual sequence path still carries full gradients. Detaching
        # only the historical side-path avoids retaining every prior block's
        # full activation graph at 32k.
        batch, seq, _ = x.shape
        kv = self.kv(self.norm(x.detach())).view(
            batch, seq, 2, DEPTH_HEADS, DEPTH_RANK // DEPTH_HEADS
        )
        return kv[:, :, 0], kv[:, :, 1]


class LongContextAttention(nn.Module):
    def __init__(self, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        self.q_proj = nn.Linear(DIM, QUERY_HEADS * HEAD_DIM, bias=False)
        self.kv_proj = nn.Linear(DIM, 2 * KV_HEADS * HEAD_DIM, bias=False)
        self.out_proj = nn.Linear(DIM, DIM, bias=False)

        # True pooled global summaries. This branch is linear in sequence
        # length because there are only N / SUMMARY_BLOCK summary vectors.
        self.global_q = nn.Linear(DIM, DEPTH_RANK, bias=False)
        self.global_kv = nn.Linear(DIM, 2 * DEPTH_RANK, bias=False)
        self.global_out = nn.Linear(DEPTH_RANK, DIM, bias=False)
        self.null_summary = nn.Parameter(torch.zeros(1, 1, DIM))
        # Names intentionally contain mix_logit/content_gate so the existing
        # Cubic Muon builder assigns the validated gate learning-rate schedule.
        self.global_mix_logit = nn.Parameter(torch.full((DIM,), -2.0))

        self.has_depth = layer_idx > 0
        if self.has_depth:
            self.depth_q = nn.Linear(DIM, DEPTH_RANK, bias=False)
            self.depth_up = nn.Linear(DEPTH_RANK, DIM, bias=False)
            self.depth_mix_logit = nn.Parameter(torch.full((DIM,), math.atanh(0.15)))
            self.depth_content_gate = nn.Linear(DIM, 1)
            nn.init.zeros_(self.depth_content_gate.weight)
            nn.init.zeros_(self.depth_content_gate.bias)
        self.local_mask = None
        self.global_mask = None

    def set_masks(self, local_mask, global_mask):
        self.local_mask = local_mask
        self.global_mask = global_mask

    def local_branch(self, x, cos, sin):
        batch, seq, _ = x.shape
        q = self.q_proj(x).view(batch, seq, QUERY_HEADS, HEAD_DIM).transpose(1, 2)
        kv = self.kv_proj(x).view(batch, seq, 2, KV_HEADS, HEAD_DIM)
        k, v = kv.unbind(2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        q = apply_rope(q, cos, sin)
        k_rope = apply_rope(k, cos, sin)
        local = compiled_flex_attention(
            q.contiguous(), k_rope.contiguous(), v.contiguous(),
            block_mask=self.local_mask, enable_gqa=True,
        )
        return local.transpose(1, 2).reshape(batch, seq, DIM)

    def global_branch(self, x):
        batch, seq, _ = x.shape
        blocks = seq // SUMMARY_BLOCK
        summaries = x.view(batch, blocks, SUMMARY_BLOCK, DIM).mean(dim=2)
        summaries = torch.cat((self.null_summary.expand(batch, -1, -1), summaries), dim=1)
        q = self.global_q(x).view(batch, seq, DEPTH_HEADS, DEPTH_RANK // DEPTH_HEADS).transpose(1, 2)
        kv = self.global_kv(summaries).view(
            batch, blocks + 1, 2, DEPTH_HEADS, DEPTH_RANK // DEPTH_HEADS
        )
        k, v = kv.unbind(2)
        global_out = compiled_flex_attention(
            q.contiguous(), k.transpose(1, 2).contiguous(), v.transpose(1, 2).contiguous(),
            block_mask=self.global_mask,
        )
        global_out = global_out.transpose(1, 2).reshape(batch, seq, DEPTH_RANK)
        return self.global_out(global_out)

    def depth_branch(self, x, history_k, history_v):
        batch, seq, _ = x.shape
        layers = len(history_k)
        dim = DEPTH_RANK // DEPTH_HEADS
        q = self.depth_q(x).view(batch * seq, DEPTH_HEADS, 1, dim)
        k = torch.stack(history_k, dim=2).permute(0, 1, 3, 2, 4).reshape(
            batch * seq, DEPTH_HEADS, layers, dim
        )
        v = torch.stack(history_v, dim=2).permute(0, 1, 3, 2, 4).reshape(
            batch * seq, DEPTH_HEADS, layers, dim
        )
        depth = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)
        depth = depth.reshape(batch, seq, DEPTH_RANK)
        depth = self.depth_up(depth)
        gate = 2.0 * torch.sigmoid(self.depth_content_gate(x))
        return gate * self.depth_mix_logit.tanh() * depth

    def forward(self, x, cos, sin, history_k, history_v):
        local = self.local_branch(x, cos, sin)
        global_out = self.global_branch(x)
        out = local + torch.sigmoid(self.global_mix_logit) * global_out
        if self.has_depth:
            out = out + self.depth_branch(x, history_k, history_v)
        return self.out_proj(out)


class LongContextBlock(nn.Module):
    def __init__(self, layer_idx: int):
        super().__init__()
        self.norm1 = nn.RMSNorm(DIM)
        self.attn = LongContextAttention(layer_idx)
        self.norm2 = nn.RMSNorm(DIM)
        self.mlp = SwiGLU(DIM)
        self.ls1 = nn.Parameter(torch.ones(DIM))
        self.ls2 = nn.Parameter(torch.ones(DIM))

    def forward(self, x, cos, sin, history_k, history_v):
        x = x + self.ls1 * self.attn(self.norm1(x), cos, sin, history_k, history_v)
        return x + self.ls2 * self.mlp(self.norm2(x))


class CubicV11LongContext(nn.Module):
    def __init__(self, seq_len: int, use_checkpoint: bool = True, vocab_size: int = VOCAB_SIZE):
        super().__init__()
        if seq_len % SUMMARY_BLOCK:
            raise ValueError(f"seq_len must be divisible by {SUMMARY_BLOCK}")
        self.seq_len = seq_len
        self.vocab_size = vocab_size
        self.use_checkpoint = use_checkpoint
        self.embed = nn.Embedding(vocab_size, DIM)
        self.rope = RotaryEmbedding(HEAD_DIM, ROPE_BASE)
        self.depth_memory = CompressedDepthMemory()
        self.blocks = nn.ModuleList([LongContextBlock(index) for index in range(DEPTH)])
        self.norm = nn.RMSNorm(DIM)
        self.head = nn.Linear(DIM, vocab_size, bias=False)
        self.apply(self._init_weights)
        residual_std = 0.02 / math.sqrt(2 * DEPTH)
        for block in self.blocks:
            nn.init.normal_(block.attn.out_proj.weight, mean=0.0, std=residual_std)
            nn.init.normal_(block.mlp.down.weight, mean=0.0, std=residual_std)
            if block.attn.has_depth:
                nn.init.zeros_(block.attn.depth_content_gate.weight)
                nn.init.zeros_(block.attn.depth_content_gate.bias)
        self.head.weight = self.embed.weight

    @staticmethod
    def _init_weights(module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def set_masks(self, local_mask, global_mask):
        for block in self.blocks:
            block.attn.set_masks(local_mask, global_mask)

    def run_block(self, block, x, cos, sin, history_k, history_v):
        if not self.use_checkpoint or not self.training:
            return block(x, cos, sin, history_k, history_v)

        layer_count = len(history_k)

        def inner(x_value, cos_value, sin_value, *memory):
            keys = memory[:layer_count]
            values = memory[layer_count:]
            return block(x_value, cos_value, sin_value, keys, values)

        return checkpoint(
            inner, x, cos, sin, *history_k, *history_v,
            use_reentrant=False, preserve_rng_state=False,
        )

    def forward(self, tokens):
        batch, seq = tokens.shape
        if seq != self.seq_len:
            raise ValueError(f"Expected fixed sequence length {self.seq_len}, got {seq}")
        x = self.embed(tokens)
        cos, sin = self.rope(seq, x.device, x.dtype)
        history_k, history_v = [], []
        for index, block in enumerate(self.blocks):
            if index < DEPTH - 1:
                new_k, new_v = self.depth_memory(x)
            x = self.run_block(block, x, cos, sin, history_k, history_v)
            if index < DEPTH - 1:
                history_k.append(new_k)
                history_v.append(new_v)
        return self.head(self.norm(x))


def make_masks(seq_len: int, device):
    def local_causal_mask(batch, head, q_idx, kv_idx):
        return (q_idx >= kv_idx) & ((q_idx - kv_idx) < LOCAL_WINDOW)

    # kv_idx=0 is a learned null summary. kv_idx=1 summarizes tokens
    # [0, SUMMARY_BLOCK), and becomes visible starting at q=SUMMARY_BLOCK.
    def global_summary_mask(batch, head, q_idx, kv_idx):
        return (kv_idx == 0) | ((kv_idx * SUMMARY_BLOCK) <= q_idx)

    print("Creating FlexAttention block masks ...")
    local = create_block_mask(
        local_causal_mask, B=None, H=None, Q_LEN=seq_len, KV_LEN=seq_len,
        device=device, BLOCK_SIZE=128,
    )
    global_mask = create_block_mask(
        global_summary_mask, B=None, H=None, Q_LEN=seq_len,
        KV_LEN=seq_len // SUMMARY_BLOCK + 1, device=device,
        BLOCK_SIZE=128,
    )
    return local, global_mask


def configure_muon(base, seq_len: int):
    base.SEQ_LEN = seq_len
    base.BATCH_SIZE = 1
    base.STEPS = 100
    base.DIM = DIM
    base.DEPTH = DEPTH
    base.HEADS = QUERY_HEADS
    base.DEPTH_RANK = DEPTH_RANK
    base.VOCAB_SIZE = VOCAB_SIZE


def benchmark(seq_len: int):
    if DEVICE != "cuda":
        raise RuntimeError("The 16k/32k training benchmark requires CUDA")
    base = load_base()
    configure_muon(base, seq_len)
    torch.manual_seed(SEED)
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    local_mask, global_mask = make_masks(seq_len, DEVICE)
    model = CubicV11LongContext(seq_len, use_checkpoint=True).to(DEVICE)
    model.set_masks(local_mask, global_mask)
    optimizer = base.build_optimizer(model, f"CUBIC-V11-{seq_len // 1024}K")
    params = sum(parameter.numel() for parameter in model.parameters())
    tokens = torch.randint(0, VOCAB_SIZE, (1, seq_len + 1), device=DEVICE)
    x, targets = tokens[:, :-1], tokens[:, 1:]
    print("=" * 108)
    print(f"CUBIC V11 LONG CONTEXT | seq={seq_len:,} | params={params:,} | bf16 | checkpointing=on")
    print(f"local window={LOCAL_WINDOW} | summary block={SUMMARY_BLOCK} | GQA={QUERY_HEADS}:{KV_HEADS} | depth rank={DEPTH_RANK}")
    print("=" * 108)
    model.train()
    optimizer.zero_grad(set_to_none=True)
    started = time.perf_counter()
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        logits = model(x)
        loss = F.cross_entropy(logits.reshape(-1, VOCAB_SIZE), targets.reshape(-1))
    forward_seconds = time.perf_counter() - started
    started_backward = time.perf_counter()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0, foreach=True)
    optimizer.step()
    torch.cuda.synchronize()
    backward_seconds = time.perf_counter() - started_backward
    peak = torch.cuda.max_memory_allocated() / 1024**3
    reserved = torch.cuda.max_memory_reserved() / 1024**3
    total = forward_seconds + backward_seconds
    print(f"loss={loss.detach().float().item():.4f} (random-token target ~= {math.log(VOCAB_SIZE):.4f})")
    print(f"forward={forward_seconds:.2f}s | backward+Muon={backward_seconds:.2f}s | tokens/s={seq_len/total:,.0f}")
    print(f"peak allocated={peak:.2f} GiB | peak reserved={reserved:.2f} GiB")
    print("32K_TRAINING_FITS=YES" if seq_len == 32_768 else "16K_TRAINING_FITS=YES")
    return peak


def main():
    if DEVICE == "cuda":
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        torch.set_float32_matmul_precision("high")
        try:
            torch._inductor.config.triton.cudagraphs = False
        except (AttributeError, ImportError):
            pass
    last_error = None
    for length in TARGET_LENGTHS:
        try:
            benchmark(length)
            return
        except torch.cuda.OutOfMemoryError as exc:
            last_error = exc
            print(f"{length:,} OOM; clearing cache and trying the next target.")
            gc.collect()
            torch.cuda.empty_cache()
    raise RuntimeError("Neither 32k nor 16k training step fit") from last_error


if __name__ == "__main__":
    main()