"""The network of bard-small."""
import re
import math
import time
from functools import lru_cache

import torch
from torch import nn
from torch.nn import functional as F
try:
    from torch.nn.attention.flex_attention import create_block_mask, flex_attention
except ImportError:  # pragma: no cover
    create_block_mask = flex_attention = None
from torch.utils.checkpoint import checkpoint


@lru_cache(maxsize=32)
def sliding_window_block_mask(T: int, window: int, device: str):
    def mask_mod(b, h, q_idx, kv_idx):
        causal = q_idx >= kv_idx
        in_window = (q_idx - kv_idx) < window
        return causal & in_window

    return create_block_mask(mask_mod, B=None, H=None, Q_LEN=T, KV_LEN=T, device=device)


def apply_rope(x, base=10000.0, offset=0):
    T, D = x.shape[-2], x.shape[-1]
    inv_freq = base ** (-torch.arange(0, D, 2, device=x.device, dtype=torch.float32) / D)
    ang = torch.arange(offset, offset + T, device=x.device, dtype=torch.float32)[:, None] * inv_freq
    cos, sin = ang.cos().to(x.dtype), ang.sin().to(x.dtype)
    x1, x2 = x[..., : D // 2], x[..., D // 2:]
    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)


class LayerKV:

    def __init__(self, window):
        self.window = window
        self.k = self.v = None
        self.lo = self.hi = 0
        self.start = 0

    def append(self, k, v):
        T = k.shape[2]
        if self.k is None:
            cap = max(64, 2 * T)
            self.k = k.new_empty(k.shape[:2] + (cap,) + k.shape[3:])
            self.v = v.new_empty(self.k.shape)
        elif self.hi + T > self.k.shape[2]:
            live = self.hi - self.lo
            cap = self.k.shape[2]
            if live + T > cap // 2:
                cap = 2 * (live + T)
            k_buf = self.k.new_empty(self.k.shape[:2] + (cap,) + self.k.shape[3:])
            v_buf = self.v.new_empty(k_buf.shape)
            k_buf[:, :, :live] = self.k[:, :, self.lo:self.hi]
            v_buf[:, :, :live] = self.v[:, :, self.lo:self.hi]
            self.k, self.v, self.lo, self.hi = k_buf, v_buf, 0, live
        self.k[:, :, self.hi:self.hi + T] = k
        self.v[:, :, self.hi:self.hi + T] = v
        self.hi += T
        return self.k[:, :, self.lo:self.hi], self.v[:, :, self.lo:self.hi], self.start

    def trim(self):
        if self.window is not None and self.hi - self.lo > self.window - 1:
            drop = self.hi - self.lo - (self.window - 1)
            self.lo += drop
            self.start += drop


class KVCache:

    def __init__(self, transformer):
        self.windows = [block.window for block in transformer.blocks]
        self.reset()

    def reset(self):
        self.layers = [LayerKV(w) for w in self.windows]
        self.pos = 0

    def __len__(self):
        return self.pos


def attend(q, k, v, rope, window, block_mask, cache=None, pos=0):
    if cache is None:
        if rope:
            q, k = apply_rope(q), apply_rope(k)
        if block_mask is None:
            return F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return flex_attention(q, k, v, block_mask=block_mask)

    T = q.shape[2]
    if rope:
        q, k = apply_rope(q, offset=pos), apply_rope(k, offset=pos)
    k_all, v_all, k_start = cache.append(k, v)
    S = k_all.shape[2]
    if T == 1 and (window is None or S <= window):
        out = F.scaled_dot_product_attention(q, k_all, v_all)
    elif S == T and (window is None or T <= window):
        out = F.scaled_dot_product_attention(q, k_all, v_all, is_causal=True)
    else:
        q_pos = torch.arange(pos, pos + T, device=q.device)[:, None]
        k_pos = torch.arange(k_start, k_start + S, device=q.device)[None, :]
        mask = k_pos <= q_pos
        if window is not None:
            mask = mask & (q_pos - k_pos < window)
        out = F.scaled_dot_product_attention(q, k_all, v_all, attn_mask=mask)
    cache.trim()
    return out


