# ============================================================
# 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 example 15.
# ============================================================

# 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 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.")

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 — exemplo 15.
""")