src/bigru_t/model/hyp_t.py

"""hyp_t.py — V6.5-V3-hyp-synergy: Cabeça de Hipótese com Sinergia Matemática.

============================================================================
VISÃO GERAL (V6.5-V3 — Sinergia entre Hipóteses)
============================================================================

User requirement:
    "aprimorar matematicamente e logicamente os scripts de hyp_t.py
     analisando [...] aprimorar o módulo hyp_t.py para que as hipóteses
     (matematicamente e logicamente) trabalhem em sinergia enquanto
     treinam e ajustam os parâmetros".

A versão V6.4 do HypT era uma única hipótese global (delta = HypT(o)).
A versão V6.5-V3 eleva o módulo a um ENSEMBLE SINÉRGICO de N_HYPOTHESES
hipóteses paralelas que cooperam e competem durante o treino, em sinergia
com o SOM Kohonen 4D. Os quatro pilares matemáticos da sinergia:

  (1) DIVERSIDADE      — k-means++ D² scoring na inicialização dos geradores
                         garante que as N hipóteses cubram regiões distintas
                         do espaço de deltas desde o início.
  (2) COOPERAÇÃO       — regularização de consenso: penaliza deltas cuja
                         direção diverge fortemente do baricentro do ensemble
                         (Coherence Loss), evitando que hipóteses "vão para
                         o mato" e se anulem.
  (3) COMPETIÇÃO       — softmax gating: a cada batch, uma distribuição de
                         pesos π_h é aprendida sobre as N hipóteses (winner-
                         takes-more), permitindo que hipóteses promissoras
                         recebam mais градiente e hipóteses inúteis decaiam.
  (4) ORTOGONALIDADE   — Gram-Schmidt suave: o conjunto de deltas é projetado
                         num subespaço quase-ortogonal (penalidade na matriz
                         de Gram G = D Dᵀ), maximizando a cobertura direcional.

A integração com o SOM Kohonen 4D (256 neurônios em grade (4,4,4,4)) segue o
protocolo de duas fases definido pelo usuário:

  FASE1 (CONHECIMENTO) — somente o SOM organiza o mapa Kohonen; HypT NÃO é
                         ativada (sem ativar layers de hipóteses).
  FASE2 (PUNIÇÃO)      — SOM ajusta via treinamento + punição ativa; HypT é
                         treinada para gerar deltas que corrigem o SOM quando
                         o BMU não representa bem a entrada.

============================================================================
MATEMÁTICA DA SINERGIA (V6.5-V3)
============================================================================

Seja:
    x ∈ ℝ^P_som  — ativação SOM da entrada (P_som = 256 para grid (4,4,4,4))
    δ_h ∈ ℝ^P    — delta gerado pela hipótese h, h = 1..N
    D = [δ_1, …, δ_N] ∈ ℝ^{P × N}  — matriz de deltas
    π_h ∈ ℝ_+    — peso (responsabilidade) da hipótese h, Σ π_h = 1

Loss total por batch:

    L_total = L_task(D, π)                # task loss (BCE/classificação)
              + λ_coh · L_coherence(D)    # cooperação (consenso)
              + λ_orth · L_orth(D)        # ortogonalidade (diversidade)
              + λ_div · L_diversity(π)    # entropia de π (evita colapso)
              + λ_l2 · ||D||²_F           # regularização L2 (estabilidade)

Onde:
    L_coherence = (1/N) · Σ_h ||δ_h − δ̄||²  · exp(−‖δ_h‖)  # exp evita
                                                            # penalizar
                                                            # hipóteses ousadas
    L_orth      = ||G − diag(G)||²_F ,  G = D̂ D̂ᵀ  (D̂ normalizado)
                  → minimiza correlações entre deltas
    L_diversity = −H(π) = Σ_h π_h log π_h  (entropy bonus, promove uso
                                            # uniforme das N hipóteses)

Forward (sinergia):
    π = softmax(gating(x))                  # (N,) — responsabilidades
    D = [g_1(x), …, g_N(x)] · delta_scale   # (P, N)
    δ_ensemble = Σ_h π_h · δ_h              # (P,) — delta cooperativo
    → δ_ensemble é aplicado ao SOM

O gradiente flui para todos os g_h via δ_ensemble (cooperação) e
também via L_coherence/L_orth (diversidade+ortogonalidade).

============================================================================
COMPATIBILIDADE REVERSA
============================================================================

A classe `HypT` (single-hypothesis, V6.4) é mantida para scripts legados
que ainda fazem `HypT(...)` e chamam `forward(o, stop_grad=True)`.
Internamente, ela agora delega para `SynergyHypothesisEnsemble` com
n_hypotheses=1 (modo single-hyp), preservando a API antiga:

    >>> hyp = HypT(d_input=256, d_model=256, nhead=4, d_ff=512,
    ...            output_dim=256, num_layers=2, dropout=0.1)
    >>> delta = hyp(o, stop_grad=True)  # → (batch, output_dim)

A nova API sinérgica é exposta via `SynergyHypothesisEnsemble`:

    >>> ens = SynergyHypothesisEnsemble(
    ...     input_dim=256, output_dim=1024,
    ...     n_hypotheses=16, max_n_hypotheses=32,
    ...     hidden_dim=256,
    ...     lam_coh=0.1, lam_orth=0.05, lam_div=0.01, lam_l2=1e-4,
    ... )
    >>> out = ens(x)  # → dict com delta_ensemble, deltas, pi, losses

============================================================================
AUTOR: V6.5-V3-hyp-synergy
============================================================================
"""
from __future__ import annotations

