"""Engram module — conditional N-gram memory via scalable lookup.

Adapted from DeepSeek's Engram (Conditional Memory via Scalable Lookup, Jan 2026,
https://github.com/deepseek-ai/Engram) for genomic LMs operating over a small
single-nucleotide vocabulary.

Key adaptations vs. the NLP reference:
  * No tokenizer compression — single-nucleotide ids are already collision-free.
  * Hashing is vectorised on the GPU (no NumPy CPU roundtrip).
  * N-gram orders default to {3, 4, 5, 6, 8} (codons, splice, TFBS, Kozak…).
  * Per-head table sizes auto-sized to next prime ≥ target.
"""
from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import List, Sequence

import torch
import torch.nn as nn
import torch.nn.functional as F


# ---------------------------------------------------------------------------
# prime utilities
# ---------------------------------------------------------------------------
def _is_prime(n: int) -> bool:
    if n < 2:
        return False
    if n < 4:
        return True
    if n % 2 == 0:
        return False
    i = 3
    while i * i <= n:
        if n % i == 0:
            return False
        i += 2
    return True


def next_prime(n: int, *, exclude: set[int] | None = None) -> int:
    exclude = exclude or set()
    candidate = max(2, n)
    if candidate % 2 == 0 and candidate != 2:
        candidate += 1
    while True:
        if _is_prime(candidate) and candidate not in exclude:
            return candidate
        candidate += 1 if candidate == 2 else 2


# ---------------------------------------------------------------------------
# config
# ---------------------------------------------------------------------------
@dataclass
class EngramConfig:
    vocab_size: int = 9                    # token vocabulary (incl. specials)
    pad_id: int = 0
    ngram_orders: Sequence[int] = (3, 4, 5, 6, 8)
    n_heads_per_order: int = 4
    table_size_targets: dict | None = None  # if None → defaults below
    d_mem: int = 64                        # per-head embedding dim
    hidden_size: int = 512
    kernel_size: int = 4
    layer_inject_ids: Sequence[int] = (1, 6)  # which backbone layers receive Engram
    seed: int = 0
    use_conv: bool = True
    gate_temp: float = 1.0
    # Fusion module across cross-position taps (depthwise conv by default,
    # alternative: cross-position MLP for richer fusion). "conv" matches the
    # paper; "mlp" replaces the depthwise conv with a 2-layer MLP that takes
    # the same `kernel_size` dilated taps and outputs the residual. Default
    # "conv" preserves existing behavior and old ckpt loading.
    fusion_type: str = "conv"   # "conv" | "mlp"
    # Conv temporal direction: True = causal (left-pad only, paper default,
    # each position sees only past dilated taps); False = bidirectional
    # (symmetric pad, sees past+future). MLM encoder can use either; paper
    # (decoder LM) is causal. Default True preserves existing behavior.
    conv_causal: bool = True
    # Gapped (wildcard) n-gram patterns, ADDED on top of the exact `ngram_orders`
    # tables. Each entry is a string over {o, X} (or '_' for X): 'o' = informative
    # position (its base contributes to the hash), 'X' = wildcard (base ignored ->
    # every base there maps to the same bucket, so 4^(#X) k-mers collapse into one
    # degenerate family). E.g. ["oXo", "Xoo", "ooX"] adds the three single-gap
    # order-3 tables. A gapped pattern owns its own K hash tables, sized by its
    # informative width (span - #wildcards), so wildcard tables are naturally
    # small. None (default) -> exact-only, byte-identical to prior behavior.
    gapped_patterns: Sequence[str] | None = None


def _default_table_targets(orders: Sequence[int]) -> dict:
    """Default table-size targets per n-gram order.

    Heuristic: for order n with vocab_size 9, the universe is 9**n.
    We exhaustively cover small orders and over-provision higher orders.
    """
    defaults = {2: 251, 3: 1009, 4: 4099, 5: 16411, 6: 65537, 7: 131101, 8: 262147,
                9: 524309, 10: 1048583, 12: 2097169}
    return {n: defaults.get(n, 65537) for n in orders}


