# ============================================================
# 13 — X-Factor: reduce dimensions without retraining
#
# Colab:  !pip install winnex-madhava winnex-xfactor
# ============================================================

# 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-madhava", "winnex-xfactor"]:
    _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 numpy as np
import winnex_madhava as mad
import winnex_xfactor as xf

np.random.seed(13)

# 1. simula a tabela de embeddings de um modelo (vocabulário x dimensão)
VOCAB, DIM = 4000, 256
embed_tokens = np.random.randn(VOCAB, DIM).astype(np.float32)

print(f"embedding table : {embed_tokens.shape[0]:,} tokens x {embed_tokens.shape[1]} dims")

# 2. deriva a base do manifold — sem treino, sem rodar o modelo
P, rank, residual = xf.build_xfactor(embed_tokens, variance_tau=0.95)
print(f"manifold basis  : {P.shape}  (rank efetivo {rank})")
print(f"residual        : {residual:.4f}  (variância preservada)")

# 3. projeta textos para o subespaço
model = xf.XFactor(embed_tokens, VOCAB, DIM, 0.95)
# foldable() testa se o subespaço de dimensão k preserva tau da variância
k_alvo = 128
print(f"foldable @ k={k_alvo}  : {model.foldable(k_alvo, 0.95)}")

frases = [
    "o contrato prevê multa por atraso",
    "o paciente apresenta febre alta",
    "a receita do trimestre cresceu",
]
# embedding simples de demonstração (substitua por encoder real)

def emb(t):
    v = np.zeros(DIM, dtype=np.float32)
    for tok in t.split():
        for i, ch in enumerate(tok):
            v[(ord(ch) * 31 + i * 7) % DIM] += 1.0
    n = np.linalg.norm(v)
    return v / n if n else v

orig = np.stack([emb(f) for f in frases])
proj = model.project_batch(orig)
print(f"\noriginal        : {orig.shape}")
print(f"projetado       : {np.asarray(proj).shape}")

# 4. a busca continua funcionando no espaço projetado
corpus = np.asarray(model.project_batch(embed_tokens[:800])).astype(np.float32)
engine = mad.build_engine(corpus, dim=corpus.shape[1], k=3,
                          k1_fraction=1.0, postfilter=True,
                          audit_exhaustive=True)
q = np.asarray(model.project(orig[0])).astype(np.float32)
r = mad.retrieve(engine, q, k=3)
print(f"\nbusca no subespaço: {len(r['indices'])} resultados")
print(f"bound violations  : {r['bound_violations']}")

assert r["bound_violations"] == 0

# leitura honesta do resultado
if corpus.shape[1] < DIM:
    print(f"\nOK — dimensão reduzida de {DIM} para {corpus.shape[1]}, sem treinar nada.")
else:
    print(f"""
NOTA HONESTA
────────────
A projeção não reduziu a dimensão: a base derivada tem rank {rank} de {DIM},
e project() devolve o vetor completo reconstruído no espaço original.

Isso significa que este embed_tokens sintético (ruído gaussiano puro) não tem
estrutura de baixa dimensão. Um modelo real tem: os embeddings de tokens
ocupam um subespaço muito menor que D.

Para ver redução de verdade, use os pesos de um modelo real:
    from transformers import AutoModel
    E = AutoModel.from_pretrained('<modelo>').get_input_embeddings().weight.detach().numpy()

O X-Factor não inventa estrutura — ele a encontra quando ela existe.""")

print(f"\nrank efetivo para 95% da variância: {rank} de {DIM}")
print(f"residual da reconstrução: {residual:.4f}")