"""Verificacoes matematicas do HAKO-V2 em runtime (teoremas corrigidos).

  V2-M1 (T-ATTVAR / Minkowski): alpha e combinacao convexa => Var(z_attn) <=
        max_e Var(v_e), qualquer tau. Verificacao empirica por banco.
  V2-M2 (Dirichlet / passo seguro): energia de Dirichlet do Tree-GNN publicado
        deve ser monotona nao-crescente nos passos registrados; ajuste do passo
        com eta <= 2/lambda_max (lambda_max estimado por iteracao de potencia
        no operador de mensagem W_msg).
  V2-M3 (cota de Rademacher): R <= sqrt(1+(K-1)*rho_bar)/sqrt(K) com rho_bar a
        correlacao media entre valores de especialistas; reportada por banco.
  V2-M4 (Lemma 1 BPE): treino sharded == sequencial (verificador publicado).
"""
from __future__ import annotations

from typing import Dict, List

import numpy as np
import torch


def verify_moe_contraction(sys_, zs_list: List[dict], n_max: int = 128) -> Dict:
    """V2-M1: para cada banco presente, Var_i(z_attn) <= max_e Var_i(v_e)."""
    moe = sys_.moe
    out: Dict[str, Dict] = {}
    with torch.no_grad():
        for name in list(moe.banks):
            zs = [d[name] for d in zs_list if name in d][:n_max]
            if len(zs) < 8:
                continue
            Z = torch.stack(zs)                      # (n, d_src)
            bankd = moe.banks[name]
            K, A = bankd["K"], bankd["n_adapters"]
            Wv = getattr(moe, f"Wv_{name}")
            V = torch.stack([torch.stack(
                [Wv @ (Z[i] @ getattr(moe, f"Ad_{name}_{e % A}"))
                 for i in range(Z.shape[0])]) for e in range(K)])  # (K,n,N)
            var_e = V.var(dim=1).mean(dim=1)          # (K,)
            var_max = float(var_e.max())
            # z_attn com alpha real do mecanismo publicado
            zN = torch.stack([sys_.forward_sample({name: z}, train=False)["z"]
                              for z in zs[:32]])
            zt = []
            for zi, z_s in zip(zN, zs[:32]):
                out_f, _ = moe(zi, {name: z_s})
                zt.append(out_f)
            var_attn = float(torch.stack(zt).var(dim=0).mean())
            ok = var_attn <= var_max * 1.05 + 1e-9
            out[name] = {
                "var_attn": round(var_attn, 6),
                "var_max_especialista": round(var_max, 6),
                "razao": round(var_attn / max(var_max, 1e-12), 4),
                "ok": bool(ok),
            }
    return {"teorema": "T-ATTVAR (Minkowski, combinacao convexa, tau livre)",
            "bancos": out, "passou": all(v["ok"] for v in out.values())}


def verify_dirichlet(sys_) -> Dict:
    """V2-M2: energia registrada pelo Tree-GNN publicado deve nao-crescer;
    lambda_max do operador estimado por iteracao de potencia em W_msg."""
    tgnn = sys_.tgnn
    energies = [float(e) for e in tgnn.last_energy] if tgnn.last_energy else []
    W = tgnn.W_msg.detach()
    v = torch.randn(W.shape[1])
    v /= v.norm()
    lam = 0.0
    for _ in range(64):
        w = W @ v
        nrm = w.norm()
        if nrm < 1e-12:
            break
        v = w / nrm
        lam = float(nrm)
    eta_safe = 2.0 / max(lam, 1e-9)
    mono = all(energies[i + 1] <= energies[i] + 1e-9
               for i in range(len(energies) - 1))
    return {
        "teorema": "T-GNN Dirichlet monotono com eta <= 2/lambda_max",
        "energias": [round(e, 6) for e in energies[:8]],
        "monotona": bool(mono) if energies else None,
        "lambda_max_estimado": round(lam, 4),
        "eta_seguro_publicado": round(min(0.35, eta_safe), 4),
        "passou": bool(mono),
    }


def verify_rademacher(sys_, zs_list: List[dict]) -> Dict:
    """V2-M3: cota efetiva por banco com rho_bar medido."""
    moe = sys_.moe
    out = {}
    with torch.no_grad():
        for name in list(moe.banks):
            zs = [d[name] for d in zs_list if name in d][:64]
            if len(zs) < 16:
                continue
            Z = torch.stack(zs)
            bankd = moe.banks[name]
            K, A = bankd["K"], bankd["n_adapters"]
            Wv = getattr(moe, f"Wv_{name}")
            V = torch.stack([torch.stack(
                [Wv @ (Z[i] @ getattr(moe, f"Ad_{name}_{e % A}"))
                 for i in range(Z.shape[0])]) for e in range(K)])
            Vf = V.reshape(K, -1)                      # (K, n*N)
            Vf = Vf - Vf.mean(dim=1, keepdim=True)
            Vn2 = Vf / (Vf.norm(dim=1, keepdim=True) + 1e-9)
            C = Vn2 @ Vn2.T                            # (K, K)
            off = C - torch.eye(C.shape[0]) * torch.diag(C)
            rho = float(off.abs().mean())
            Keff = K
            bound = float(np.sqrt(1.0 + (Keff - 1) * rho) / np.sqrt(Keff))
            out[name] = {"rho_bar": round(rho, 4), "K": Keff,
                         "cota_rademacher": round(bound, 4)}
    return {"teorema": "cota de Rademacher sqrt(1+(K-1)rho)/sqrt(K)",
            "bancos": out, "passou": True}


def verify_bpe_lemma(corpus_docs) -> Dict:
    """V2-M4: verificador publicada do Lemma 1 (BPE aditivo)."""
    from hako.tokenizer.byte_bpe import ByteBPETokenizer
    small = [d[:120] for d in corpus_docs[:60] if d]
    tok_full = ByteBPETokenizer(vocab_size=300, min_freq=2, shards=1)
    tok_full.train(small, log_every=0)
    tok_shard = ByteBPETokenizer(vocab_size=300, min_freq=2, shards=4)
    tok_shard.train(small, log_every=0)
    ok = tok_full.merge_history == tok_shard.merge_history
    return {"teorema": "Lemma 1 (BPE chunk-atomico, sharded == sequencial)",
            "merges": len(tok_full.merge_history), "passou": bool(ok)}