# ---------------------------------------------------------------------------
# vectorised n-gram hashing on GPU
# ---------------------------------------------------------------------------
class NgramHasher:
    """Compute multi-head multiplicative-XOR n-gram hashes for every position.

    Causal: positions before the start are filled with `pad_id`.
    Per layer-id × per order × per head, an independent random odd multiplier
    vector and prime modulus is used. Identical recipe to the DeepSeek demo
    but evaluated entirely with torch ops on the input device.
    """

    def __init__(self, cfg: EngramConfig, layer_id: int):
        self.cfg = cfg
        self.layer_id = layer_id
        self.orders: list[int] = list(cfg.ngram_orders)
        self.n_heads: int = cfg.n_heads_per_order

        # Build the full pattern list. Each pattern is (span, wild) where `wild`
        # is a per-position tuple of bools (True = wildcard, base ignored). The
        # exact `ngram_orders` come FIRST as no-wildcard patterns, so when no
        # gapped patterns are configured the prime/RNG draw order — and thus old
        # checkpoints — is byte-identical to the previous order-based code. Any
        # explicit gapped patterns are appended after.
        patterns: list[tuple[int, tuple[bool, ...]]] = [
            (n, (False,) * n) for n in self.orders
        ]
        for spec in (cfg.gapped_patterns or []):
            span = len(spec)
            # convention: spec is written left->right = oldest (t-span+1) -> current (t).
            # internally column j of `tokens` holds position t-j (column 0 = current),
            # so reverse the spec to index columns: wild[j] = spec[span-1-j] is wildcard.
            wild = tuple(c in "Xx_" for c in reversed(spec))
            patterns.append((span, wild))
        self.patterns = patterns
        self.n_patterns = len(patterns)

        # Each pattern's table size is governed by its informative width
        # (span - #wildcards): a gapped pattern collapses 4^(#wild) k-mers into a
        # single bucket, so it needs a far smaller table than its span implies.
        defaults = {2: 251, 3: 1009, 4: 4099, 5: 16411, 6: 65537,
                    7: 131101, 8: 262147, 9: 524309, 10: 1048583, 12: 2097169}
        override = cfg.table_size_targets

        def _target_for(inf: int) -> int:
            if override is not None and inf in override:
                return override[inf]
            return defaults.get(inf, 65537)

        # per-pattern, per-head prime modulus (avoid duplicates across heads)
        seen: set[int] = set()
        self.pattern_mods: list[list[int]] = []
        for (span, wild) in patterns:
            inf = span - sum(wild)
            tgt = _target_for(inf)
            mods = []
            cur = tgt - 1
            for _ in range(self.n_heads):
                p = next_prime(cur, exclude=seen)
                seen.add(p)
                mods.append(p)
                cur = p
            self.pattern_mods.append(mods)

        # per-pattern, per-head, per-position multipliers (deterministic from
        # seed+layer). Wildcard columns are zeroed so that position contributes
        # 0 to the multiply-XOR mix -> every base there maps to the same bucket.
        # HF packaging note: NgramHasher is NOT an nn.Module, and
        # `from_pretrained` builds the model under a meta-device context — so we
        # must NOT create these tensors in __init__ (they would be meta and
        # unrecoverable). Store only the seed; build the multipliers lazily,
        # per real device, in _to(). Values are byte-identical to before.
        self._mult_seed = cfg.seed + 10007 * (layer_id + 1)
        self._mult_upper = 1 << 50   # keeps multipliers*vocab_size < 2**63 mid-XOR

        # backward-compat alias: head_mods keyed by exact order (first len(orders)
        # patterns are the exact orders, in order).
        self.head_mods: dict[int, list[int]] = {
            n: self.pattern_mods[i] for i, n in enumerate(self.orders)
        }

        self._device_multipliers: dict[str, list[torch.Tensor]] = {}
        self._device_mods: dict[str, list[torch.Tensor]] = {}

    def _build_multipliers(self, device):
        gen = torch.Generator(device="cpu")
        gen.manual_seed(self._mult_seed)
        mults = []
        for (span, wild) in self.patterns:
            # shape [n_heads, span] of odd ints
            r = torch.randint(0, self._mult_upper // 2, (self.n_heads, span),
                              generator=gen, dtype=torch.int64)
            r = r * 2 + 1   # odd
            if any(wild):
                keep = torch.tensor([0 if w else 1 for w in wild], dtype=torch.int64)
                r = r * keep.view(1, span)   # zero wildcard columns
            mults.append(r.to(device))
        return mults

    # ------------------------------------------------------------------
    def _to(self, device: torch.device):
        key = str(device)
        if key not in self._device_multipliers:
            self._device_multipliers[key] = self._build_multipliers(device)
            self._device_mods[key] = [
                torch.tensor(m, device=device, dtype=torch.int64)
                for m in self.pattern_mods
            ]
        return self._device_multipliers[key], self._device_mods[key]

    # ------------------------------------------------------------------
    @torch.no_grad()
    def hash(self, input_ids: torch.Tensor) -> torch.Tensor:
        """Return hash ids of shape [B, T, n_patterns*K] dtype int64.

        Layout along the last dim: pattern-major then head-minor —
            [(pattern[0], h=0..K-1), (pattern[1], h=0..K-1), ...]
        Exact orders come first, gapped patterns after.
        """
        assert input_ids.dim() == 2, "expected [B, T] input ids"
        B, T = input_ids.shape
        device = input_ids.device
        x64 = input_ids.to(torch.int64)
        mults, mods = self._to(device)
        pad = self.cfg.pad_id

        # Pre-compute left-shifted views of the input. shifts[k] = x shifted right by k,
        # so position t holds the token that was at position t-k (pad if out of range).
        max_n = max(span for span, _ in self.patterns)
        shifts: list[torch.Tensor] = [x64]
        for k in range(1, max_n):
            shifted = torch.full_like(x64, pad)
            shifted[:, k:] = x64[:, :-k]
            shifts.append(shifted)

        out_chunks: list[torch.Tensor] = []
        for i, (span, wild) in enumerate(self.patterns):
            mult = mults[i]                       # [K, span]  int64 (wildcard cols = 0)
            head_mods = mods[i]                   # [K]        int64
            tokens = torch.stack(shifts[:span], dim=-1)  # [B, T, span]

            # mix[b,t,k] = (tokens[b,t,0] * mult[k,0]) XOR (tokens[b,t,1] * mult[k,1]) XOR ...
            # broadcast: tokens [B,T,1,span] * mult [1,1,K,span] -> [B,T,K,span], reduce-XOR.
            # Wildcard columns have mult=0 -> contribute 0 -> XOR-identity (base ignored).
            scaled = tokens.unsqueeze(2) * mult.view(1, 1, self.n_heads, span)
            mix = scaled[..., 0]
            for j in range(1, span):
                mix = torch.bitwise_xor(mix, scaled[..., j])

            # mod per head
            head_hash = mix % head_mods.view(1, 1, self.n_heads)   # [B, T, K]
            out_chunks.append(head_hash)

        return torch.cat(out_chunks, dim=-1)      # [B, T, n_patterns*K]


