src/bigru_t/training/hypothesis_monitor.py

"""hypothesis_monitor.py — Detector de ativação das hipóteses (L1-L4).

Filosofia: "se as hipóteses não são ativadas durante punição, há um bug".

Cada um dos 4 lemas tem uma condição de ativação observável:
  L1 (ModuleSelector): α = softmax(logits/T) — ativa quando há variação em α
      (entropy < log(K) - ε). Se entropy ≈ log(K), nenhum módulo é selecionado
      (todos com peso igual) — possível bug no temperature ou nos logits.
  L2 (GradientSurgery): ativa quando use_hyp=True e apply_gradient_surgery
      modifica os gradientes. Se ||g_hyp|| ≈ 0 ou projeção ortogonal ≈ 0,
      a cirurgia é no-op (possível bug: loss_hyp não tem gradiente).
  L3 (HypT): ativa quando loss > τ (modelo está "lutando"). Se loss > τ mas
      use_hyp=False (não foi ativada), é bug. Se use_hyp=True mas ||delta|| ≈ 0,
      a hipótese não contribui (possível bug: HypT morta ou stop_grad excessivo).
  L4 (MetaConfigurator): ativa a cada meta_interval passos. Se T e τ não mudam
      entre chamadas, o meta-otimizador está estagnado (possível bug: lr=0 ou
      gradientes None).

O monitor mantém contadores e dispara alertas quando:
  - Punishment_count > 0 mas activation_count = 0 (hipótese nunca ativada)
  - Activation_count > 0 mas delta_norm_medio < epsilon (HypT morta)
  - T ou τ não mudam em N chamadas do meta (estagnação)
  - Entropy do L1 saturada (> 0.95 * log(K)) — sem seleção real
"""
from __future__ import annotations

import logging
import math
from collections import deque
from dataclasses import dataclass, field
from typing import Optional

import torch

logger = logging.getLogger(__name__)


@dataclass
class LemmaState:
    """Snapshot do estado de ativação de cada lema em um step."""
    step: int
    # L1
    alpha_entropy: float = 0.0
    alpha_max: float = 0.0
    alpha_active_modules: int = 0
    temperature: float = 1.0
    l1_activated: bool = False  # True se entropy < 0.95 * log(K)
    # L2
    l2_applied: bool = False
    grad_main_norm: float = 0.0
    grad_hyp_norm: float = 0.0
    orth_proj_norm: float = 0.0
    # L3
    loss_main: float = 0.0
    tau: float = float("inf")
    use_hyp_activated: bool = False  # True se use_hyp=True
    delta_norm: float = 0.0  # ||delta||₂ — se 0, HypT morta
    punishment_triggered: bool = False  # True se loss > τ
    # L4
    T_new: float = 1.0
    tau_new: float = 1.0
    sharpness: float = 0.0
    l4_meta_called: bool = False
    l4_T_moved: bool = False
    l4_tau_moved: bool = False
    # Circular Reasoning (Teorema 4)
    crw_converged: bool = False
    crw_w2_final: float = 0.0
    crw_q_avg: float = 1.0
    crw_contraction_valid: bool = False