import math
import logging
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F

logger = logging.getLogger(__name__)


# ============================================================================
# 1. Configuração canônica (NUNCA alterar — diretivas do usuário)
# ============================================================================
N_HYPOTHESES_CANONICAL: int = 16
MAX_N_HYPOTHESES_CANONICAL: int = 32
HYP_TRAIN_STEPS_CANONICAL: int = 30
HYP_HIDDEN_DIM_CANONICAL: int = 256


# ============================================================================
# 2. SynergyConfig — configuração da sinergia
# ============================================================================
@dataclass
class SynergyConfig:
    """Configuração dos pesos das regularizações de sinergia.

    Defaults calibrados para P_som=256 (grid (4,4,4,4)) e N=16 hipóteses.
    """
    # Pesos das regularizações (todos ≥ 0; 0 desativa)
    lam_coh: float = 0.1         # cooperação (consenso de direção)
    lam_orth: float = 0.05       # ortogonalidade (diversidade direcional)
    lam_div: float = 0.01        # entropia de π (uso uniforme)
    lam_l2: float = 1.0e-4       # regularização L2 dos deltas
    # Temperature do softmax gating (controla "winner-takes-more")
    gating_temperature: float = 1.0
    # Epsilon para estabilidade numérica
    eps: float = 1.0e-8
    # Delta scale inicial (pequeno para evitar explosão)
    delta_scale_init: float = 0.01
    # Ativar/desativar componentes individualmente
    enable_coherence: bool = True
    enable_orthogonality: bool = True
    enable_diversity: bool = True
    enable_l2: bool = True
    # OOM-guard: máximo de bytes por tensor de deltas (n_hyp × P × 4)
    max_delta_tensor_bytes: int = 256 * 1024 * 1024  # 256 MiB

    def to_dict(self) -> Dict[str, Any]:
        return {
            "lam_coh": float(self.lam_coh),
            "lam_orth": float(self.lam_orth),
            "lam_div": float(self.lam_div),
            "lam_l2": float(self.lam_l2),
            "gating_temperature": float(self.gating_temperature),
            "eps": float(self.eps),
            "delta_scale_init": float(self.delta_scale_init),
            "enable_coherence": bool(self.enable_coherence),
            "enable_orthogonality": bool(self.enable_orthogonality),
            "enable_diversity": bool(self.enable_diversity),
            "enable_l2": bool(self.enable_l2),
            "max_delta_tensor_bytes": int(self.max_delta_tensor_bytes),
        }


# ============================================================================
# 3. DeltaGenerator — uma hipótese individual
# ============================================================================
class DeltaGenerator(nn.Module):
    """Gerador de delta: g_h(x) → δ_h ∈ ℝ^P.

    Arquitetura MLP com 3 camadas (input → hidden → hidden → output).
    Inicialização Xavier com gain=0.01 garante deltas pequenos no início.

    Sinergia V6.5-V3: usa TransformerEncoder (em vez de MLP puro) quando
    `use_transformer=True` (default False — MLP é mais leve em CPU).

    Args:
        input_dim: dimensão da entrada x (ativação SOM flatten).
        output_dim: dimensão do delta (nº de parâmetros do SOM achatados).
        hidden_dim: largura das camadas internas (HYP_HIDDEN_DIM=256).
        use_transformer: se True, usa TransformerEncoder (mais expressivo,
                         porém mais custoso).
    """

    def __init__(
        self,
        input_dim: int,
        output_dim: int,
        hidden_dim: int = HYP_HIDDEN_DIM_CANONICAL,
        use_transformer: bool = False,
        nhead: int = 4,
        num_layers: int = 1,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.input_dim = int(input_dim)
        self.output_dim = int(output_dim)
        self.hidden_dim = int(hidden_dim)
        self.use_transformer = bool(use_transformer)

        if use_transformer:
            # Encoder + proj_out (compatível com HypT V6.4)
            self.proj_in = nn.Linear(input_dim, hidden_dim)
            encoder_layer = nn.TransformerEncoderLayer(
                hidden_dim, nhead, hidden_dim * 4, dropout=dropout,
                batch_first=True,
            )
            self.encoder = nn.TransformerEncoder(encoder_layer, num_layers)
            self.proj_out = nn.Linear(hidden_dim, output_dim)
        else:
            # MLP leve (default — otimizado para CPU/Xeon)
            self.net = nn.Sequential(
                nn.Linear(input_dim, hidden_dim),
                nn.GELU(),
                nn.Linear(hidden_dim, hidden_dim),
                nn.GELU(),
                nn.Linear(hidden_dim, output_dim),
            )

        # Inicialização Xavier com gain=0.01 (deltas pequenos no início)
        self._init_weights_xavier_small()

    def _init_weights_xavier_small(self) -> None:
        """Inicialização Xavier com gain=0.01 para garantir deltas pequenos."""
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight, gain=0.01)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """x: (batch, input_dim) → δ: (batch, output_dim)."""
        if self.use_transformer:
            # (batch, 1, input_dim) → encoder → (batch, 1, hidden) → proj
            h = self.proj_in(x).unsqueeze(1)        # (B, 1, H)
            h = self.encoder(h)                      # (B, 1, H)
            return self.proj_out(h.squeeze(1))       # (B, output_dim)
        return self.net(x)                           # (B, output_dim)