# ---------------------------------------------------------------------------
# multi-head embedding: one nn.Embedding sharing the underlying buffer
# ---------------------------------------------------------------------------
class MultiHeadEmbedding(nn.Module):
    """Wrap multiple variable-size embedding tables in a single nn.Embedding.

    For head k with size N_k and dim d, addresses live in [Σ_{<k} N_j, Σ_{≤k} N_j).
    Forward expects per-position per-head ids in their *local* (per-head) range.
    """

    def __init__(self, table_sizes: Sequence[int], dim: int):
        super().__init__()
        self.table_sizes = list(table_sizes)
        self.dim = dim
        offsets = [0]
        for n in self.table_sizes[:-1]:
            offsets.append(offsets[-1] + n)
        self.register_buffer("offsets", torch.tensor(offsets, dtype=torch.long))
        self.total = sum(self.table_sizes)
        self.emb = nn.Embedding(self.total, dim)
        # Modest init — defaults to N(0, 1) which is too large for residual injection.
        nn.init.normal_(self.emb.weight, mean=0.0, std=0.02)

    def forward(self, head_ids: torch.Tensor) -> torch.Tensor:
        """head_ids: [..., H] in local per-head range. Returns [..., H, dim]."""
        shifted = head_ids + self.offsets
        return self.emb(shifted)