@dataclass
class HypothesisMonitorReport:
    """Relatório final do monitor."""
    total_steps: int = 0
    punishment_steps: int = 0       # steps com loss > τ
    hyp_activated_steps: int = 0    # steps com use_hyp=True
    hyp_activated_during_punishment: int = 0  # ⚠️ chave: punição E ativação
    delta_zero_when_activated: int = 0  # ⚠️ bug: HypT morta
    l1_entropy_saturated_steps: int = 0  # ⚠️ bug: sem seleção
    l2_zero_grad_hyp_steps: int = 0  # ⚠️ bug: loss_hyp sem grad
    l4_T_stagnation: int = 0  # ⚠️ bug: T não move
    l4_tau_stagnation: int = 0  # ⚠️ bug: τ não move
    crw_converged_steps: int = 0
    crw_contraction_invalid_steps: int = 0  # ⚠️ bug: W₂ cresceu
    last_T: float = 1.0
    last_tau: float = 1.0
    history: list = field(default_factory=list)

    @property
    def has_bug(self) -> bool:
        """True se algum bug provável foi detectado."""
        # Punição ocorreu mas hipótese nunca ativada — BUG CRÍTICO
        if self.punishment_steps > 0 and self.hyp_activated_during_punishment == 0:
            return True
        # Hipótese ativada mas delta sempre zero — HypT morta
        if self.hyp_activated_steps > 0 and self.hyp_activated_steps == self.delta_zero_when_activated:
            return True
        # L1 saturada em mais de 80% dos steps
        if self.total_steps > 5 and self.l1_entropy_saturated_steps > 0.8 * self.total_steps:
            return True
        # CRW com contração inválida em mais de 50% dos steps
        if self.crw_contraction_invalid_steps > 0.5 * max(1, self.total_steps):
            return True
        return False

    @property
    def bug_summary(self) -> list[str]:
        bugs = []
        if self.punishment_steps > 0 and self.hyp_activated_during_punishment == 0:
            bugs.append(
                f"CRÍTICO: {self.punishment_steps} passos de punição (loss>τ) "
                f"mas hipótese nunca ativada — verifique tau e use_hypothesis"
            )
        if self.hyp_activated_steps > 0 and self.hyp_activated_steps == self.delta_zero_when_activated:
            bugs.append(
                f"CRÍTICO: hipótese ativada {self.hyp_activated_steps}x mas "
                f"||delta|| sempre zero — HypT morta ou stop_grad excessivo"
            )
        if self.total_steps > 5 and self.l1_entropy_saturated_steps > 0.8 * self.total_steps:
            bugs.append(
                f"ALERTA: L1 entropy saturada em {self.l1_entropy_saturated_steps}/{self.total_steps} "
                f"passos — module_selector não está selecionando"
            )
        if self.l4_T_stagnation > 3:
            bugs.append(
                f"ALERTA: T não mudou em {self.l4_T_stagnation} chamadas meta — "
                f"meta_otim estagnado"
            )
        if self.crw_contraction_invalid_steps > 0.5 * max(1, self.total_steps):
            bugs.append(
                f"ALERTA: CRW contração inválida em "
                f"{self.crw_contraction_invalid_steps}/{self.total_steps} passos — "
                f"W₂ cresceu em vez de contrair"
            )
        return bugs


