#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""features_v3.py — تجزئة متجهة متكافئة بايت-بايت مع phase_p2_slot.hashed_features_slot.

الاستراتيجية (بدل MD5 لكل ميزة في كل استدعاء):
  1) مخزون لكل كلمة: وحداتها + جميع ثلاثيات أحرفها (u: و c:) بمصطلحات (bucket, ±sign).
  2) مخزون لكل ثنائية كلمات (b:) — المفتاح نصي "t1\x00t2".
  3) مخزون لكل رقم للنوافذ N:<num>#0/1.
  4) التجميع بnumpy: np.unique(return_index) لترتيب أول-ظهور + np.bincount للمجاميع
     (القيم كلها صحيحة: مجموع إشارات ±1 ⇒ مطابقة تامة بلا أخطاء عائمة).
  أمان الترتيب: مرحلتان (أساس ثم أرقام) مطابقةً لدلالات الدالة الأصلية (التي تُسقط
  الأصفار الأساسية قبل إضافة ميزات الأرقام).
"""
import hashlib
import re
import numpy as np

_FIXED_F = 1 << 15
_TOK = re.compile(r"\w+", re.UNICODE)
_NUM = re.compile(r"\d+")

_TOK_CACHE = {}
_BI_CACHE = {}
_NUM_CACHE = {}


def _h32(f):
    """int(md5(f).hexdigest()[:8], 16) — نفس القيمة بقراءة digest مباشرة."""
    return int.from_bytes(hashlib.md5(f.encode()).digest()[:4], "big")


def _entry(feats, F):
    """يبني تسلسل (b،s) مسطّحًا كقائمة [b0,s0,b1,s1,...] لاستدعاء extend واحدة."""
    out = []
    for f in feats:
        h = _h32(f)
        out.append(h % F)
        out.append(1 if (h >> 16) & 1 else -1)
    return out


def _tok_entry(t, F):
    """[وحدات u:t، ثلاثيات #t#، قاموس الثنائيات التالية] — واحد لأن مستهلكًا متسلسلًا."""
    u = _entry(["u:" + t], F)
    w = "#" + t + "#"
    cs = []
    for j in range(len(w) - 2):
        cs += _entry(["c:" + w[j:j + 3]], F)
    return [u, cs, {}]


def _bi_entry(t1, t2, F):
    return _entry(["b:" + t1 + "_" + t2], F)


def _num_entry(t, F):
    return _entry(["N:" + t + "#0", "N:" + t + "#1"], F)


def _accumulate(flat):
    """flat = [b0,s0,b1,s1,...] → (idx, val) بترتيب أول-ظهور، متطابق مع dict الأصلي."""
    B = np.array(flat[0::2], dtype=np.int64)
    S = np.array(flat[1::2], dtype=np.int64)
    if B.size == 0:
        return np.array([], dtype=np.int64), np.array([], dtype=np.float32)
    u, first = np.unique(B, return_index=True)
    order = np.argsort(first, kind="stable")
    idx = u[order]
    vals = np.bincount(B, weights=S)[idx]
    nz = vals != 0
    return idx[nz].astype(np.int64), vals[nz].astype(np.float32)


def hashed_features_slot_v3(text, F=2 ** 15):
    """بديل متكافئ بايت-بايت لـhashed_features_slot (نفس الترتيب ونفس القيم)."""
    if F != _FIXED_F:
        raise ValueError("features_v3 مُحسَّنة لـF=32768 فقط (قيمة كل الأوزان المنشورة)")
    toks = _TOK.findall(text.lower())
    flat = []
    ext = flat.extend
    tc = _TOK_CACHE
    for i, t in enumerate(toks):
        e = tc.get(t)
        if e is None:
            e = _tok_entry(t, F)
            tc[t] = e
        ext(e[0])                                   # u:t
        if i + 1 < len(toks):
            bd = e[2]
            nxt = toks[i + 1]
            e2 = bd.get(nxt)
            if e2 is None:
                e2 = _bi_entry(t, nxt, F)
                bd[nxt] = e2
            ext(e2)                                 # b: (قبل الأحرف — كالأصل)
        ext(e[1])                                   # c:#t#[...]
    idx, val = _accumulate(flat)
    nums = _NUM.findall(text)
    if not nums:
        return idx, val
    # المرحلة 2: مطابقة دلالات الأصل (إسقاط الأصفار الأساسية قبل إضافة N:)
    flat2 = [0] * (2 * len(idx))
    flat2[0::2] = idx.tolist()
    flat2[1::2] = val.tolist()
    ext2 = flat2.extend
    nc = _NUM_CACHE
    for t in nums:
        e = nc.get(t)
        if e is None:
            e = _num_entry(t, F)
            nc[t] = e
        ext2(e)
    return _accumulate(flat2)


def cache_stats():
    return {"tokens": len(_TOK_CACHE), "bigrams": len(_BI_CACHE), "numbers": len(_NUM_CACHE)}


def clear_caches():
    _TOK_CACHE.clear(); _BI_CACHE.clear(); _NUM_CACHE.clear()


# ------------------------------------------------------------------ المُرمِّز الدفعي
def flat_all_with_counts(texts):
    """قائمة مسطّحة واحدة [b0,s0,b1,s1,...] لكل النصوص + عدّ الميزات لكل نص."""
    flat = []
    counts = np.empty(len(texts), dtype=np.int64)
    tc = _TOK_CACHE
    for idx, text in enumerate(texts):
        mark = len(flat)
        ext = flat.extend
        toks = _TOK.findall(text.lower())
        n = len(toks)
        for i, t in enumerate(toks):
            e = tc.get(t)
            if e is None:
                e = _tok_entry(t, _FIXED_F)
                tc[t] = e
            ext(e[0])
            if i + 1 < n:
                bd = e[2]
                nxt = toks[i + 1]
                e2 = bd.get(nxt)
                if e2 is None:
                    e2 = _bi_entry(t, nxt, _FIXED_F)
                    bd[nxt] = e2
                ext(e2)
            ext(e[1])
        counts[idx] = (len(flat) - mark) >> 1
    return flat, counts


def _accumulate_flat(flat, counts):
    """تراكم دفعة كاملة من قائمة مسطّحة واحدة: يُرجع (text, bucket, value) مرتّبة
    حسب (النص، أول-ظهور) — مكافئ رياضيًا لتجميع كل نص على حدة (القيم صحيحة)."""
    if not flat:
        z64 = np.empty(0, np.int64); z32 = np.empty(0, np.float32)
        return z64, z64.copy(), z32
    arr = np.fromiter(flat, dtype=np.int64, count=len(flat))
    B = arr[0::2].astype(np.int32)
    S = arr[1::2]
    T = np.repeat(np.arange(len(counts), dtype=np.int32), counts)
    key = T.astype(np.int64) * _FIXED_F + B
    sorter = np.argsort(key, kind="stable")
    sk = key[sorter]
    fm = np.empty(sk.size, dtype=bool)
    fm[0] = True
    np.not_equal(sk[1:], sk[:-1], out=fm[1:])
    u = sk[fm]
    first = sorter[fm]
    sums = np.add.reduceat(S[sorter], np.flatnonzero(fm))
    order = np.lexsort((first, u // _FIXED_F))
    return (u // _FIXED_F)[order], (u % _FIXED_F)[order], sums[order].astype(np.float32)


CHUNK = 4000   # حد أقصى للنصوص في الدفعة الواحدة (ذاكرة محكومة)


def encode_batch_v3(enc, texts):
    """بديل دفعي متكافئ بايت-بايت لـenc.encode_batch (يعيد مصفوفة (n، d) مُطبَّعة).
    يعالج على دفعات CHUNK نصًا للحفاظ على الذاكرة — المخزون يستمر عبر الدفعات."""
    if int(enc.F) != _FIXED_F:
        raise ValueError("features_v3 مُحسَّنة لـF=32768")
    if len(texts) > CHUNK:
        parts = [encode_batch_v3(enc, texts[i:i + CHUNK]) for i in range(0, len(texts), CHUNK)]
        return np.concatenate(parts, axis=0)
    n = len(texts)
    W = enc.W
    d = enc.d
    flat, counts = flat_all_with_counts(texts)
    t_s, b_s, v_s = _accumulate_flat(flat, counts)
    m = v_s != 0
    t_s, b_s, v_s = t_s[m], b_s[m], v_s[m]
    counts = np.bincount(t_s, minlength=n)
    offs = np.zeros(n + 1, dtype=np.int64)
    np.cumsum(counts, out=offs[1:])

    num_texts = [i for i in range(n) if _NUM.search(texts[i])]
    if num_texts:
        flat2 = []
        c2 = np.empty(len(num_texts), dtype=np.int64)
        for j, i in enumerate(num_texts):
            mark = len(flat2)
            a, bnd = int(offs[i]), int(offs[i + 1])
            tlen = bnd - a
            if tlen:
                seg = np.empty(2 * tlen, dtype=np.int64)
                seg[0::2] = b_s[a:bnd]
                seg[1::2] = v_s[a:bnd]
                flat2.extend(seg.tolist())
            ex = flat2.extend
            for tk in _NUM.findall(texts[i]):
                e = _NUM_CACHE.get(tk)
                if e is None:
                    e = _num_entry(tk, _FIXED_F)
                    _NUM_CACHE[tk] = e
                ex(e)
            c2[j] = (len(flat2) - mark) // 2
        t2s, b2s, v2s = _accumulate_flat(flat2, c2)
        m2 = v2s != 0
        t2s, b2s, v2s = t2s[m2], b2s[m2], v2s[m2]
        counts2 = np.bincount(t2s, minlength=len(num_texts))
        offs2 = np.zeros(len(num_texts) + 1, dtype=np.int64)
        np.cumsum(counts2, out=offs2[1:])
        loc = {i: j for j, i in enumerate(num_texts)}
    else:
        loc = {}

    V = np.zeros((n, d), dtype=np.float32)
    for i in range(n):
        if i in loc:
            j = loc[i]
            a, bnd = int(offs2[j]), int(offs2[j + 1])
            ii, vv = b2s[a:bnd], v2s[a:bnd]
        else:
            a, bnd = int(offs[i]), int(offs[i + 1])
            ii, vv = b_s[a:bnd], v_s[a:bnd]
        if ii.size == 0:
            continue
        z = vv @ W[ii]
        nz = float(np.linalg.norm(z))
        V[i] = (z / nz).astype(np.float32) if nz > 0 else z.astype(np.float32)
    return V