# ---------------------------------------------------------------------------
# the Engram block
# ---------------------------------------------------------------------------
class Engram(nn.Module):
    """A single Engram injection layer.

    Forward:
        hidden_states: [B, T, hidden]
        input_ids:     [B, T]   (token ids; required because hashing uses raw tokens)
    Returns:
        delta:         [B, T, hidden]   (to be added to the residual stream)
    """

    def __init__(self, cfg: EngramConfig, layer_id: int):
        super().__init__()
        self.cfg = cfg
        self.layer_id = layer_id
        self.hasher = NgramHasher(cfg, layer_id=layer_id)

        flat_table_sizes: list[int] = []
        for mods in self.hasher.pattern_mods:
            flat_table_sizes.extend(mods)

        self.embedding = MultiHeadEmbedding(flat_table_sizes, cfg.d_mem)

        n_heads_total = self.hasher.n_patterns * cfg.n_heads_per_order
        engram_hidden = n_heads_total * cfg.d_mem
        self.value_proj = nn.Linear(engram_hidden, cfg.hidden_size, bias=False)
        self.key_proj = nn.Linear(engram_hidden, cfg.hidden_size, bias=False)
        self.norm_q = nn.RMSNorm(cfg.hidden_size)
        self.norm_k = nn.RMSNorm(cfg.hidden_size)

        self.conv = None
        self.fusion_mlp = None
        self.conv_pad = 0
        if cfg.use_conv:
            dilation = max(span for span, _ in self.hasher.patterns)
            self.dilation = dilation
            self.conv_pad = (cfg.kernel_size - 1) * dilation
            if cfg.fusion_type == "conv":
                # depthwise 1-D causal conv over the sequence axis (paper default)
                self.conv = nn.Conv1d(
                    in_channels=cfg.hidden_size,
                    out_channels=cfg.hidden_size,
                    kernel_size=cfg.kernel_size,
                    groups=cfg.hidden_size,
                    bias=False,
                    dilation=dilation,
                )
                # zero-init: identity at initialisation (conv contributes 0 to residual)
                nn.init.zeros_(self.conv.weight)
            elif cfg.fusion_type == "mlp":
                # 2-layer MLP that fuses `kernel_size` dilated taps cross-channel.
                # Input: concat(v_t, v_{t-d}, v_{t-2d}, ..., v_{t-(k-1)d}) along channel dim.
                fusion_in = cfg.kernel_size * cfg.hidden_size
                self.fusion_mlp = nn.Sequential(
                    nn.Linear(fusion_in, 2 * cfg.hidden_size, bias=False),
                    nn.GELU(),
                    nn.Linear(2 * cfg.hidden_size, cfg.hidden_size, bias=False),
                )
                # zero-init the last layer so the fusion contributes 0 at init
                nn.init.zeros_(self.fusion_mlp[-1].weight)
            else:
                raise ValueError(
                    f"unknown fusion_type={cfg.fusion_type!r}; "
                    f"expected 'conv' or 'mlp'"
                )

    # ------------------------------------------------------------------
    def forward(self, hidden_states: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
        B, T, _ = hidden_states.shape
        # [B, T, H_total] where H_total = O * K
        hash_ids = self.hasher.hash(input_ids)
        # [B, T, H_total, d_mem] → [B, T, H_total*d_mem]
        emb = self.embedding(hash_ids).flatten(start_dim=-2)

        # context-aware gate: query from hidden_states, key from emb
        q = self.norm_q(hidden_states)
        k = self.norm_k(self.key_proj(emb))
        gate_score = (q * k).sum(dim=-1) / math.sqrt(self.cfg.hidden_size)
        # DeepSeek's nonlinear soft gate: sign·sqrt(|.|)·sigmoid
        gate = (gate_score.abs().clamp_min(1e-6).sqrt() * gate_score.sign())
        gate = torch.sigmoid(gate / self.cfg.gate_temp).unsqueeze(-1)   # [B, T, 1]

        v = self.value_proj(emb)           # [B, T, hidden]
        v_gated = gate * v                 # [B, T, hidden]

        if self.conv is not None:
            # conv expects [B, C, T]. Pad keeps output length = T.
            #   causal: all pad on the left (only past taps).
            #   bidirectional: split pad left/right (past + future taps).
            x = v_gated.transpose(1, 2)
            if self.cfg.conv_causal:
                x = F.pad(x, (self.conv_pad, 0))
            else:
                left = self.conv_pad // 2
                x = F.pad(x, (left, self.conv_pad - left))
            y = self.conv(x).transpose(1, 2)        # [B, T, hidden]
            return v_gated + y

        if self.fusion_mlp is not None:
            # left-pad causally, then extract `kernel_size` dilated taps along T
            # and concat along channel dim before MLP fusion.
            B, T, H = v_gated.shape
            padded = F.pad(v_gated.transpose(1, 2), (self.conv_pad, 0)).transpose(1, 2)
            # padded: [B, T + conv_pad, H]
            taps = [padded[:, i * self.dilation : i * self.dilation + T, :]
                    for i in range(self.cfg.kernel_size)]
            concat = torch.cat(taps, dim=-1)        # [B, T, kernel*hidden]
            y = self.fusion_mlp(concat)             # [B, T, hidden]
            return v_gated + y

        return v_gated