"""Single-nucleotide tokenizer for genomic sequences.

Vocabulary (small, fixed):
  0: <pad>
  1: <bos>
  2: <eos>
  3: <mask>
  4: A
  5: C
  6: G
  7: T
  8: N    (unknown / soft-masked nucleotide)

Lower- and upper-case nucleotides collapse. Anything not in {A,C,G,T} maps to N.
"""
from __future__ import annotations

import numpy as np
import torch

PAD_ID = 0
BOS_ID = 1
EOS_ID = 2
MASK_ID = 3
A_ID = 4
C_ID = 5
G_ID = 6
T_ID = 7
N_ID = 8

VOCAB_SIZE = 9
SPECIAL_IDS = (PAD_ID, BOS_ID, EOS_ID, MASK_ID)
NUCLEOTIDE_IDS = (A_ID, C_ID, G_ID, T_ID, N_ID)


def _build_lookup() -> np.ndarray:
    """Byte → token-id table. Length 256."""
    table = np.full(256, N_ID, dtype=np.int64)
    table[ord("A")] = A_ID
    table[ord("a")] = A_ID
    table[ord("C")] = C_ID
    table[ord("c")] = C_ID
    table[ord("G")] = G_ID
    table[ord("g")] = G_ID
    table[ord("T")] = T_ID
    table[ord("t")] = T_ID
    table[ord("U")] = T_ID
    table[ord("u")] = T_ID
    table[ord("N")] = N_ID
    table[ord("n")] = N_ID
    return table


_LOOKUP = _build_lookup()


def encode(seq: str | bytes) -> np.ndarray:
    """Encode a DNA string to a numpy int64 array of token-ids (no BOS/EOS)."""
    if isinstance(seq, str):
        seq = seq.encode("ascii", errors="replace")
    arr = np.frombuffer(seq, dtype=np.uint8)
    return _LOOKUP[arr]


def encode_torch(seq: str | bytes, device: str | torch.device = "cpu") -> torch.Tensor:
    return torch.from_numpy(encode(seq).copy()).to(device)


_DECODE_TABLE = {
    PAD_ID: "_",
    BOS_ID: "[",
    EOS_ID: "]",
    MASK_ID: "?",
    A_ID: "A",
    C_ID: "C",
    G_ID: "G",
    T_ID: "T",
    N_ID: "N",
}


def decode(ids) -> str:
    if torch.is_tensor(ids):
        ids = ids.detach().cpu().tolist()
    return "".join(_DECODE_TABLE.get(int(i), "?") for i in ids)