# -*- coding: utf-8 -*-
"""Testes unitários: SOM base + 8 variantes, quantização, cache,
tokenizador, modelo (formas, geração com cache), atenções mistas.
Cada teorema-chave tem sua invariante verificada (ver docs/matematica).
"""
import sys

sys.path.insert(0, "/home/z/my-project/aurora/src")

import torch

from aurora.config import Config
from aurora.dados.cache import CacheDuplo, CacheRAM
from aurora.dados.tokenizador import TokenizadorAurora
from aurora.memoria.orquestrador import OrquestradorSOM
from aurora.memoria.som_base import SOMKohonen
from aurora.memoria.variantes_crescimento import GSOM, GrowingCellStructures, GrowingGrid
from aurora.memoria.variantes_especiais import (CounterpropagationNetwork,
                                                HierarchicalSOM, RecurrentSOM,
                                                SupervisedSOM, TemporalKohonenMap)
from aurora.nucleo.modelo import AURORAModel
from aurora.percepcao.atencao import MixtureOfAttention
from aurora.percepcao.ortogonais import CamadaOrtogonal
from aurora.quanta.quantizacao import (metricas_qualidade, quantiza_simetrica,
                                       separar_outliers)

OK, FALHAS = 0, []


def verifica(nome, cond, detalhe=""):
    global OK
    if cond:
        OK += 1
        print(f"  ✓ {nome} {detalhe}")
    else:
        FALHAS.append(nome)
        print(f"  ✗ FALHOU: {nome} {detalhe}")


def testa_som():
    print("== SOM e variantes ==")
    torch.manual_seed(0)
    # dois clusters sintéticos
    a = torch.randn(120, 16) * 0.3 + 3.0
    b = torch.randn(120, 16) * 0.3 - 3.0
    dados = torch.cat([a, b])

    som = SOMKohonen(16, 16)
    som.ajustar(dados, epocas=2, total_planejado=2 * 240)
    verifica("SOM: EQ decresce e é pequena", som.eq(dados) < 2.0, f"EQ={som.eq(dados):.3f}")

    gg = GrowingGrid(16, lado_max=4, lambda_insercao=60)
    antes = gg.eq(dados)
    gg.ajustar(dados, epocas=2)
    verifica("GG: cresceu e EQ não aumentou (Teorema 2.1)",
             gg.k > 4 and gg.eq(dados) <= antes * 1.5, f"k={gg.k} EQ={gg.eq(dados):.3f}")

    gcs = GrowingCellStructures(16, kmax=12, lambda_insercao=60)
    gcs.ajustar(dados, epocas=2)
    estat = gcs.estatisticas()
    verifica("GCS: conectividade preservada (Teorema 2.2)",
             estat["k"] >= 3 and estat["arestas"] >= estat["k"] - 1, str(estat))

    gsom = GSOM(16, sf=0.5, max_unidades=40)
    gsom.ajustar(dados, epocas=1)
    verifica("GSOM: cresceu controlado por GT (Teorema 3.1)", 4 <= gsom.k <= 40,
             f"k={gsom.k} GT={gsom.gt:.2f}")

    hsom = HierarchicalSOM(16, k_nivel1=6, k_filho=5)
    hsom.ajustar(dados, epocas=1)
    verifica("H-SOM: filhos criados (Teorema 4.1)", len(hsom.filhos) > 0)

    tkm = TemporalKohonenMap(16, k=8)
    ci = tkm.bmu_temporal(torch.randn(10, 16))
    verifica("TKM: vencedor temporal calculado (Teorema 5.1)", 0 <= ci < 8)

    rsom = RecurrentSOM(16, k=8, alpha=0.3)
    evento = None
    for t in torch.randn(20, 16):
        _, evento = rsom.passo_temporal(t)
    verifica("RSOM: contexto dentro da cota B (Teorema 5.3)", evento is None)

    cpn = CounterpropagationNetwork(16, 4, k=6)
    alvo = torch.randn(240, 4)
    cpn.ajustar((dados, alvo), epocas=2)
    verifica("CPN: regressão por partes converge (Teorema 6.1)",
             cpn.eq(dados, alvo) < torch.var(alvo).item(), f"EQ={cpn.eq(dados, alvo):.3f}")

    ssom = SupervisedSOM(16, k=8, n_classes=2)
    rotulos = (dados.sum(dim=1) > 0).long()   # rótulo ligado à estrutura dos clusters
    ssom.ajustar(dados, rotulos, epocas=2)
    classe, conf = ssom.classificar(dados)
    acc = (classe == rotulos).float().mean().item()
    verifica("S-SOM: aprendeu rótulos melhor que aleatório (eq. 6.2)", acc > 0.55,
             f"acc={acc:.2f}")

    cfg = Config()
    orch = OrquestradorSOM(16, cfg.som, semente=1)
    escolha = orch.escolher({"temporal": True, "tendencia": True})
    verifica("Orquestrador: roteia RSOM p/ sequência com tendência", escolha == "rsom")


