"""Unicode class features for the character encoder (E2).

A character's identity embedding says nothing about a code point the model has
never seen; its Unicode general category and script do. Every character gets a
(category, script) pair, and the parser adds one learned embedding for each to
the identity embedding, so an unseen separator or bracket still arrives as
"opening punctuation" or "dash".

Class 0 is reserved for positions that are not characters ([CLS], [SEP],
padding). The table is built from this interpreter's `unicodedata` and exported
into the Go artifact as ranges, so both runtimes use the same assignment even
where Unicode versions differ.
"""

import bisect
import unicodedata

import numpy as np

CATEGORIES = ("Lu", "Ll", "Lt", "Lm", "Lo", "Mn", "Mc", "Me", "Nd", "Nl", "No",
              "Pc", "Pd", "Ps", "Pe", "Pi", "Pf", "Po", "Sm", "Sc", "Sk", "So",
              "Zs", "Zl", "Zp", "Cc", "Cf", "Cs", "Co", "Cn")
# Block ranges, not the Unicode Script property: `unicodedata` does not expose
# scripts, and blocks are what separate the signals a filename cares about
# (fullwidth vs ASCII punctuation, kana vs Han). Unlisted code points are "other".
SCRIPT_BLOCKS = (
    ("latin", ((0x0000, 0x024F), (0x1E00, 0x1EFF), (0x2C60, 0x2C7F), (0xA720, 0xA7FF),
               (0xAB30, 0xAB6F))),
    ("ipa-modifier", ((0x0250, 0x036F),)),
    ("greek", ((0x0370, 0x03FF), (0x1F00, 0x1FFF))),
    ("cyrillic", ((0x0400, 0x052F), (0x1C80, 0x1C8F), (0x2DE0, 0x2DFF), (0xA640, 0xA69F))),
    ("armenian-hebrew", ((0x0530, 0x05FF),)),
    ("arabic", ((0x0600, 0x06FF), (0x0750, 0x077F), (0x08A0, 0x08FF), (0xFB50, 0xFDFF),
                (0xFE70, 0xFEFF))),
    ("indic", ((0x0900, 0x0DFF),)),
    ("southeast-asian", ((0x0E00, 0x0FFF), (0x1000, 0x109F), (0x1780, 0x17FF))),
    ("hangul", ((0x1100, 0x11FF), (0x3130, 0x318F), (0xA960, 0xA97F), (0xAC00, 0xD7FF))),
    ("general-punctuation", ((0x2000, 0x206F),)),
    ("symbols", ((0x2070, 0x2BFF),)),
    ("cjk-symbols", ((0x3000, 0x303F), (0x3200, 0x33FF), (0xFE10, 0xFE1F), (0xFE30, 0xFE4F))),
    ("hiragana", ((0x3040, 0x309F),)),
    ("katakana", ((0x30A0, 0x30FF), (0x31F0, 0x31FF), (0xFF66, 0xFF9F))),
    ("han", ((0x2E80, 0x2FDF), (0x3400, 0x4DBF), (0x4E00, 0x9FFF), (0xF900, 0xFAFF),
             (0x20000, 0x3FFFF))),
    ("fullwidth", ((0xFF00, 0xFF65), (0xFFE0, 0xFFEF))),
    ("emoji", ((0x1F000, 0x1FAFF),)),
)
SCRIPTS = tuple(name for name, _ in SCRIPT_BLOCKS) + ("other",)
N_CATEGORIES, N_SCRIPTS = len(CATEGORIES) + 1, len(SCRIPTS) + 1
LIMIT = 0x110000
# Categories that carry structure rather than content; identity dropout hits
# them harder (see Parser.character_embeddings).
STRUCTURAL = frozenset(i + 1 for i, name in enumerate(CATEGORIES) if name[0] in "PSZ")

_TABLES = None


def tables():
    """(category ids, script ids) for every code point, as uint8 arrays."""
    global _TABLES
    if _TABLES is None:
        category = np.empty(LIMIT, dtype=np.uint8)
        index = {name: i + 1 for i, name in enumerate(CATEGORIES)}
        for cp in range(LIMIT):
            category[cp] = index[unicodedata.category(chr(cp))]
        script = np.full(LIMIT, SCRIPTS.index("other") + 1, dtype=np.uint8)
        for name, ranges in SCRIPT_BLOCKS:
            for a, b in ranges:
                script[a:b + 1] = SCRIPTS.index(name) + 1
        _TABLES = category, script
    return _TABLES


def classes(char):
    category, script = tables()
    return int(category[ord(char)]), int(script[ord(char)])


def ranges():
    """Run-length [start, end, category, script] rows covering 0..0x10FFFF."""
    category, script = tables()
    combined = category.astype(np.uint16) << 8 | script
    change = np.flatnonzero(np.diff(combined)) + 1
    starts = np.concatenate(([0], change))
    ends = np.concatenate((change - 1, [LIMIT - 1]))
    return [[int(a), int(b), int(category[a]), int(script[a])] for a, b in zip(starts, ends)]


def lookup(table, cp):
    """Reference implementation of the Go range lookup (for parity tests)."""
    position = bisect.bisect_right([row[0] for row in table], cp) - 1
    return table[position][2], table[position][3]