class Attention(nn.Module):
    def __init__(self, d_model, heads, window=None, rope=False):
        super().__init__()
        assert d_model % heads == 0
        self.d_model = d_model
        self.heads = heads
        self.head_dim = d_model // heads
        self.window = window
        self.rope = rope

        self.qkvg = nn.Linear(d_model, 4 * d_model, bias=False)
        self.out = nn.Linear(d_model, d_model, bias=False)

        self.qnorm = nn.RMSNorm(self.head_dim)
        self.knorm = nn.RMSNorm(self.head_dim)

    def forward(self, x, block_mask=None, cache=None, pos=0):
        B, T, E = x.shape
        H, D = self.heads, self.head_dim

        q, k, v, g = self.qkvg(x).split(self.d_model, dim=-1)

        q = self.qnorm(q.view(B, T, H, D)).transpose(1, 2)
        k = self.knorm(k.view(B, T, H, D)).transpose(1, 2)
        v = v.view(B, T, H, D).transpose(1, 2)
        out = attend(q, k, v, self.rope, self.window, block_mask, cache, pos)

        out = out.transpose(1, 2).reshape(B, T, E)

        out = out * torch.sigmoid(g)
        return self.out(out)


class FFN(nn.Module):
    def __init__(self, d_model, expansion=4, match_naive_params=True):
        super().__init__()
        hidden = d_model * expansion
        if match_naive_params:
            hidden = int(hidden * 2 / 3)
        hidden = max(128, (hidden // 128) * 128)
        self.hidden = hidden

        self.W_gate_up = nn.Linear(d_model, 2 * hidden, bias=False)
        self.W_down = nn.Linear(hidden, d_model, bias=False)

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


class Block(nn.Module):
    def __init__(self, d_model, heads, ffn_expansion=4, window=None, rope=False):
        super().__init__()
        self.attn = Attention(d_model, heads, window, rope)
        self.ffn = FFN(d_model, ffn_expansion)
        self.norm1 = nn.RMSNorm(d_model)
        self.norm2 = nn.RMSNorm(d_model)

    @property
    def window(self):
        return self.attn.window

    def forward(self, x, block_mask=None, cache=None, pos=0):
        x = x + self.attn(self.norm1(x), block_mask, cache, pos)
        x = x + self.ffn(self.norm2(x))
        return x


class ParallelBlock(nn.Module):

    def __init__(self, d_model, heads, ffn_expansion=4, window=None, rope=False):
        super().__init__()
        assert d_model % heads == 0
        self.d_model, self.heads, self.head_dim = d_model, heads, d_model // heads
        self.window = window
        self.rope = rope
        hidden = int(d_model * ffn_expansion * 2 / 3)
        self.hidden = max(128, (hidden // 128) * 128)

        self.norm = nn.RMSNorm(d_model)
        self.W_in = nn.Linear(d_model, 4 * d_model + 2 * self.hidden, bias=False)
        self.W_out = nn.Linear(d_model + self.hidden, d_model, bias=False)
        self.qnorm = nn.RMSNorm(self.head_dim)
        self.knorm = nn.RMSNorm(self.head_dim)

    def forward(self, x, block_mask=None, cache=None, pos=0):
        B, T, E = x.shape
        H, D = self.heads, self.head_dim
        q, k, v, g, gate, up = self.W_in(self.norm(x)).split(
            [E, E, E, E, self.hidden, self.hidden], dim=-1)

        q = self.qnorm(q.view(B, T, H, D)).transpose(1, 2)
        k = self.knorm(k.view(B, T, H, D)).transpose(1, 2)
        v = v.view(B, T, H, D).transpose(1, 2)
        a = attend(q, k, v, self.rope, self.window, block_mask, cache, pos)
        a = a.transpose(1, 2).reshape(B, T, E) * torch.sigmoid(g)

        f = F.silu(gate) * up
        return x + self.W_out(torch.cat([a, f], dim=-1))


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)


class CausalTransformer(nn.Module):
    def __init__(
        self,
        d_model: int,
        heads: int,
        n_blocks: int,
        ffn_expansion: int = 4,
        window: int = 1024,
        global_every: int = 4,
        grad_checkpoint: bool = True,
        compile_blocks: bool = True,
        parallel: bool = False,
        rope: str = "local",
    ):
        super().__init__()
        self.grad_checkpoint = grad_checkpoint

        windows = []
        for i in range(n_blocks):
            is_global = ((i + 1) % global_every == 0) or (i == n_blocks - 1)
            windows.append(None if is_global else window)
        self.windows = windows

        block_cls = ParallelBlock if parallel else Block
        self.blocks = nn.ModuleList(
            [block_cls(d_model, heads, ffn_expansion, w,
                       rope=(rope == "all" or (rope == "local" and w is not None)))
             for w in windows]
        )
        self.norm = nn.RMSNorm(d_model)

        self.apply(_init_weights)
        residual_std = 0.02 / math.sqrt(2 * n_blocks)
        for block in self.blocks:
            if parallel:
                nn.init.normal_(block.W_out.weight, mean=0.0, std=residual_std)
            else:
                nn.init.normal_(block.attn.out.weight, mean=0.0, std=residual_std)
                nn.init.normal_(block.ffn.W_down.weight, mean=0.0, std=residual_std)

        if compile_blocks:
            for block in self.blocks:
                block.compile()

    def block_masks(self, T, device):
        return {
            w: sliding_window_block_mask(T, w, str(device))
            for w in set(self.windows)
            if w is not None and w < T
        }

    cache_offset = 0
    """While a cached forward runs: the absolute position of x[:, 0] (hooks that index
    positions, like steering, need it). 0 otherwise."""

    def forward(self, x, masks=None, cache=None):
        if cache is not None:
            pos = cache.pos
            self.cache_offset = pos
            try:
                for block, layer_cache in zip(self.blocks, cache.layers):
                    x = block(x, None, layer_cache, pos)
            finally:
                self.cache_offset = 0
            cache.pos += x.shape[1]
            return self.norm(x)

        if masks is None:
            masks = self.block_masks(x.shape[1], x.device)

        for block in self.blocks:
            m = masks.get(block.window)
            if self.grad_checkpoint and self.training:
                x = checkpoint(block, x, m, use_reentrant=False)
            else:
                x = block(x, m)

        return self.norm(x)


MAX_DIGITS = 8
LEN_SHIFT = 27
VALUE_MASK = (1 << LEN_SHIFT) - 1

NUMBER_RE = re.compile(r"( ?)([0-9]+)")


def pack(digits: str) -> int:
    assert 0 < len(digits) <= MAX_DIGITS
    return int(digits) | (len(digits) << LEN_SHIFT)


def unpack(packed: int) -> str:
    return str(packed & VALUE_MASK).zfill(packed >> LEN_SHIFT)


def split_numbers(text):
    out, pos = [], 0
    for m in NUMBER_RE.finditer(text):
        if m.start() > pos:
            out.append(text[pos:m.start()])
        spaced, digits = bool(m.group(1)), m.group(2)
        for i in range(0, len(digits), MAX_DIGITS):
            out.append((spaced and i == 0, pack(digits[i:i + MAX_DIGITS])))
        pos = m.end()
    if pos < len(text):
        out.append(text[pos:])
    return out


def _lengths_values(packed):
    packed = packed.long()
    return packed >> LEN_SHIFT, packed & VALUE_MASK


def _pow10(n, device):
    return 10 ** torch.arange(n, device=device, dtype=torch.long)


class NumberEncoder(nn.Module):
    def __init__(self, d_model, n_freq=4):
        super().__init__()
        self.n_freq = n_freq
        n_feat = MAX_DIGITS * (2 * n_freq + 1) + MAX_DIGITS + 1
        self.mlp = nn.Sequential(nn.Linear(n_feat, d_model), nn.SiLU(), nn.Linear(d_model, d_model))
        for m in self.mlp:
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, std=0.02)
                nn.init.zeros_(m.bias)

    def forward(self, packed):
        L, value = _lengths_values(packed)
        digits = value.unsqueeze(-1) // _pow10(MAX_DIGITS, packed.device) % 10
        present = (torch.arange(MAX_DIGITS, device=packed.device) < L.unsqueeze(-1)).float()
        freqs = torch.arange(1, self.n_freq + 1, dtype=torch.float32, device=packed.device)
        angle = 2 * math.pi * digits.float().unsqueeze(-1) / 10 * freqs
        places = torch.cat([angle.sin(), angle.cos()], -1) * present.unsqueeze(-1)
        feats = torch.cat([
            places.flatten(-2),
            present,
            F.one_hot((L - 1).clamp(min=0), MAX_DIGITS).float() * (L > 0).unsqueeze(-1),
            (torch.log10(value.float() + 1) / MAX_DIGITS).unsqueeze(-1),
        ], -1)
        out = self.mlp(feats.to(self.mlp[0].weight.dtype))
        return out * (packed != 0).unsqueeze(-1).to(out.dtype)


class NumberDecoder(nn.Module):
    def __init__(self, d_model, width=256, heads=4, layers=2):
        super().__init__()
        self.len_head = nn.Sequential(nn.Linear(d_model, width, bias=False), nn.SiLU(),
                                      nn.Linear(width, MAX_DIGITS, bias=False))
        self.ctx = nn.Linear(d_model, width, bias=False)
        self.len_emb = nn.Embedding(MAX_DIGITS, width)
        self.place_emb = nn.Embedding(MAX_DIGITS, width)
        self.prev_emb = nn.Embedding(11, width)
        self.blocks = nn.ModuleList([Block(width, heads) for _ in range(layers)])
        self.norm = nn.RMSNorm(width)
        self.out = nn.Linear(width, 10, bias=False)
        for m in self.modules():
            if isinstance(m, (nn.Linear, nn.Embedding)):
                nn.init.normal_(m.weight, std=0.02)
        for b in self.blocks:
            nn.init.normal_(b.attn.out.weight, std=0.02 / math.sqrt(2 * layers))
            nn.init.normal_(b.ffn.W_down.weight, std=0.02 / math.sqrt(2 * layers))

    def _digit_logits(self, h, L, prev):
        j = torch.arange(MAX_DIGITS, device=h.device)
        power = (L.unsqueeze(-1) - 1 - j).clamp(min=0)
        x = (self.ctx(h) + self.len_emb(L - 1)).unsqueeze(1) + self.place_emb(power) + self.prev_emb(prev)
        for b in self.blocks:
            x = b(x)
        return self.out(self.norm(x))

    @staticmethod
    def _msb_digits(packed):
        L, value = _lengths_values(packed)
        power = (L.unsqueeze(-1) - 1 - torch.arange(MAX_DIGITS, device=packed.device)).clamp(min=0)
        digits = value.unsqueeze(-1) // (10 ** power) % 10
        valid = torch.arange(MAX_DIGITS, device=packed.device) < L.unsqueeze(-1)
        return L, digits, valid

    def nll(self, h, packed, weight=None):
        L, digits, valid = self._msb_digits(packed)
        prev = torch.cat([torch.full_like(digits[:, :1], 10), digits[:, :-1]], 1)
        len_nll = F.cross_entropy(self.len_head(h).float(), L - 1, reduction="none")
        logits = self._digit_logits(h, L, prev).float()
        dig_nll = F.cross_entropy(logits.view(-1, 10), digits.view(-1), reduction="none")
        per_num = len_nll + (dig_nll.view_as(digits) * valid).sum(-1)
        return per_num.sum() if weight is None else (per_num * weight).sum()

    @torch.no_grad()
    def generate(self, h, temperature=0.0):
        def pick(logits):
            if temperature <= 0:
                return logits.argmax(-1)
            return torch.multinomial(F.softmax(logits.float() / temperature, -1), 1).squeeze(-1)

        L = pick(self.len_head(h)) + 1
        prev = torch.full((h.shape[0], MAX_DIGITS), 10, dtype=torch.long, device=h.device)
        digits = torch.zeros_like(prev)
        for j in range(MAX_DIGITS):
            digits[:, j] = pick(self._digit_logits(h, L, prev)[:, j])
            if j + 1 < MAX_DIGITS:
                prev[:, j + 1] = digits[:, j]
        power = (L.unsqueeze(-1) - 1 - torch.arange(MAX_DIGITS, device=h.device)).clamp(min=0)
        valid = torch.arange(MAX_DIGITS, device=h.device) < L.unsqueeze(-1)
        value = (digits * 10 ** power * valid).sum(-1)
        return value | (L << LEN_SHIFT)


LOGIT_SOFTCAP = 30.0


def softcap(logits):

    return LOGIT_SOFTCAP * torch.tanh(logits.float() / LOGIT_SOFTCAP)


class LM(nn.Module):

    def __init__(self, vocab, d_model, heads, n_blocks, window, global_every, parallel=True, rope="local"):
        super().__init__()
        self.emb = nn.Embedding(vocab, d_model)
        self.num_enc = NumberEncoder(d_model)
        self.body = CausalTransformer(d_model, heads, n_blocks, window=window, global_every=global_every,
                                      grad_checkpoint=False, parallel=parallel, rope=rope, compile_blocks=False)
        self.head = nn.Linear(d_model, vocab, bias=False)
        self.num_dec = NumberDecoder(d_model)

    def hidden(self, idx, vals, masks=None, cache=None):
        return self.body(self.emb(idx) + self.num_enc(vals), masks, cache)