# ============================================================
# 14 — winnex-nano: deterministic tokenizer and weight blending
#
# Colab:  !pip install winnex-nano
#
# WHAT THIS SHOWS
#   Two components of winnex-nano that run with no model download and
#   no GPU:
#
#     1. SpectralTokenizer — character to quaternion spectrum, with a
#        perfect round-trip and no trained vocabulary
#     2. WeightBalancer   — fuse the weights of two models without
#        any gradient step
#
#   For real inference (ForwardEngine, prefill/step, bf16 on a 16 GB
#   card) see the winnex-nano documentation section on docs.winnex.ai —
#   none of the 17 examples here runs the inference path.
# ============================================================

# instalação automática — Colab, Jupyter e script local.
# No Colab, "!pip" dentro de comentário não executa: a instalação
# precisa acontecer de dentro da célula, antes dos imports.
import importlib
import subprocess
import sys

for _pkg in ["winnex-nano"]:
    _mod = _pkg.replace("-", "_")
    try:
        __import__(_mod)
        continue
    except ImportError:
        pass
    _r = subprocess.run(
        [sys.executable, "-m", "pip", "install", "-q", "--upgrade", _pkg],
        capture_output=True, text=True
    )
    # O interpretador do Colab já está em execução quando o pip termina.
    # Sem invalidar o cache, o import abaixo pode não enxergar o pacote.
    importlib.invalidate_caches()
    try:
        __import__(_mod)
    except ImportError:
        sys.stderr.write(_r.stdout or "")
        sys.stderr.write(_r.stderr or "")
        raise SystemExit(
            "\n\nNão foi possível importar " + _pkg + " após instalar.\n"
            "No Colab: menu 'Ambiente de execução' -> 'Reiniciar sessão',\n"
            "e execute esta célula novamente.\n"
        )

import time

import winnex_nano as wn

# ── 1. the spectral tokenizer ───────────────────────────────
# No BPE, no vocabulary file, no training. Characters map to a
# quaternion spectrum and decode back exactly.

tok = wn.SpectralTokenizer(embed_dim=64)

frases = [
    "Winnex AI",
    "Contrato de locação",
    "Paciente com febre alta",
    "Ñoño café — acentuação preservada",
    "日本語のテスト",
]

print("round-trip test (encode then decode must be identical)\n")
falhas = 0
for f in frases:
    estados = tok.encode(f)
    volta = tok.decode(estados)
    ok = volta == f
    falhas += 0 if ok else 1
    marca = "OK  " if ok else "FALHA"
    print(f"  [{marca}] {len(estados):3} estados  {f!r} -> {volta!r}")

assert falhas == 0, f"{falhas} round-trips falharam"
print(f"\n{len(frases)} frases, 0 falhas de round-trip.")

# ── 2. weight blending ──────────────────────────────────────
# W' = Σ αᵢ · R(qᵢ) · Wᵢ  — fuse two models without training.
# Useful when a bank has a general model and a domain model and cannot
# fine-tune either one.
bal = wn.WeightBalancer()

modelo_generalista = {"layer.0.weight": [1.0, 2.0, 3.0, 4.0]}
modelo_dominio     = {"layer.0.weight": [10.0, 20.0, 30.0, 40.0]}

peso_a = wn.BlendWeight(); peso_a.alpha = 0.7   # 70% generalista
peso_b = wn.BlendWeight(); peso_b.alpha = 0.3   # 30% domínio

fundido = bal.blend([modelo_generalista, modelo_dominio], [peso_a, peso_b])

esperado = [0.7 * 1 + 0.3 * 10, 0.7 * 2 + 0.3 * 20,
            0.7 * 3 + 0.3 * 30, 0.7 * 4 + 0.3 * 40]

print(f"\nweight blending")
print(f"  modelo A (70%): {modelo_generalista['layer.0.weight']}")
print(f"  modelo B (30%): {modelo_dominio['layer.0.weight']}")
print(f"  fundido       : {fundido['layer.0.weight']}")
print(f"  esperado      : {esperado}")

# comparação com tolerância: o blend trabalha em float32
import math
obtido = fundido["layer.0.weight"]
assert all(math.isclose(a, b, rel_tol=1e-6) for a, b in zip(obtido, esperado))
print("\nOK — round-trip exato e fusão de pesos correta, sem treinar nada.")

# ── 3. o custo real, medido ─────────────────────────────────
# "50 a 70x mais lento" é vago sem uma referência. O que dá para medir
# sem ambiguidade é a EXPANSÃO: o tokenizador gera embed_dim estados por
# BYTE UTF-8. Um BPE típico usa ~0,3 token por byte.
texto = ("O contrato de locação estabelece o pagamento mensal de aluguel "
         "corrigido anualmente pelo IGP-M. ") * 20
n_bytes = len(texto.encode("utf-8"))
estados = tok.encode(texto)
expansao = len(estados) / n_bytes

t0 = time.perf_counter()
tok.encode(texto)
t_enc = time.perf_counter() - t0

print(f"\ncost, measured not estimated:")
print(f"  text                   : {len(texto):,} chars, {n_bytes:,} UTF-8 bytes")
print(f"  states generated       : {len(estados):,}")
print(f"  expansion              : {expansao:.0f} states per byte "
      f"(BPE típico: ~0,3 token por byte)")
print(f"  encode                 : {t_enc*1000:.2f} ms "
      f"({len(texto)/t_enc:,.0f} chars/s)")
print(f"  expansion scales with embed_dim={64}: that parameter defines it.")

print("""
Escopo, dito com clareza
────────────────────────
O SpectralTokenizer é 50 a 70 vezes mais lento e cerca de 1000 vezes
menos compacto que um BPE treinado. Ele não substitui o tokenizador de
um modelo pré-treinado — é uma codificação autônoma para o motor nativo.

Para inferência real com os mesmos pesos e vocabulário do Hugging Face,
use ForwardEngine (seção winnex-nano em docs.winnex.ai). Nenhum dos 17
exemplos daqui executa o caminho de inferência — este mostra apenas o
tokenizador e o balanceador, que rodam sem pesos.
""")