class HypothesisMonitor:
    """Monitora ativação dos 4 lemas + CircularReasoning durante o treino.

    Uso:
        monitor = HypothesisMonitor(max_history=100)
        # No loop de treino, após cada micro-batch:
        state = monitor.record(
            step=global_step,
            alpha=aux["alpha"],            # (K,) tensor
            entropy_reg=aux["entropy_reg"],
            loss_main=loss_main.item(),
            tau=model.tau.item(),
            use_hyp=use_hyp,
            delta=delta,
            temperature=meta_cfg.temperature,
            circular_info=aux.get("circular_info"),
        )
        # Após cada meta step:
        monitor.record_meta(T_new, tau_new, sharpness)
        # Ao final:
        report = monitor.summary()
        if report.has_bug:
            for bug in report.bug_summary:
                logger.error(bug)
    """

    def __init__(
        self,
        max_modules: int = 8,
        delta_epsilon: float = 1e-4,
        meta_stagnation_epsilon: float = 1e-4,
        history_size: int = 200,
    ):
        self.max_modules = max_modules
        self.delta_epsilon = delta_epsilon
        self.meta_stagnation_epsilon = meta_stagnation_epsilon
        self.history_size = history_size
        self._history: deque = deque(maxlen=history_size)
        # Acumuladores para relatório
        self._punishment_steps = 0
        self._hyp_activated_steps = 0
        self._hyp_activated_during_punishment = 0
        self._delta_zero_when_activated = 0
        self._l1_entropy_saturated_steps = 0
        self._l2_zero_grad_hyp_steps = 0
        self._crw_converged_steps = 0
        self._crw_contraction_invalid_steps = 0
        # Meta tracking (L4)
        self._last_T: float = 1.0
        self._last_tau: float = 1.0
        self._T_stagnation_count: int = 0
        self._tau_stagnation_count: int = 0
        self._meta_calls: int = 0
        # L2 tracking
        self._last_l2_applied: bool = False
        self._last_grad_hyp_norm: float = 0.0
        self._last_orth_proj_norm: float = 0.0

    def record(
        self,
        step: int,
        alpha: Optional[torch.Tensor] = None,
        entropy_reg: float = 0.0,
        loss_main: float = 0.0,
        tau: float = float("inf"),
        use_hyp: bool = False,
        delta: Optional[torch.Tensor] = None,
        temperature: float = 1.0,
        circular_info: Optional[dict] = None,
        l2_applied: bool = False,
        grad_main_norm: float = 0.0,
        grad_hyp_norm: float = 0.0,
        orth_proj_norm: float = 0.0,
    ) -> LemmaState:
        """Registra estado de um step. Retorna LemmaState para inspeção imediata."""
        # L1: entropy e max(alpha)
        if alpha is not None:
            alpha = alpha.detach().float()
            # alpha pode ser (K,) ou (B, K) — reduzir para (K,)
            if alpha.dim() > 1:
                alpha = alpha.mean(dim=0)
            alpha_entropy = float(-(alpha * torch.log(alpha + 1e-10)).sum().item())
            alpha_max = float(alpha.max().item())
            alpha_active = int((alpha > 0.01).sum().item())
            # L1 ativa se entropy < 0.95 * log(K) (seleção real está acontecendo)
            max_entropy = math.log(self.max_modules)
            l1_activated = alpha_entropy < 0.95 * max_entropy
            if not l1_activated:
                self._l1_entropy_saturated_steps += 1
        else:
            alpha_entropy = 0.0
            alpha_max = 0.0
            alpha_active = 0
            l1_activated = False

        # L3: punição e delta
        punishment_triggered = loss_main > tau
        if punishment_triggered:
            self._punishment_steps += 1
        if use_hyp:
            self._hyp_activated_steps += 1
            if punishment_triggered:
                self._hyp_activated_during_punishment += 1
            # Verifica se delta é significativo
            if delta is not None:
                delta_norm = float(delta.detach().float().norm().item())
                if delta_norm < self.delta_epsilon:
                    self._delta_zero_when_activated += 1
            else:
                delta_norm = 0.0
                self._delta_zero_when_activated += 1
        else:
            delta_norm = 0.0

        # L2 tracking (será preenchido por record_l2 após backward)
        self._last_l2_applied = l2_applied
        self._last_grad_hyp_norm = grad_hyp_norm
        self._last_orth_proj_norm = orth_proj_norm
        if l2_applied and grad_hyp_norm < 1e-10:
            self._l2_zero_grad_hyp_steps += 1

        # Circular Reasoning
        crw_converged = False
        crw_w2_final = 0.0
        crw_q_avg = 1.0
        crw_contraction_valid = True
        if circular_info is not None:
            crw_converged = bool(circular_info.get("converged", False))
            crw_w2_final = float(circular_info.get("w2_final", 0.0))
            crw_q_avg = float(circular_info.get("q_empirical_avg", 1.0))
            crw_contraction_valid = bool(circular_info.get("contraction_valid", True))
            if crw_converged:
                self._crw_converged_steps += 1
            if not crw_contraction_valid:
                self._crw_contraction_invalid_steps += 1

        state = LemmaState(
            step=step,
            alpha_entropy=alpha_entropy,
            alpha_max=alpha_max,
            alpha_active_modules=alpha_active,
            temperature=temperature,
            l1_activated=l1_activated,
            l2_applied=l2_applied,
            grad_main_norm=grad_main_norm,
            grad_hyp_norm=grad_hyp_norm,
            orth_proj_norm=orth_proj_norm,
            loss_main=loss_main,
            tau=tau,
            use_hyp_activated=use_hyp,
            delta_norm=delta_norm,
            punishment_triggered=punishment_triggered,
            T_new=self._last_T,
            tau_new=self._last_tau,
            l4_meta_called=False,
            crw_converged=crw_converged,
            crw_w2_final=crw_w2_final,
            crw_q_avg=crw_q_avg,
            crw_contraction_valid=crw_contraction_valid,
        )
        self._history.append(state)
        return state

    def record_meta(
        self,
        T_new: float,
        tau_new: float,
        sharpness: float = 0.0,
    ) -> None:
        """Registra resultado de um passo do MetaConfigurator (L4)."""
        self._meta_calls += 1
        # Detecta estagnação
        if abs(T_new - self._last_T) < self.meta_stagnation_epsilon:
            self._T_stagnation_count += 1
        else:
            self._T_stagnation_count = 0
        if abs(tau_new - self._last_tau) < self.meta_stagnation_epsilon:
            self._tau_stagnation_count += 1
        else:
            self._tau_stagnation_count = 0
        self._last_T = T_new
        self._last_tau = tau_new

    def summary(self) -> HypothesisMonitorReport:
        """Gera relatório final com detecção de bugs."""
        report = HypothesisMonitorReport(
            total_steps=len(self._history),
            punishment_steps=self._punishment_steps,
            hyp_activated_steps=self._hyp_activated_steps,
            hyp_activated_during_punishment=self._hyp_activated_during_punishment,
            delta_zero_when_activated=self._delta_zero_when_activated,
            l1_entropy_saturated_steps=self._l1_entropy_saturated_steps,
            l2_zero_grad_hyp_steps=self._l2_zero_grad_hyp_steps,
            l4_T_stagnation=self._T_stagnation_count,
            l4_tau_stagnation=self._tau_stagnation_count,
            crw_converged_steps=self._crw_converged_steps,
            crw_contraction_invalid_steps=self._crw_contraction_invalid_steps,
            last_T=self._last_T,
            last_tau=self._last_tau,
            history=[s.__dict__ for s in list(self._history)[-20:]],  # últimos 20
        )
        return report

    def log_summary(self) -> None:
        """Imprime relatório no log."""
        r = self.summary()
        logger.info("=" * 60)
        logger.info("HIPÓTESE MONITOR — RELATÓRIO DE ATIVAÇÃO")
        logger.info("=" * 60)
        logger.info(f"  Total steps:                       {r.total_steps}")
        logger.info(f"  Punição (loss > τ):                {r.punishment_steps}")
        logger.info(f"  Hipótese ativada (use_hyp=True):   {r.hyp_activated_steps}")
        logger.info(f"  Hipótese ativada DURANTE punição:  {r.hyp_activated_during_punishment}")
        logger.info(f"  ||delta||≈0 quando ativada:        {r.delta_zero_when_activated}")
        logger.info(f"  L1 entropy saturada:               {r.l1_entropy_saturated_steps}")
        logger.info(f"  L2 grad_hyp≈0:                     {r.l2_zero_grad_hyp_steps}")
        logger.info(f"  L4 T estagnado:                    {r.l4_T_stagnation}")
        logger.info(f"  L4 τ estagnado:                    {r.l4_tau_stagnation}")
        logger.info(f"  CRW convergiu:                     {r.crw_converged_steps}")
        logger.info(f"  CRW contração inválida:            {r.crw_contraction_invalid_steps}")
        logger.info(f"  T final:                           {r.last_T:.4f}")
        logger.info(f"  τ final:                           {r.last_tau:.4f}")
        if r.has_bug:
            logger.error("=" * 60)
            logger.error("⚠️  BUGS DETECTADOS:")
            for bug in r.bug_summary:
                logger.error(f"  - {bug}")
            logger.error("=" * 60)
        else:
            logger.info("  ✓ Nenhum bug crítico detectado pelo monitor")
        logger.info("=" * 60)


__all__ = ["HypothesisMonitor", "HypothesisMonitorReport", "LemmaState"]