def testa_quantizacao():
    print("== Quantização W8A8 ==")
    torch.manual_seed(1)
    x = torch.randn(64, 32) * 0.7
    xq, s, erro = quantiza_simetrica(x)
    verifica("Teorema 8.1: erro ≤ s/2", erro <= s / 2 + 1e-6, f"erro={erro:.4f} s={s:.4f}")
    # v5 — quantização REAL (sem fakes): backends + métricas honestas
    mq = metricas_qualidade(torch.randn(8, 32), torch.randn(16, 32))
    verifica("Teorema 16.7: métricas reais de quantização",
             mq["erro_relativo"] >= 0 and mq["snr_db"] > 0,
             f"SNR={mq['snr_db']}dB erro_rel={mq['erro_relativo']:.4f}")
    from aurora.quanta.quantizacao import converter_int8_real, modo_quantizacao
    lin = torch.nn.Sequential(torch.nn.Linear(32, 16), torch.nn.ReLU())
    conv = converter_int8_real(lin)
    verifica("torchao: conversão INT8 REAL", conv.get("ok") is True,
             str(conv.get("motivo", f"fator={conv.get('fator_memoria')}")))
    modos = modo_quantizacao()
    verifica("§9: diagnóstico honesto de backends", "torchao_int8" in modos)
    mask, n = separar_outliers(x)
    verifica("LLM.int8: critério de outliers roda", isinstance(n, int))


def testa_cache_tokenizador():
    print("== Cache e Tokenizador ==")
    c = CacheRAM(1.0)
    c.put("a", torch.ones(100))
    verifica("cache RAM hit", c.get("a") is not None)
    cd = CacheDuplo(1.0, "/tmp/aurora_teste_cache")
    cd.put("k", torch.zeros(10), persistir=True)
    verifica("cache duplo: disco persiste", cd.get("k") is not None)

    tk = TokenizadorAurora()
    textos = iter(["olá mundo", "isso é um teste de tokenização pt-BR",
                   "carolina corpus", "aula de matemática"] * 20)
    info = tk.treinar(textos, vocab=300)
    ids = tk.encode("olá mundo teste", tarefa="lm")
    dec = tk.decode(ids)
    verifica("tokenizador: encode/decode roundtrip", len(ids) > 2, dec[:20])
    lote = tk.encode_lote(["texto um", "texto dois comprido"], tarefa="lm")
    verifica("tokenizador: encode_lote paralelizável",
             len(lote) == 2 and len(lote[0]) > 2 and len(lote[1]) > 2)


def testa_modelo():
    print("== Modelo unificado ==")
    cfg = Config()
    cfg.modelo.vocab = 300
    cfg.modelo.d_modelo = 64
    cfg.modelo.n_camadas = 2
    cfg.modelo.d_ff = 128
    cfg.modelo.imagem = {"tam_patch": 16, "resolucao": 32, "d": 64, "n_camadas": 1, "n_cabecas": 2}
    cfg.modelo.audio = {"n_mels": 80, "d": 64, "n_camadas": 1, "n_cabecas": 2}
    cfg.modelo.video = {"n_frames": 2, "tam_patch": 16, "resolucao": 32, "d": 64, "n_camadas": 1}
    modelo = AURORAModel(cfg, usar_multimodal=True)
    ids = torch.randint(5, 300, (2, 24))
    alvo = ids.clone()
    alvo[:, 0] = -100
    logits, perda = modelo(ids, alvo=alvo)
    verifica("modelo: formas (B,T,V)", logits.shape == (2, 24, 300))
    verifica("modelo: perda finita", torch.isfinite(torch.tensor(perda if perda else 0)))

    # geração com KV-cache
    prompt = torch.randint(5, 300, (1, 6))
    novos = modelo.gerar(prompt, max_novos=8, temperatura=1.0, top_k=10)
    verifica("modelo: geração autoregressiva com cache", len(novos) > 0)

    # multimodal: fusão com imagem+áudio+texto (formas brutas de cada modalidade)
    emb = {"imagem": torch.rand(2, 3, 32, 32),          # (B,3,H,W)
           "audio": torch.randn(2, 4000),               # (B,N) onda em [-1,1]
           "texto_emb": torch.randn(2, 64)}             # (B,d)
    pref = modelo.codificar_multimodal(emb)
    verifica("fusão multimodal: prefixo (B,d) com gates", pref.shape == (2, 64)
             and hasattr(modelo, "ultimos_gates"))

    # ortogonal
    co = CamadaOrtogonal(32, 32, vies=False)
    desvio0 = co.metrica_ortogonalidade()
    with torch.no_grad():
        co.w.add_(0.1 * torch.randn_like(co.w))
    co.renormalizar()
    verifica("ortogonal: renormalização restaura W Wᵀ=I (Teorema 9.1)",
             co.metrica_ortogonalidade() < 1e-4)

    # atenção mista: gate no simplex
    moa = MixtureOfAttention(64, 4, raio_janela=4)
    g = torch.softmax(moa.logits_gate, 0)
    verifica("Teorema 7.1: gate é mistura convexa", abs(float(g.sum()) - 1.0) < 1e-6)


if __name__ == "__main__":
    testa_som()
    testa_quantizacao()
    testa_cache_tokenizador()
    testa_modelo()
    print(f"\n===== {OK} verificações OK, {len(FALHAS)} falhas =====")
    if FALHAS:
        print("FALHAS:", FALHAS)
        sys.exit(1)