#!/usr/bin/env python3
"""
MATRIZ COMPLETA — HISTOGRAMA vs BINDING DE ORDEM
=================================================

Mede a representação de texto para inferência por geometria, em dois modelos
reais e dois regimes de corpus. A pergunta: a agregação com posição (binding
unitário) recupera melhor o token do BPE do que o histograma de bytes?

Duas representações comparadas, ambas derivadas da MESMA geometria do modelo
(base de bytes B, resolvida por mínimos quadrados no espaço dos pesos):

  H  histograma     v = sum_b B[b]              (comutativo — saco de bytes)
  B  binding        v = sum_k R_k·B[b_k]        (R_k = isometria da posição k)

R_k é uma rotação em pares de coordenadas por um ângulo irracional
(k·2π·φ⁻¹). Por ser isometria, preserva a norma exatamente — é a forma
"Parseval" desta agregação: ‖R_k x‖ = ‖x‖ para todo x.

O ângulo irracional garante que posições diferentes gerem rotações
incomensuráveis: E[⟨R_j x, R_k y⟩] = 0 para j ≠ k quando x, y são
isotrópicos. É isso que separa anagramas — o histograma tem ⟨·,·⟩ = 1.

Modelos:  Qwen2.5-0.5B-Instruct (896d), GPT-2 (768d)
Corpora:  GloVe 6B (palavras), SIFT 1M (vetores)

Métricas:
  M1  top-k de identidade de token (k = 1, 5, 10, 50)
  M2  separação de anagramas (cos entre pares)
  M3  determinismo (repetição bit a bit)
  M4  custo (build do índice, latência por consulta)
  M5  Fold Limit com a representação vencedora

Uso:
    python3 matriz_binding.py
    python3 matriz_binding.py --rapido

BSL 1.1 | (c) Winnex Brasil Soluções Empresariais LTDA-ME
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time

os.environ.setdefault("OMP_NUM_THREADS", "4")

import numpy as np

CACHE = "/tmp/models"
RESULTADOS: dict[str, object] = {}
PHI = 2 * np.pi * (np.sqrt(5) - 1) / 2      # ângulo áureo (irracional)


def titulo(t: str) -> None:
    print(f"\n{'=' * 76}\n {t}\n{'=' * 76}", flush=True)


def teste(nome: str, afirma: str, falsifica: str) -> None:
    print(f"\n--- {nome} ---", flush=True)
    print(f"    afirmação : {afirma}", flush=True)
    print(f"    falsifica : {falsifica}", flush=True)
    print("    " + "-" * 62, flush=True)


# ===========================================================================
# Geometria: base de bytes B na geometria do modelo
# ===========================================================================

def base_bytes(E: np.ndarray, vocab: dict[str, int], ridge: float = 1e-2):
    """Resolve A·B ≈ E  ->  B ∈ R^{256×D} é a base de bytes na geometria
    do modelo. Regularização de Tikhonov porque bytes correlacionados deixam
    AᵀA quase-singular."""
    V, D = E.shape
    En = E / (np.linalg.norm(E, axis=1, keepdims=True) + 1e-12)
    Ainc = np.zeros((V, 256), dtype=np.float32)
    mask = np.zeros(V, dtype=bool)
    for s, i in vocab.items():
        if i >= V:
            continue
        disp = s.replace("Ġ", " ").replace("Ċ", "\n")
        try:
            b = disp.encode("utf-8")
        except (UnicodeEncodeError, AttributeError):
            continue
        if not (0 < len(b) <= 32):
            continue
        mask[i] = True
        for byte in b:
            Ainc[i, byte] += 1.0
    idx = np.where(mask)[0]
    G = Ainc[idx].T @ Ainc[idx] + ridge * np.eye(256, dtype=np.float32)
    B = np.linalg.solve(G, Ainc[idx].T @ En[idx])
    return B, Ainc, idx


# ===========================================================================
# As duas agregações
# ===========================================================================

def _bytes(texto: str) -> np.ndarray:
    return np.frombuffer(texto.replace("Ġ", " ").replace("Ċ", "\n").encode("utf-8"),
                         dtype=np.uint8)


def agg_histograma(texto: str, B: np.ndarray) -> np.ndarray:
    """v = Σ_b B[b] — comutativo. Anagramas são idênticos por construção."""
    b = _bytes(texto)
    if len(b) == 0:
        return np.zeros(B.shape[1], dtype=np.float32)
    v = B[b].sum(0)
    return (v / (np.linalg.norm(v) + 1e-12)).astype(np.float32)


def agg_binding(texto: str, B: np.ndarray) -> np.ndarray:
    """v = Σ_k R_k·B[b_k] — a posição entra como isometria.

    R_k rotaciona cada par de coordenadas por k·PHI. Como PHI/(2π) é
    irracional, posições distintas produzem rotações incomensuráveis, e os
    produtos cruzados ⟨R_j x, R_k y⟩ se cancelam na soma. É isso que
    distingue 'need' de 'ened', que o histograma não pode distinguir.
    """
    b = _bytes(texto)
    d = B.shape[1]
    if len(b) == 0:
        return np.zeros(d, dtype=np.float32)
    v = np.zeros(d, dtype=np.float64)
    for k, byte in enumerate(b):
        u = B[byte].astype(np.float64)
        ang = PHI * k
        c, s = np.cos(ang), np.sin(ang)
        r = np.empty(d)
        r[0::2] = u[0::2] * c - u[1::2] * s
        r[1::2] = u[0::2] * s + u[1::2] * c
        v += r
    n = np.linalg.norm(v)
    return (v / (n + 1e-12)).astype(np.float32)


AGREGACOES = {"H histograma": agg_histograma, "B binding": agg_binding}


# ===========================================================================
# M1 — identidade de token
# ===========================================================================

def m1_identidade(B, Ainc, idx, tk, palavras, nome_modelo: str) -> dict:
    titulo(f"M1 — IDENTIDADE DE TOKEN  [{nome_modelo}]")
    teste("M1.1 O top-k devolve o token do BPE?",
          "o caminho texto → assinatura → busca devolve o ID que o BPE daria",
          "acurácia no nível do acaso (k/|V|)")

    simples = [(w, tk.encode(w).ids[0]) for w in palavras
               if len(tk.encode(w).ids) == 1]
    print(f"    palavras de 1 token no BPE: {len(simples)}")
    out = {}
    for nome, f in AGREGACOES.items():
        # índice: a MESMA função aplicada ao texto do token — coerência de
        # espaço é o que faz a busca funcionar
        toks = [tk.id_to_token(int(i)) or "" for i in idx]
        St = np.stack([f(t, B) for t in toks])
        St /= (np.linalg.norm(St, axis=1, keepdims=True) + 1e-12)
        St = St.astype(np.float32)

        t0 = time.time()
        linhas = {}
        for k in (1, 5, 10, 50):
            ac = 0
            for w, t in simples:
                v = f(w, B)
                s = St @ v
                top = idx[np.argsort(-s)[:k]]
                ac += int(t in set(top.tolist()))
            linhas[str(k)] = round(ac / len(simples), 4)
        dt = time.time() - t0
        out[nome] = {"topk": linhas, "tempo_s": round(dt, 1)}
        print(f"    {nome:<14} " +
              "  ".join(f"top-{k} {v * 100:5.1f}%" for k, v in linhas.items()) +
              f"   ({dt:.0f}s)")
    return out


# ===========================================================================
# M2 — anagramas
# ===========================================================================

def m2_anagramas(B) -> dict:
    titulo("M2 — SEPARAÇÃO DE ANAGRAMAS")
    teste("M2.1 Anagramas produzem vetores distintos?",
          "a representação codifica a ordem, não só o conjunto de bytes",
          "cos = 1,0 — a representação é comutativa e não pode distinguir")
    pares = [("need", "ened"), ("pc", "cp"), ("reward", "drawer"),
             ("fit", "ift"), ("use", "eus"), ("low", "owl"), ("cash", "chas")]
    out = {}
    for nome, f in AGREGACOES.items():
        cs = []
        for a, b in pares:
            va, vb = f(a, B), f(b, B)
            cs.append(float(va @ vb))
        out[nome] = {"media": round(float(np.mean(cs)), 4),
                     "min": round(float(np.min(cs)), 4),
                     "max": round(float(np.max(cs)), 4)}
        print(f"    {nome:<14} cos médio {np.mean(cs):+.4f}  "
              f"[{min(cs):+.4f}, {max(cs):+.4f}]")
    sep = out["B binding"]["max"] < 0.99
    print(f"\n    {'✓' if sep else '✗'} o binding separa todos os anagramas")
    return out


# ===========================================================================
# M3 — determinismo
# ===========================================================================

def m3_determinismo(B) -> dict:
    titulo("M3 — DETERMINISMO")
    teste("M3.1 Cinco execuções produzem o mesmo vetor?",
          "a representação é função pura do texto",
          "qualquer diferença bit a bit entre execuções")
    out = {}
    amostras = ["the", "computer", "international", "ação", "dog"]
    for nome, f in AGREGACOES.items():
        ok = all(all(np.array_equal(f(w, B), f(w, B)) for _ in range(5))
                 for w in amostras)
        out[nome] = bool(ok)
        print(f"    {'✓' if ok else '✗'} {nome:<14} idênticas em 5 execuções")
    return out


# ===========================================================================
# M4 — custo
# ===========================================================================

def m4_custo(B, Ainc, idx, tk) -> dict:
    titulo("M4 — CUSTO")
    teste("M4.1 A representação é utilizável em produção?",
          "a codificação de um texto é barata o suficiente para indexar",
          "acima de 1 ms por palavra inviabiliza indexação em escala")
    out = {}
    amostras = ["internationalization", "supercalifragilistic",
                "counterrevolutionary", "the", "dog"]
    for nome, f in AGREGACOES.items():
        t0 = time.time()
        for _ in range(200):
            for w in amostras:
                f(w, B)
        ms = (time.time() - t0) / (200 * len(amostras)) * 1000
        out[nome] = round(ms, 4)
        print(f"    {nome:<14} {ms:.4f} ms por palavra "
              f"({len('internationalization')} bytes no pior caso)")
    return out


# ===========================================================================
# main
# ===========================================================================

# nomes de arquivo por modelo: o cache local usa prefixos diferentes
MODELOS = {
    "qwen": ("Qwen/Qwen2.5-0.5B-Instruct", "model.embed_tokens.weight",
             {"config.json": "qwen_config.json",
              "model.safetensors": "qwen.safetensors",
              "tokenizer.json": "qwen_tokenizer.json"}),
    "gpt2": ("openai-community/gpt2", "wte.weight",
             {"config.json": "gpt2_config.json",
              "model.safetensors": "gpt2.safetensors",
              "tokenizer.json": "gpt2_tokenizer.json"}),
}


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--rapido", action="store_true")
    ap.add_argument("--json", default="matriz_binding.json")
    ap.add_argument("--glove", default="/tmp/glove.6B.100d.txt")
    args = ap.parse_args()

    titulo("MATRIZ COMPLETA — HISTOGRAMA vs BINDING DE ORDEM")
    print(" A pergunta: a posição, codificada como isometria, melhora a")
    print(" recuperação de identidade de token que o histograma fornece?")

    from safetensors.torch import load_file
    from tokenizers import Tokenizer

    # corpus de palavras
    palavras = []
    with open(args.glove, encoding="utf-8", errors="ignore") as f:
        for i, line in enumerate(f):
            if i >= (8000 if args.rapido else 50000):
                break
            w = line.split(" ", 1)[0]
            if 0 < len(w) <= 24:
                palavras.append(w)
    rng = np.random.default_rng(11)
    n_aval = 400 if args.rapido else 1200
    amostra = [palavras[j] for j in rng.choice(len(palavras),
                                               min(n_aval, len(palavras)),
                                               replace=False)]
    print(f"\n corpus: {len(palavras)} palavras do GloVe, "
          f"{len(amostra)} na avaliação")

    for chave, (repo, embed, nomes) in MODELOS.items():
        titulo(f"MODELO: {repo}")
        arquivos = {}
        for remoto, local in nomes.items():
            caminho = os.path.join(CACHE, local)
            if not os.path.exists(caminho):
                print(f"    (ausente: {caminho} — pulando)", flush=True)
                break
            arquivos[remoto] = caminho
        if len(arquivos) != 3:
            continue

        sd = load_file(arquivos["model.safetensors"])
        E = sd[embed].float().numpy()
        tk = Tokenizer.from_file(arquivos["tokenizer.json"])
        with open(arquivos["tokenizer.json"], encoding="utf-8") as f:
            tv = json.load(f)["model"]["vocab"]
        vocab = {k: i for k, i in
                 (tv.items() if isinstance(tv, dict) else enumerate(tv))}
        print(f"    {E.shape[0]} tokens × {E.shape[1]}d")

        B, Ainc, idx = base_bytes(E, vocab)
        print(f"    base de bytes: {B.shape} | tokens indexáveis: {len(idx)}")

        res = {
            "M1": m1_identidade(B, Ainc, idx, tk, amostra, repo),
            "M2": m2_anagramas(B),
            "M3": m3_determinismo(B),
            "M4": m4_custo(B, Ainc, idx, tk),
        }
        RESULTADOS[chave] = res

    # ---------------- relatório ----------------
    titulo("RELATÓRIO — as duas representações lado a lado")
    print(f"  {'modelo':<10}{'representação':<16}{'top-1':>8}{'top-5':>8}"
          f"{'top-10':>8}{'top-50':>8}{'anagramas':>12}{'ms':>8}")
    for chave, res in RESULTADOS.items():
        m1, m2, m4 = res["M1"], res["M2"], res["M4"]
        for nome in AGREGACOES:
            t = m1[nome]["topk"]
            print(f"  {chave:<10}{nome:<16}"
                  f"{t['1'] * 100:>7.1f}%{t['5'] * 100:>7.1f}%"
                  f"{t['10'] * 100:>7.1f}%{t['50'] * 100:>7.1f}%"
                  f"{m2[nome]['max']:>12.4f}{m4[nome]:>8.4f}")

    print("\n  leitura:")
    for chave, res in RESULTADOS.items():
        h = res["M1"]["H histograma"]["topk"]
        b = res["M1"]["B binding"]["topk"]
        ganho = (b["1"] - h["1"]) * 100
        anag_h = res["M2"]["H histograma"]["max"]
        anag_b = res["M2"]["B binding"]["max"]
        print(f"    {chave}: top-1 {h['1'] * 100:.1f}% → {b['1'] * 100:.1f}% "
              f"({ganho:+.1f} pts) | anagrama max {anag_h:.3f} → {anag_b:.3f}")
        det = res["M3"]["B binding"]
        print(f"    {'':>{len(chave)}}  binding determinístico: "
              f"{'sim' if det else 'não'} | custo {res['M4']['B binding']:.4f} ms")

    with open(args.json, "w", encoding="utf-8") as f:
        json.dump(RESULTADOS, f, indent=1, ensure_ascii=False, default=str)
    print(f"\n[ok] {args.json}")
    return 0


if __name__ == "__main__":
    sys.exit(main())