# ============================================================================
# 4. k-means++ Diversidade na Inicialização
# ============================================================================
def _kmeans_pp_init_delta_directions(
    n_hypotheses: int,
    output_dim: int,
    device: torch.device,
    seed: Optional[int] = None,
) -> torch.Tensor:
    """Inicializa N direções quasi-ortogonais usando k-means++ (Arthur &
    Vassilvitskii, 2007).

    Gera um conjunto de vetores unitários em ℝ^P tais que cada novo vetor
    é selecionado com probabilidade proporcional a D(i)² (distância ao
    quadrado do vetor mais próximo já escolhido). Isto maximiza a cobertura
    angular do conjunto inicial, fornecendo diversidade desde o step 0.

    Args:
        n_hypotheses: número de direções a gerar.
        output_dim: dimensionalidade do espaço (P).
        device: dispositivo do tensor retornado.
        seed: seed opcional para reprodutibilidade.

    Returns:
        Tensor (n_hypotheses, output_dim) com direções unitárias.
    """
    if seed is not None:
        g = torch.Generator(device=device).manual_seed(int(seed))
    else:
        g = None

    n = max(1, int(n_hypotheses))
    P = int(output_dim)
    directions = torch.zeros(n, P, device=device, dtype=torch.float32)

    # 1. Primeira direção: aleatória uniforme na esfera unitária P-dim
    v0 = torch.randn(P, device=device, generator=g)
    v0 = v0 / (v0.norm() + 1.0e-8)
    directions[0] = v0

    # 2. Demais direções: k-means++ com D² sampling
    for i in range(1, n):
        # Distância² de cada candidato (aleatório) ao conjunto já escolhido
        # Amostra um pool de candidatos (4× n para diversificar)
        n_candidates = max(16, 4 * n)
        cand = torch.randn(n_candidates, P, device=device, generator=g)
        cand = cand / (cand.norm(dim=-1, keepdim=True) + 1.0e-8)
        # D² ao vetor mais próximo já escolhido
        # dist²[j, k] = ||cand[j] - directions[k]||²
        #             = 2 − 2·<cand[j], directions[k]>  (unitários)
        chosen = directions[:i]                                # (i, P)
        dots = cand @ chosen.t()                               # (n_cand, i)
        dist_sq = (1.0 - dots).clamp(min=0.0).sum(dim=-1) * 2.0  # (n_cand,)
        # Probabilidade k-means++: p(j) ∝ D(j)²
        probs = dist_sq / (dist_sq.sum() + 1.0e-8)
        # Amostra 1 índice segundo probs
        idx = torch.multinomial(probs, num_samples=1, generator=g).item()
        directions[i] = cand[idx]

    return directions


# ============================================================================
# 5. Gating Network — aprende π (responsabilidades) sobre as N hipóteses
# ============================================================================
class GatingNetwork(nn.Module):
    """Rede de gating: π = softmax(gating(x) / T).

    Aprende uma distribuição de responsabilidades π sobre as N hipóteses
    com base na entrada x. Inicializada para produzir π ≈ uniforme no
    início (sem viés), permitindo que o treino desloque a distribuição
    conforme as hipóteses ganham/ perdem utilidade.

    Args:
        input_dim: dimensão da entrada x.
        n_hypotheses: número de hipóteses (= nº de logits de saída).
        hidden_dim: largura interna.
        temperature: temperatura do softmax (controla sharpness).
    """

    def __init__(
        self,
        input_dim: int,
        n_hypotheses: int,
        hidden_dim: int = 64,
        temperature: float = 1.0,
    ):
        super().__init__()
        self.temperature = float(temperature)
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, n_hypotheses),
        )
        # Inicialização zero → logits=0 → π uniforme no início
        for m in self.net.modules():
            if isinstance(m, nn.Linear):
                nn.init.zeros_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """x: (batch, input_dim) → π: (batch, n_hypotheses), Σ π_h = 1."""
        logits = self.net(x) / max(self.temperature, 1.0e-6)
        return F.softmax(logits, dim=-1)


