# Têtu Code 2 — inference code (tetu-AI). License: Apache 2.0, see README.md.
import numpy as np
try:
    from numba import njit
except ImportError:

    def njit(*a, **k):
        return (lambda f: f) if not a or not callable(a[0]) else a[0]
C_OTHER, C_ALNUM, C_SPACE, C_TAB, C_NL, C_DQ, C_SQ, C_HASH, C_BSL, C_OPEN, C_CLOSE = range(11)
T_NORMAL, T_RESET, T_FIM_SUF, T_FIM_MID = (0, 1, 2, 3)

class Tables:

    def __init__(self, tok):
        V = tok.get_vocab_size()
        specials = {t.content for t in tok.get_added_tokens_decoder().values()} if hasattr(tok, 'get_added_tokens_decoder') else set()
        bs = list(range(ord('!'), ord('~') + 1)) + list(range(ord('¡'), ord('¬') + 1)) + list(range(ord('®'), ord('ÿ') + 1))
        cs, n = (bs[:], 0)
        for b in range(256):
            if b not in bs:
                bs.append(b)
                cs.append(256 + n)
                n += 1
        dec = {chr(c): b for b, c in zip(bs, cs)}
        off, ln, flag, ev = (np.zeros(V, np.int64), np.zeros(V, np.int32), np.zeros(V, np.int8), [])
        for i in range(V):
            s = tok.id_to_token(i) or ''
            if s in specials or (s.startswith('<|') and s.endswith('|>')):
                flag[i] = {'<|fim_suffix|>': T_FIM_SUF, '<|fim_middle|>': T_FIM_MID}.get(s, T_RESET)
                off[i], ln[i] = (len(ev), 0)
                continue
            raw = bytes((dec.get(ch, 63) for ch in s))
            off[i] = len(ev)
            for b in raw:
                c = chr(b)
                ev.append(C_ALNUM if c.isalnum() or c == '_' or b >= 128 else C_SPACE if c == ' ' else C_TAB if c == '\t' else C_NL if c == '\n' else C_DQ if c == '"' else C_SQ if c == "'" else C_HASH if c == '#' else C_BSL if c == '\\' else C_OPEN if c in '([{' else C_CLOSE if c in ')]}' else C_OTHER)
            ln[i] = len(raw)
        pref = np.zeros(max(1, len(ev)), np.int8)
        for i in range(V):
            if ln[i]:
                raw = bytes((dec.get(ch, 63) for ch in tok.id_to_token(i) or ''))
                for j, b in enumerate(raw):
                    pref[off[i] + j] = 1 if chr(b) in 'frbuFRBU' else 0
        self.arrays = (off, ln, flag, np.asarray(ev, np.int8), pref)

@njit(cache=True)
def features(ids, off, ln, flag, ev, pref):
    n = ids.shape[0]
    out = np.zeros(n, np.int16)
    mode = 0
    q = 0
    qrun = 0
    opened = 0
    esc = 0
    depth = 0
    indent = 0
    ls = 1
    run = 0
    run_pref = 1
    s_mode = s_depth = s_indent = s_ls = 0
    for t in range(n):
        i = ids[t]
        f = flag[i]
        if f == 1:
            mode = q = qrun = opened = esc = depth = indent = run = 0
            ls = 1
            run_pref = 1
        elif f == 2:
            s_mode, s_depth, s_indent, s_ls = (mode, depth, indent, ls)
            mode = q = qrun = opened = esc = depth = indent = run = 0
            ls = 1
            run_pref = 1
        elif f == 3:
            mode, depth, indent, ls = (s_mode, s_depth, s_indent, s_ls)
            q = qrun = opened = esc = run = 0
            run_pref = 1
        else:
            o = off[i]
            for j in range(ln[i]):
                c = ev[o + j]
                if c == 4:
                    if mode == 1 or mode == 3:
                        mode = 0
                    opened = qrun = esc = run = 0
                    run_pref = 1
                    ls = 1
                    indent = 0
                    continue
                if ls == 1:
                    if c == 2:
                        indent += 1
                        continue
                    if c == 3:
                        indent += 4
                        continue
                    ls = 0
                if mode == 0:
                    if c == 1:
                        run += 1
                        if pref[o + j] == 0:
                            run_pref = 0
                    else:
                        if c == 5 or c == 6:
                            if c == 5 or run == 0 or (run_pref == 1 and run <= 2):
                                mode, q, opened, qrun = (1, c, 1, 1)
                        elif c == 7:
                            mode = 3
                        elif c == 9:
                            depth += 1
                        elif c == 10:
                            depth = depth - 1 if depth > 0 else 0
                        run = 0
                        run_pref = 1
                elif mode == 1:
                    if esc == 1:
                        esc = 0
                        opened = 0
                    elif c == q:
                        if opened == 1:
                            qrun += 1
                            if qrun == 3:
                                mode, opened, qrun = (2, 0, 0)
                        else:
                            mode, qrun = (0, 0)
                    elif opened == 1 and qrun == 2:
                        mode, opened, qrun = (0, 0, 0)
                        if c == 9:
                            depth += 1
                        elif c == 10:
                            depth = depth - 1 if depth > 0 else 0
                        elif c == 7:
                            mode = 3
                    else:
                        opened = 0
                        qrun = 0
                        if c == 8:
                            esc = 1
                elif mode == 2:
                    if esc == 1:
                        esc = 0
                    elif c == 8:
                        esc = 1
                        qrun = 0
                    elif c == q:
                        qrun += 1
                        if qrun == 3:
                            mode, qrun = (0, 0)
                    else:
                        qrun = 0
        lvl = indent // 4
        if lvl > 15:
            lvl = 15
        d = depth if depth < 15 else 15
        m = 1 if mode == 1 else 2 if mode == 2 else 3 if mode == 3 else 0
        out[t] = lvl | d << 4 | m << 8 | ls << 10
    return out

def describe(v):
    v = int(v)
    return dict(indent=v & 15, depth=v >> 4 & 15, mode=('code', 'chaîne', 'triple', 'comm.')[v >> 8 & 3], ls=v >> 10 & 1)