# ============================================================================
# 6. SynergyHypothesisEnsemble — Ensemble Sinérgico (núcleo do V6.5-V3)
# ============================================================================
class SynergyHypothesisEnsemble(nn.Module):
    """Ensemble sinérgico de N hipóteses para correção delta do SOM.

    Implementa os 4 pilares matemáticos da sinergia:
        (1) DIVERSIDADE      — k-means++ na inicialização das direções
        (2) COOPERAÇÃO       — coherence loss (consenso de direção)
        (3) COMPETIÇÃO       — gating softmax com entropia bonus
        (4) ORTOGONALIDADE   — Gram-Schmidt suave via Frobenius em G = D̂ D̂ᵀ

    Args:
        input_dim: dimensão da ativação SOM (x).
        output_dim: dimensão do delta (P = I*J*K*L*4 para SOM 4D).
        n_hypotheses: nº ATIVO inicial de hipóteses (default 16).
        max_n_hypotheses: nº MÁXIMO pré-alocado (default 32).
        hidden_dim: largura interna dos geradores (default 256).
        config: SynergyConfig com pesos das regularizações.
        use_transformer: se True, geradores usam TransformerEncoder.

    Constantes canônicas (NUNCA alterar):
        N_HYPOTHESES = 16, MAX_N_HYPOTHESES = 32,
        HYP_TRAIN_STEPS = 30, HYP_HIDDEN_DIM = 256.
    """

    def __init__(
        self,
        input_dim: int,
        output_dim: int,
        n_hypotheses: int = N_HYPOTHESES_CANONICAL,
        max_n_hypotheses: int = MAX_N_HYPOTHESES_CANONICAL,
        hidden_dim: int = HYP_HIDDEN_DIM_CANONICAL,
        config: Optional[SynergyConfig] = None,
        use_transformer: bool = False,
        device: Optional[torch.device] = None,
        seed: Optional[int] = 42,
    ):
        super().__init__()
        # Sanity checks
        if input_dim < 1 or output_dim < 1:
            raise ValueError(
                f"SynergyHypothesisEnsemble: input_dim={input_dim} e "
                f"output_dim={output_dim} devem ser ≥ 1"
            )
        n_hypotheses = max(1, min(int(n_hypotheses), int(max_n_hypotheses)))

        self.input_dim = int(input_dim)
        self.output_dim = int(output_dim)
        self.n_hypotheses = int(n_hypotheses)
        self.max_n_hypotheses = int(max_n_hypotheses)
        self.active_count = int(n_hypotheses)
        self.hidden_dim = int(hidden_dim)
        self.config = config or SynergyConfig()
        self.use_transformer = bool(use_transformer)
        self._seed = int(seed) if seed is not None else None

        # Pré-aloca max_n_hypotheses geradores
        self.generators = nn.ModuleList([
            DeltaGenerator(
                input_dim=self.input_dim,
                output_dim=self.output_dim,
                hidden_dim=self.hidden_dim,
                use_transformer=self.use_transformer,
            )
            for _ in range(self.max_n_hypotheses)
        ])

        # Rede de gating (competição)
        self.gating = GatingNetwork(
            input_dim=self.input_dim,
            n_hypotheses=self.max_n_hypotheses,
            hidden_dim=max(32, self.hidden_dim // 4),
            temperature=self.config.gating_temperature,
        )

        # delta_scale: parâmetro treinável (começa pequeno, cresce conforme
        # o treino aprende a usar os deltas)
        self.delta_scale = nn.Parameter(
            torch.tensor(self.config.delta_scale_init, dtype=torch.float32)
        )

        # k-means++ diversity init: aplica direções quasi-ortogonais à
        # última camada linear de cada gerador (bias inicial das direções).
        self._apply_kmeans_pp_init(device=device or torch.device("cpu"))

        # Estado interno (estatísticas para monitoramento)
        self._reset_runtime_stats()

    # ------------------------------------------------------------------
    # k-means++ init
    # ------------------------------------------------------------------
    def _apply_kmeans_pp_init(self, device: torch.device) -> None:
        """Aplica direções quasi-ortogonais à última camada de cada gerador.

        As direções k-means++ são geradas no espaço ``input_dim`` (espaço de
        entrada dos geradores). Cada gerador ``h`` recebe um ``vetor direção``
        distinto quasi-ortogonal aos demais, e esse vetor é ``misturado`` a
        uma fração das ``linhas`` da matriz de pesos da camada Linear final
        (shape ``(output_dim, input_dim)``). Como cada linha de ``W`` é um
        vetor no espaço ``input_dim``, substituir parcialmente essas linhas
        pelas direções k-means++ faz com que os ``geradores`` projetem a
        entrada em subespaços distintos — produzindo deltas diversificados
        desde o início do treino.
        """
        try:
            # Direções no espaço input_dim (não output_dim!)
            directions = _kmeans_pp_init_delta_directions(
                n_hypotheses=self.max_n_hypotheses,
                output_dim=self.input_dim,
                device=device,
                seed=self._seed,
            )  # (max_n, input_dim)
            for h, gen in enumerate(self.generators):
                if self.use_transformer:
                    last_layer = gen.proj_out
                else:
                    # A última camada do nn.Sequential é a Linear final
                    last_layer = None
                    for m in gen.net:
                        if isinstance(m, nn.Linear):
                            last_layer = m
                if last_layer is None:
                    continue
                with torch.no_grad():
                    # W: (output_dim, input_dim) — linhas são vetores em input_dim
                    W = last_layer.weight.data
                    n_rows = W.shape[0]  # = output_dim
                    # Substitui 10% das linhas pela direção k-means++ do gerador h
                    n_replace = max(1, n_rows // 10)
                    idx = torch.randperm(n_rows, device=W.device)[:n_replace]
                    # Escala pequena para manter deltas pequenos
                    scale = 0.05
                    dir_h = directions[h].to(W.device)  # (input_dim,)
                    # Mistura: W[i] ← scale · (dir_h ⊙ ruído)
                    for i in idx.tolist():
                        noise = torch.randn(W.shape[1], device=W.device) * 0.1
                        W[i, :] = scale * (dir_h + noise)
            logger.info(
                "[SynergyHypothesisEnsemble] k-means++ diversity init applied "
                f"to {self.max_n_hypotheses} generators (input_dim="
                f"{self.input_dim}, output_dim={self.output_dim})."
            )
        except Exception as e:
            logger.warning(
                f"[SynergyHypothesisEnsemble] k-means++ init failed: {e}. "
                f"Falling back to Xavier-only init."
            )

    # ------------------------------------------------------------------
    # Runtime stats
    # ------------------------------------------------------------------
    def _reset_runtime_stats(self) -> None:
        """Reseta estatísticas internas para monitoramento."""
        self._last_pi: Optional[torch.Tensor] = None
        self._last_deltas: Optional[torch.Tensor] = None
        self._last_losses: Dict[str, float] = {}
        self._forward_count: int = 0

    def set_active_count(self, n: int) -> None:
        """Ajusta dinamicamente o número de hipóteses ativas."""
        n = max(1, min(int(n), self.max_n_hypotheses))
        self.active_count = int(n)
        self.n_hypotheses = int(n)

    # ------------------------------------------------------------------
    # Forward principal (sinergia)
    # ------------------------------------------------------------------
    def forward(
        self,
        x: torch.Tensor,
        return_components: bool = True,
        stop_grad_x: bool = False,
    ) -> Dict[str, torch.Tensor]:
        """Forward sinérgico do ensemble.

        Args:
            x: (batch, input_dim) — ativação SOM da entrada.
            return_components: se True, retorna dict com todos os componentes
                               (deltas, π, perdas). Se False, retorna apenas
                               {'delta_ensemble': ...}.
            stop_grad_x: se True, aplica x.detach() (Lema 3 — isola correção
                         do gradiente principal).

        Returns:
            Dict com:
                'delta_ensemble': (batch, output_dim) — delta cooperativo
                'deltas':          (batch, n_active, output_dim) — todos os δ_h
                'pi':              (batch, n_active) — responsabilidades
                'delta_scale':     escalar — escala aplicada
                'losses':          dict de regularizações (se return_components)
        """
        # OOM-guard: verifica tamanho esperado do tensor de deltas
        # (batch × n_active × output_dim × 4 bytes)
        est_bytes = (
            x.shape[0] * self.active_count * self.output_dim * 4
        )
        if est_bytes > self.config.max_delta_tensor_bytes:
            # Reduz n_active temporariamente para caber no budget
            new_n = max(
                1,
                int(self.config.max_delta_tensor_bytes
                    // (x.shape[0] * self.output_dim * 4)),
            )
            new_n = min(new_n, self.active_count)
            if new_n < self.active_count:
                logger.warning(
                    f"[SynergyHypothesisEnsemble] OOM-guard: reducing "
                    f"n_active {self.active_count}→{new_n} (est. "
                    f"{est_bytes/1e6:.1f}MB > budget "
                    f"{self.config.max_delta_tensor_bytes/1e6:.1f}MB)."
                )
                # Ajusta temporariamente (não persistente)
                original_active = self.active_count
                self.active_count = new_n
            else:
                original_active = self.active_count
        else:
            original_active = self.active_count

        try:
            # Lema 3: stop_grad isola o sinal de correção
            if stop_grad_x:
                x = x.detach()

            batch_size = x.shape[0]
            n_active = self.active_count

            # 1. Gera N deltas: (batch, n_active, output_dim)
            deltas = torch.stack(
                [self.generators[h](x) for h in range(n_active)],
                dim=1,
            )  # (batch, n_active, output_dim)

            # 2. Aplica delta_scale (escalar treinável)
            scale = self.delta_scale
            deltas = scale * deltas

            # 3. Gating: π = softmax(gating(x)) → (batch, n_active)
            # Pegamos apenas os primeiros n_active logits
            pi_full = self.gating(x)  # (batch, max_n)
            pi = pi_full[:, :n_active]
            # Renormaliza π após o slice (já que cortamos os inativos)
            pi = pi / (pi.sum(dim=-1, keepdim=True) + self.config.eps)

            # 4. Delta ensemble (cooperação): δ_ens = Σ_h π_h · δ_h
            # pi: (batch, n_active) → (batch, n_active, 1)
            delta_ensemble = (pi.unsqueeze(-1) * deltas).sum(dim=1)  # (batch, output_dim)

            # 5. Computa regularizações de sinergia
            losses: Dict[str, torch.Tensor] = {}
            if return_components:
                losses = self._compute_synergy_losses(deltas, pi)

            # Atualiza stats
            self._last_pi = pi.detach()
            self._last_deltas = deltas.detach()
            self._last_losses = {
                k: float(v.item()) if torch.is_tensor(v) else float(v)
                for k, v in losses.items()
            }
            self._forward_count += 1

            result: Dict[str, torch.Tensor] = {
                "delta_ensemble": delta_ensemble,
                "deltas": deltas,
                "pi": pi,
                "delta_scale": scale,
            }
            if return_components:
                result["losses"] = losses
            return result

        finally:
            # Restaura active_count se foi temporariamente reduzido
            self.active_count = original_active

    # ------------------------------------------------------------------
    # Regularizações de sinergia
    # ------------------------------------------------------------------
    def _compute_synergy_losses(
        self,
        deltas: torch.Tensor,  # (batch, n_active, output_dim)
        pi: torch.Tensor,       # (batch, n_active)
    ) -> Dict[str, torch.Tensor]:
        """Computa as 4 regularizações de sinergia.

        L_coherence: penaliza divergência do baricentro (cooperação)
        L_orthogonal: penaliza correlação entre deltas (diversidade)
        L_diversity:  entropia negativa de π (uso uniforme)
        L_l2:         norma Frobenius de D (estabilidade)
        """
        cfg = self.config
        eps = cfg.eps
        losses: Dict[str, torch.Tensor] = {}

        # (1) COOPERAÇÃO — coherence loss
        if cfg.enable_coherence and cfg.lam_coh > 0:
            # Baricentro do ensemble: δ̄ = (1/N) Σ_h δ_h
            delta_bar = deltas.mean(dim=1, keepdim=True)  # (batch, 1, output_dim)
            # Distância² ao baricentro, com decaimento exponencial para
            # não penalizar hipóteses ousadas (||δ_h|| grande)
            diff = deltas - delta_bar  # (batch, n_active, output_dim)
            diff_sq = (diff * diff).sum(dim=-1)  # (batch, n_active)
            # Peso exponencial: w_h = exp(-||δ_h||)
            delta_norm = deltas.norm(dim=-1)  # (batch, n_active)
            w = torch.exp(-delta_norm)  # (batch, n_active)
            # L_coh = mean_{b,h} w_{b,h} · ||δ_h - δ̄||²
            l_coh_raw = (w * diff_sq).mean()
            losses["loss_coherence"] = cfg.lam_coh * l_coh_raw
        else:
            losses["loss_coherence"] = torch.tensor(0.0, device=deltas.device)

        # (2) ORTOGONALIDADE — Gram-Schmidt suave
        if cfg.enable_orthogonality and cfg.lam_orth > 0:
            # Normaliza deltas para o cálculo de G (estabilidade)
            d_norm = deltas.norm(dim=-1, keepdim=True).clamp(min=eps)
            d_hat = deltas / d_norm  # (batch, n_active, output_dim) — unitários
            # G = D̂ D̂ᵀ ∈ ℝ^{n_active × n_active} (matriz de Gram por batch)
            G = torch.matmul(d_hat, d_hat.transpose(1, 2))  # (batch, n_active, n_active)
            # Off-diagonal: penaliza correlações
            I = torch.eye(G.shape[1], device=G.device).unsqueeze(0)
            off_diag = G - I
            l_orth_raw = (off_diag * off_diag).sum(dim=(-1, -2)).mean()
            losses["loss_orthogonality"] = cfg.lam_orth * l_orth_raw
        else:
            losses["loss_orthogonality"] = torch.tensor(0.0, device=deltas.device)

        # (3) DIVERSIDADE — entropia de π (entropia bonus)
        if cfg.enable_diversity and cfg.lam_div > 0:
            # H(π) = -Σ π log π  → maximizar H = minimizar -H
            log_pi = torch.log(pi.clamp(min=eps))
            entropy = -(pi * log_pi).sum(dim=-1)  # (batch,)
            # log(N) é a entropia máxima (distribuição uniforme)
            log_n = math.log(max(self.active_count, 1))
            # Normaliza: l_div ∈ [0, 1] (0 = uniforme, 1 = colapsada)
            l_div_raw = (1.0 - entropy / max(log_n, eps)).mean()
            losses["loss_diversity"] = cfg.lam_div * l_div_raw
        else:
            losses["loss_diversity"] = torch.tensor(0.0, device=deltas.device)

        # (4) L2 — norma Frobenius de D
        if cfg.enable_l2 and cfg.lam_l2 > 0:
            l_l2_raw = (deltas * deltas).sum()
            losses["loss_l2"] = cfg.lam_l2 * l_l2_raw
        else:
            losses["loss_l2"] = torch.tensor(0.0, device=deltas.device)

        # Total
        losses["loss_synergy_total"] = (
            losses["loss_coherence"]
            + losses["loss_orthogonality"]
            + losses["loss_diversity"]
            + losses["loss_l2"]
        )

        return losses

    # ------------------------------------------------------------------
    # API utilitária para o trainer
    # ------------------------------------------------------------------
    def forward_stacked(self, x: torch.Tensor) -> torch.Tensor:
        """Compat com HypothesisEnsemble.forward_stacked (legado).

        Retorna (batch, n_active, output_dim) — apenas os deltas, sem π.
        """
        out = self.forward(x, return_components=False, stop_grad_x=False)
        return out["deltas"]

    def get_delta_ensemble(
        self, x: torch.Tensor, stop_grad_x: bool = False
    ) -> torch.Tensor:
        """Atalho para obter apenas o delta cooperativo."""
        return self.forward(
            x, return_components=False, stop_grad_x=stop_grad_x
        )["delta_ensemble"]

    def get_last_stats(self) -> Dict[str, Any]:
        """Retorna estatísticas do último forward (para monitoramento)."""
        pi = self._last_pi
        deltas = self._last_deltas
        stats: Dict[str, Any] = {
            "forward_count": int(self._forward_count),
            "active_count": int(self.active_count),
            "delta_scale": float(self.delta_scale.item()),
            "losses": dict(self._last_losses),
        }
        if pi is not None:
            # Entropia normalizada de π (1.0 = uniforme, 0.0 = colapsada)
            eps = self.config.eps
            log_pi = torch.log(pi.clamp(min=eps))
            entropy = -(pi * log_pi).sum(dim=-1).mean().item()
            log_n = math.log(max(self.active_count, 1))
            stats["pi_entropy"] = float(entropy)
            stats["pi_normalized_entropy"] = float(entropy / max(log_n, eps))
            stats["pi_max"] = float(pi.max(dim=-1).values.mean().item())
            stats["pi_min"] = float(pi.min(dim=-1).values.mean().item())
        if deltas is not None:
            # Norma média dos deltas
            stats["delta_norm_mean"] = float(
                deltas.norm(dim=-1).mean().item()
            )
            # Diversidade: 1 - cos(δ_i, δ_j) médio
            if deltas.shape[1] >= 2:
                d_norm = deltas.norm(dim=-1, keepdim=True).clamp(min=eps)
                d_hat = deltas / d_norm
                G = torch.matmul(d_hat, d_hat.transpose(1, 2))
                n = G.shape[1]
                # Off-diagonal mean
                mask = ~torch.eye(n, dtype=torch.bool, device=G.device)
                off_mean = float(G[:, mask].mean().item())
                stats["delta_correlation_mean"] = off_mean  # ideal: ~0
        return stats


# ============================================================================
# 7. HypT — Wrapper legado (compatibilidade reversa com V6.4)
# ============================================================================
class HypT(nn.Module):
    """HypT (V6.4 compat) — wrapper para SynergyHypothesisEnsemble em modo
    single-hypothesis.

    Mantém a API original para scripts legados:
        >>> hyp = HypT(d_input=256, d_model=256, nhead=4, d_ff=512,
        ...            output_dim=256, num_layers=2, dropout=0.1)
        >>> delta = hyp(o, stop_grad=True)  # → (batch, output_dim)

    Internamente, delega para SynergyHypothesisEnsemble com n_hypotheses=1
    (modo single-hyp, sem regularizações de sinergia — apenas L2 leve).

    Args:
        d_input: dimensão de entrada (saída do OrqCell = d_cache).
        d_model: dimensão interna (mantido para compat, mapeado para hidden_dim).
        nhead: nº de cabeças do Transformer (mantido para compat).
        d_ff: dimensão do FFN (mantido para compat, ignorado no modo MLP).
        output_dim: dimensão de saída (igual ao TrainT).
        num_layers: nº de camadas (mantido para compat).
        dropout: probabilidade de dropout.

    Forward:
        o: (batch, d_input), stop_grad: bool → delta: (batch, output_dim)
    """

    def __init__(
        self,
        d_input: int,
        d_model: int,
        nhead: int,
        d_ff: int,
        output_dim: int,
        num_layers: int = 2,
        dropout: float = 0.1,
    ):
        super().__init__()
        # Config single-hyp: sem sinergia (apenas L2 leve para estabilidade)
        cfg = SynergyConfig(
            lam_coh=0.0,
            lam_orth=0.0,
            lam_div=0.0,
            lam_l2=1.0e-4,
            enable_coherence=False,
            enable_orthogonality=False,
            enable_diversity=False,
            enable_l2=True,
            delta_scale_init=1.0,  # scale=1 para preservar magnitude V6.4
        )
        self._ensemble = SynergyHypothesisEnsemble(
            input_dim=int(d_input),
            output_dim=int(output_dim),
            n_hypotheses=1,
            max_n_hypotheses=1,
            hidden_dim=int(d_model) if d_model else HYP_HIDDEN_DIM_CANONICAL,
            config=cfg,
            use_transformer=True,  # Mantém TransformerEncoder para compat V6.4
        )
        # Guarda parâmetros originais para introspecção
        self.d_input = int(d_input)
        self.d_model = int(d_model)
        self.output_dim = int(output_dim)
        self.nhead = int(nhead)
        self.d_ff = int(d_ff)
        self.num_layers = int(num_layers)
        self.dropout = float(dropout)

    def forward(
        self, o: torch.Tensor, stop_grad: bool = True
    ) -> torch.Tensor:
        """o: (batch, d_input) → delta: (batch, output_dim).

        V6.5-V3: delega para SynergyHypothesisEnsemble (single-hyp mode).
        """
        out = self._ensemble.forward(
            o, return_components=False, stop_grad_x=stop_grad
        )
        # delta_ensemble = Σ_h π_h · δ_h, com n_hyp=1 → π_1=1, δ_ens=δ_1
        return out["delta_ensemble"]

    # ------------------------------------------------------------------
    # Exposição de estado interno (para inspeção/depuração)
    # ------------------------------------------------------------------
    def get_ensemble(self) -> SynergyHypothesisEnsemble:
        """Retorna o ensemble interno (para upgrade sinérgico posterior)."""
        return self._ensemble


# ============================================================================
# 8. Factory helpers
# ============================================================================
def create_synergy_ensemble(
    input_dim: int,
    output_dim: int,
    n_hypotheses: int = N_HYPOTHESES_CANONICAL,
    max_n_hypotheses: int = MAX_N_HYPOTHESES_CANONICAL,
    hidden_dim: int = HYP_HIDDEN_DIM_CANONICAL,
    config: Optional[SynergyConfig] = None,
    use_transformer: bool = False,
    device: Optional[torch.device] = None,
    seed: Optional[int] = 42,
) -> SynergyHypothesisEnsemble:
    """Factory para criar um SynergyHypothesisEnsemble com defaults canônicos.

    Constantes canônicas (NUNCA alterar):
        N_HYPOTHESES = 16, MAX_N_HYPOTHESES = 32,
        HYP_TRAIN_STEPS = 30, HYP_HIDDEN_DIM = 256.
    """
    return SynergyHypothesisEnsemble(
        input_dim=input_dim,
        output_dim=output_dim,
        n_hypotheses=n_hypotheses,
        max_n_hypotheses=max_n_hypotheses,
        hidden_dim=hidden_dim,
        config=config,
        use_transformer=use_transformer,
        device=device,
        seed=seed,
    )


# ============================================================================
# 9. Self-test (executável: python -m bigru_t.model.hyp_t)
# ============================================================================
def _self_test() -> Dict[str, Any]:
    """Self-test do módulo: cria ensemble, faz forward, verifica saídas."""
    print("[hyp_t self-test] Starting...")
    torch.manual_seed(42)

    # Parâmetros canônicos
    P_som = 256  # grid (4,4,4,4) flatten
    output_dim = 256 * 4  # I*J*K*L*4 = 256*4 = 1024
    n_hyp = N_HYPOTHESES_CANONICAL
    max_n = MAX_N_HYPOTHESES_CANONICAL
    hidden = HYP_HIDDEN_DIM_CANONICAL

    # Cria ensemble
    ens = create_synergy_ensemble(
        input_dim=P_som,
        output_dim=output_dim,
        n_hypotheses=n_hyp,
        max_n_hypotheses=max_n,
        hidden_dim=hidden,
        use_transformer=False,  # MLP para self-test rápido
        device=torch.device("cpu"),
        seed=42,
    )
    print(f"  Ensemble criado: {n_hyp} hipóteses (max {max_n}), "
          f"hidden={hidden}, P_som={P_som}, output_dim={output_dim}")

    # Forward com batch de 8 amostras
    x = torch.randn(8, P_som)
    out = ens(x)
    delta_ens = out["delta_ensemble"]
    deltas = out["deltas"]
    pi = out["pi"]
    losses = out["losses"]

    # Verifica shapes
    assert delta_ens.shape == (8, output_dim), \
        f"delta_ensemble shape: {delta_ens.shape} != (8, {output_dim})"
    assert deltas.shape == (8, n_hyp, output_dim), \
        f"deltas shape: {deltas.shape} != (8, {n_hyp}, {output_dim})"
    assert pi.shape == (8, n_hyp), \
        f"pi shape: {pi.shape} != (8, {n_hyp})"

    # Verifica que π soma 1
    pi_sum = pi.sum(dim=-1)
    assert torch.allclose(pi_sum, torch.ones(8), atol=1e-5), \
        f"π não soma 1: {pi_sum}"

    # Verifica losses
    print(f"  Losses: " + ", ".join(
        f"{k}={v.item():.6f}" for k, v in losses.items()
    ))
    assert torch.isfinite(losses["loss_synergy_total"]), \
        "loss_synergy_total não é finita"

    # Stats
    stats = ens.get_last_stats()
    print(f"  Stats: pi_entropy={stats.get('pi_entropy', 0):.4f} "
          f"(norm={stats.get('pi_normalized_entropy', 0):.4f}), "
          f"delta_norm_mean={stats.get('delta_norm_mean', 0):.6f}, "
          f"corr={stats.get('delta_correlation_mean', 0):.4f}")

    # Test HypT legacy wrapper
    hyp = HypT(
        d_input=P_som, d_model=hidden, nhead=4, d_ff=hidden * 2,
        output_dim=output_dim, num_layers=1, dropout=0.1,
    )
    delta_legacy = hyp(x, stop_grad=True)
    assert delta_legacy.shape == (8, output_dim), \
        f"HypT legacy delta shape: {delta_legacy.shape}"
    assert torch.isfinite(delta_legacy).all(), "HypT legacy delta não finito"
    print(f"  HypT legacy OK: delta shape={delta_legacy.shape}, "
          f"norm_mean={delta_legacy.norm(dim=-1).mean().item():.6f}")

    print("[hyp_t self-test] ALL CHECKS PASSED ✓")
    return {
        "ensemble_n_hyp": n_hyp,
        "ensemble_max_n": max_n,
        "ensemble_hidden": hidden,
        "delta_ensemble_shape": list(delta_ens.shape),
        "deltas_shape": list(deltas.shape),
        "pi_shape": list(pi.shape),
        "pi_sum_mean": float(pi_sum.mean().item()),
        "losses": {k: float(v.item()) for k, v in losses.items()},
        "stats": stats,
        "hyp_legacy_shape": list(delta_legacy.shape),
    }


# ============================================================================
# __all__
# ============================================================================
__all__ = [
    # Config
    "SynergyConfig",
    # Classes principais
    "SynergyHypothesisEnsemble",
    "DeltaGenerator",
    "GatingNetwork",
    "HypT",  # legacy wrapper (V6.4 compat)
    # Constantes canônicas
    "N_HYPOTHESES_CANONICAL",
    "MAX_N_HYPOTHESES_CANONICAL",
    "HYP_TRAIN_STEPS_CANONICAL",
    "HYP_HIDDEN_DIM_CANONICAL",
    # Factory
    "create_synergy_ensemble",
    # k-means++ init
    "_kmeans_pp_init_delta_directions",
]


if __name__ == "__main__":
    _self_test()