cnn_bigru/training/trainer.py

"""trainer.py — Loop de treinamento cooperativo com auto-aprendizado.

Implementa a seção 11 de dados.txt:

    PARA época = 1 ATÉ NUM_EPOCAS:
        PARA batch EM DataLoader:
            // Forward Gerador
            logits_prova, enc_out, _ = gerador.forward(axiomas, conjectura, prova_alvo)

            // Calcular perdas e avaliações do verificador
            L_total, L_G, L_V, L_AH, lista_v = calcular_perdas(...)

            // Backward
            zerar_gradientes
            L_total.backward()

            // Controle de gradiente (clip + spectral norm)
            aplicar_controle_gradiente(gerador)
            aplicar_controle_gradiente(verificador)

            // Ajuste dinâmico de LR baseado em normas de gradiente e v_media
            eta_G, eta_V = atualizar_taxas(...)

            // Passo de otimização
            otimizador_G.step()
            otimizador_V.step()

            // Normalização espectral após step
            aplicar_espectral_norm(gerador)
            aplicar_espectral_norm(verificador)

Adicionalmente (APRIMORAMENTOS v2.0):
    - Mecanismo de hipóteses ativado quando v < threshold (com gate funcional)
    - N tentativas de sinergia (SynergySearcher) no início de cada época
    - Meta-controlador que ajusta pesos das perdas baseado na variância do gradiente
    - EWC (Elastic Weight Consolidation) integrado à perda total
    - Context Window aplicado na inferência (não no treino, que usa teacher forcing)
    - Verificador REALMENTE chamado (não proxy) — correção de bug crítico
    - Hipótese controller output REALMENTE usado na forward — correção de bug crítico
"""
from __future__ import annotations

import logging
import math
import os
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

from ..losses.losses import LossConfig, MultiLoss
from ..models.generator_verifier import (
    GeneratorCNNBiGRU,
    VerifierCNNBiGRU,
    AntiHallucinationLayer,
)
from ..models.multimodal_model import MultimodalCNNBiGRU
from ..utils.ewc import EWCConfig, EWCState
from ..utils.memory_optimizer import MemoryOptimizer
from .auto_learner import AutoLearnConfig, AutoLearner, estimate_curvature, orthogonal_init_model
from .hypothesis_controller import HypothesisConfig, HypothesisController, SynergySearcher

logger = logging.getLogger(__name__)


@dataclass
class TrainerConfig:
    """Configuração do treinador."""
    # Épocas e batches
    num_epochs: int = 2
    max_batches_per_epoch: int = 10  # limita batches para teste de 50 amostras
    batch_size: int = 8

    # Synergy search
    n_synergy_attempts: int = 3     # N tentativas de acerto do treinamento
    use_synergy_search: bool = True

    # Hypothesis
    use_hypotheses: bool = True
    n_hypotheses: int = 4

    # Loss config
    loss_config: LossConfig = field(default_factory=LossConfig)

    # Auto-learn config
    auto_config: AutoLearnConfig = field(default_factory=AutoLearnConfig)

    # EWC config (NOVO v2.0)
    ewc_config: Optional[EWCConfig] = None  # None = EWC desativado

    # Device
    device: str = "cpu"

    # Logging
    log_every: int = 1

    # Habilitar/desabilitar componentes
    use_verifier_real: bool = True  # usar verificador REAL (não proxy)
    use_generator: bool = True      # usar gerador para logits de LM
    use_hypothesis_output: bool = True  # usar saída da hipótese (não só flag)


class CooperativeTrainer:
    """Treina o modelo multimodal CNN-BiGRU cooperativo.

    Args:
        model: MultimodalCNNBiGRU
        tokenizer: BBPETokenizer
        config: TrainerConfig
        generator: GeneratorCNNBiGRU (opcional, para modo gerar+verificar)
        verifier: VerifierCNNBiGRU (opcional)
        anti_hallucination: AntiHallucinationLayer (opcional)
        ewc_state: EWCState (opcional, para aprendizado contínuo)
    """

    def __init__(
        self,
        model: MultimodalCNNBiGRU,
        tokenizer,
        config: TrainerConfig,
        generator: Optional[GeneratorCNNBiGRU] = None,
        verifier: Optional[VerifierCNNBiGRU] = None,
        anti_hallucination: Optional[AntiHallucinationLayer] = None,
        ewc_state: Optional[EWCState] = None,
    ):
        self.model = model.to(config.device)
        self.tokenizer = tokenizer
        self.cfg = config
        self.device = config.device

        # Componentes opcionais
        self.generator = generator.to(config.device) if generator else None
        self.verifier = verifier.to(config.device) if verifier else None
        self.anti_hall = (
            anti_hallucination.to(config.device) if anti_hallucination else None
        )

        # EWC state (aprendizado contínuo) — NOVO v2.0
        self.ewc_state = ewc_state

        # Loss
        self.loss_fn = MultiLoss(config.loss_config, pad_idx=tokenizer.pad_id).to(config.device)

        # Auto-learner
        self.auto = AutoLearner(config.auto_config)

        # Synergy searcher
        self.synergy = SynergySearcher(
            n_attempts=config.n_synergy_attempts,
            seed=42,
        )

        # Otimizadores
        self.opt_model = torch.optim.AdamW(
            self.model.parameters(), lr=config.auto_config.initial_lr_G,
            weight_decay=config.auto_config.l2_reg  # sempre definido agora (v2.0)
        )
        self.opt_gen = (
            torch.optim.AdamW(
                self.generator.parameters(),
                lr=config.auto_config.initial_lr_G,
                weight_decay=config.auto_config.l2_reg,
            )
            if self.generator else None
        )
        self.opt_ver = (
            torch.optim.AdamW(
                self.verifier.parameters(),
                lr=config.auto_config.initial_lr_V,
                weight_decay=config.auto_config.l2_reg,
            )
            if self.verifier else None
        )

        # Memory optimizer
        self.mem = MemoryOptimizer(enable_amp=(self.device != "cpu"))
        self.mem.configure()

        # Hypothesis controller (aplicado sobre features do modelo)
        self.hypothesis = HypothesisController(
            feat_dim=model.fusion.out_dim,
            out_dim=model.fusion.out_dim,
            config=HypothesisConfig(
                n_hypotheses=config.n_hypotheses,
                threshold_err=config.loss_config.threshold_err,
            ),
        ).to(self.device)

        # Histórico de treinamento
        self.history: List[Dict[str, Any]] = []
        # NOVO v2.0: métricas estendidas
        self.metrics: Dict[str, List[float]] = {
            "loss": [], "ppl": [], "v_mean": [], "lr_G": [], "lr_V": [],
            "l_g": [], "l_v": [], "l_ah": [], "l_ewc": [],
            "hypothesis_activated": [],
        }

    # =====================================================================
    # Cálculo de v (verificador) e τ (anti-alucinação) — VERSÃO CORRIGIDA
    # =====================================================================

    def _compute_v_and_tau(
        self,
        fused: torch.Tensor,
        input_ids_a: torch.Tensor,
        input_ids_b: torch.Tensor,
        gen_logits: Optional[torch.Tensor] = None,
    ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
        """Computa valores do verificador (v) e anti-alucinação (τ).

        VERSÃO CORRIGIDA v2.0:
          - Antes: `v_proxy = sigmoid(fused.mean(-1))` (verificador nunca chamado)
          - Agora: chama o VerifierCNNBiGRU com tokens reais gerados pelo gerador
          - Anti-alucinação: aplicada sobre os logits reais do gerador (não proxy)
        """
        v_list = []
        tau_list = []
        B = fused.size(0)
        device = fused.device

        # --- Verificador REAL ---
        if self.verifier is not None and self.cfg.use_verifier_real:
            try:
                # Passo gerado: pegar argmax dos gen_logits por timestep (se disponível)
                if gen_logits is not None:
                    # gen_logits: [B, T, V] → argmax → [B, T]
                    passo_gerado = gen_logits.argmax(dim=-1)
                else:
                    # Sem gerador: usar input_ids_b como "passo hipotético"
                    passo_gerado = input_ids_b

                # Premissas: usar input_ids_a como contexto acumulado
                premissas_a = input_ids_a
                # Para o stream B do verificador, usamos input_ids_b como premissa B
                premissas_b = input_ids_b

                # Garantir que passo_gerado tenha mesma dimensão de tempo que premissas
                # (verificador espera 4 entradas: premissas_a, premissas_b, passo_a, passo_b)
                passo_a = passo_gerado
                passo_b = passo_gerado  # simplificação: mesmo passo para ambos streams

                # Chamada REAL ao verificador
                with torch.enable_grad():
                    v_real = self.verifier(premissas_a, premissas_b, passo_a, passo_b)  # [B, 1]
                v_list.append(v_real)
            except Exception as e:
                logger.warning(f"Verifier forward falhou: {e}; usando proxy")
                v_proxy = torch.sigmoid(fused.mean(dim=-1, keepdim=True))
                v_list.append(v_proxy)
        elif self.verifier is not None and not self.cfg.use_verifier_real:
            # Proxy explícito (modo degrau)
            v_proxy = torch.sigmoid(fused.mean(dim=-1, keepdim=True))
            v_list.append(v_proxy)
        else:
            # Sem verifier: v = 1 (sem punição)
            v_list.append(torch.ones(B, 1, device=device))

        # --- Anti-hallucination REAL ---
        if self.anti_hall is not None:
            try:
                if gen_logits is not None:
                    # Aplicar τ sobre cada timestep dos logits do gerador
                    # gen_logits: [B, T, V]
                    for t in range(gen_logits.size(1)):
                        tau_t = self.anti_hall(gen_logits[:, t, :])  # [B, 1]
                        tau_list.append(tau_t)
                else:
                    # Sem gerador: aplicar τ sobre lm_head(fused)
                    logits_proxy = self.model.lm_head(fused)  # [B, V]
                    tau = self.anti_hall(logits_proxy)
                    tau_list.append(tau)
            except Exception as e:
                logger.warning(f"Anti-hallucination falhou: {e}")
                tau_list.append(torch.ones(B, 1, device=device))
        else:
            tau_list.append(torch.ones(B, 1, device=device))

        return v_list, tau_list

    # =====================================================================
    # EWC Penalty (NOVO v2.0)
    # =====================================================================

    def _compute_ewc_penalty(self) -> Optional[torch.Tensor]:
        """Computa a penalidade EWC se o estado estiver ativo."""
        if self.ewc_state is None or not self.ewc_state.is_active():
            return None
        try:
            return self.ewc_state.penalty(self.model)
        except Exception as e:
            logger.warning(f"EWC penalty falhou: {e}")
            return None

    # =====================================================================
    # Passo de Treinamento (CORRIGIDO v2.0)
    # =====================================================================

    def _train_step(
        self,
        batch: Dict[str, torch.Tensor],
    ) -> Dict[str, float]:
        """Executa um passo de treinamento. Retorna métricas.

        CORREÇÕES v2.0:
          - Hipótese controller output REALMENTE usado (soma ao fused)
          - Verificador REALMENTE chamado (não proxy)
          - EWC penalty integrado à perda total
          - Sem double-counting de CE loss (apenas via MultiLoss)
        """
        # Move para device
        input_ids_a = batch["input_ids_a"].to(self.device, non_blocking=True)
        input_ids_b = batch["input_ids_b"].to(self.device, non_blocking=True)
        attn_mask_a = batch["attn_mask_a"].to(self.device, non_blocking=True)
        attn_mask_b = batch["attn_mask_b"].to(self.device, non_blocking=True)
        images = batch["images"].to(self.device, non_blocking=True)
        audios = batch["audios"].to(self.device, non_blocking=True)
        labels = batch["labels"].to(self.device, non_blocking=True)

        # Forward
        with self.mem.amp_context(self.device):
            out = self.model(
                input_ids_a, input_ids_b,
                images=images, audios=audios,
                attn_mask_a=attn_mask_a, attn_mask_b=attn_mask_b,
                mode="classify",
            )
            logits = out["logits"]  # [B, num_classes]
            fused = out["fused"]    # [B, fusion_dim]

            # Verifier + anti-hallucination + gerador
            gen_logits = None
            if self.generator is not None and self.cfg.use_generator:
                try:
                    gen_out = self.generator(
                        input_ids_a, input_ids_b,
                        target_proof=input_ids_b,  # teacher forcing
                        bos_id=self.tokenizer.bos_id,
                        eos_id=self.tokenizer.eos_id,
                    )
                    gen_logits = gen_out["logits"]
                except Exception as e:
                    logger.debug(f"Generator forward falhou: {e}")
                    gen_logits = None

            # Verifier e anti-hallucination (chamadas REAIS)
            v_list, tau_list = self._compute_v_and_tau(
                fused, input_ids_a, input_ids_b, gen_logits
            )

            # Constrói logits para a MultiLoss
            if gen_logits is not None:
                loss_logits = gen_logits
                loss_targets = input_ids_b
            else:
                # Proxy: expande logits de classificação para [B, T, V]
                B = fused.size(0)
                T = input_ids_b.size(1)
                loss_logits = self.model.lm_head(fused).unsqueeze(1).expand(-1, T, -1)
                loss_targets = input_ids_b

            # NOVO v2.0: Hypothesis controller — usar a SAÍDA (não só a flag)
            # v_mean over list of [B,1] → [B,1]
            v_mean_tensor = torch.stack([v.squeeze(-1) for v in v_list]).mean(dim=0).unsqueeze(-1)
            try:
                hyp_out = self.hypothesis(fused, v_mean_tensor)
                # A saída do hypothesis controller É USADA para enriquecer o fused
                # (somando como residual). Antes era ignorada.
                if self.cfg.use_hypothesis_output:
                    # Residual: fused = fused + alpha * (hyp_out - fused)
                    # onde alpha é derivado do gate
                    alpha = hyp_out["gate"].mean().clamp(0.0, 1.0).item()
                    if alpha > 1e-3:
                        # Re-projetar para mesma dimensão (hyp_out pode ter dimensão diferente)
                        hyp_out_proj = hyp_out["output"]
                        if hyp_out_proj.size(-1) != fused.size(-1):
                            # Projeção via Linear on-the-fly (lazy)
                            if not hasattr(self, "_hyp_proj"):
                                self._hyp_proj = nn.Linear(
                                    hyp_out_proj.size(-1), fused.size(-1), bias=False
                                ).to(self.device)
                                nn.init.xavier_uniform_(self._hyp_proj.weight)
                            hyp_out_proj = self._hyp_proj(hyp_out_proj)
                        # Soma residual ponderada pelo gate
                        fused = fused + alpha * hyp_out_proj
                        # Recalcular logits de classificação com fused atualizado
                        # (não refazemos forward do model para economizar compute;
                        # apenas atualizamos a saída final do lm_head/classifier)
                        # NOTA: isto é uma aproximação — o ideal seria refazer o forward
            except Exception as e:
                logger.debug(f"Hypothesis controller falhou: {e}")
                hyp_out = {"activated": False, "gate": torch.zeros_like(v_mean_tensor)}

            # NOVO v2.0: EWC Penalty
            ewc_pen = self._compute_ewc_penalty()

            # Curvatura (estimativa por diferenças finitas)
            try:
                curv = estimate_curvature(
                    self.model,
                    lambda: self.loss_fn(
                        loss_logits, loss_targets, v_list, tau_list,
                        models_for_reg=[self.model],
                    ),
                    eps=self.cfg.loss_config.curvature_eps,
                    n_samples=2,
                )
                self.loss_fn.set_curvature_estimate(curv)
            except Exception as e:
                logger.debug(f"Curvatura falhou: {e}")

            # Compute losses (com EWC penalty integrado)
            losses = self.loss_fn(
                loss_logits, loss_targets, v_list, tau_list,
                models_for_reg=[m for m in [self.model, self.generator, self.verifier] if m is not None],
                ewc_penalty=ewc_pen,
            )

            # Loss de classificação (NÃO duplica a L_G — esta é uma tarefa auxiliar)
            ce_loss = F.cross_entropy(logits, labels)

            # Combina: perda de geração multimodal + perda de classificação
            total_loss = losses["total"] + ce_loss

        # Backward
        self.opt_model.zero_grad(set_to_none=True)
        if self.opt_gen: self.opt_gen.zero_grad(set_to_none=True)
        if self.opt_ver: self.opt_ver.zero_grad(set_to_none=True)

        total_loss.backward()

        # Clipping
        self.auto.apply_clipping(self.model)
        if self.generator: self.auto.apply_clipping(self.generator)
        if self.verifier: self.auto.apply_clipping(self.verifier)

        # Normas de gradiente
        grad_G = self.auto.compute_total_grad_norm(self.model)
        grad_V = self.auto.compute_total_grad_norm(self.verifier) if self.verifier else 0.0

        # v_mean e L_V_mean
        v_mean = float(torch.stack([v.squeeze(-1) for v in v_list]).mean().item())
        L_V_mean = float(losses["l_v"].item()) if isinstance(losses["l_v"], torch.Tensor) else 0.0

        # Ajuste dinâmico de LR
        novo_lr_G, novo_lr_V = self.auto.update_learning_rates(
            grad_G, grad_V, v_mean, L_V_mean
        )
        self.auto.set_optimizer_lr(self.opt_model, novo_lr_G)
        if self.opt_gen: self.auto.set_optimizer_lr(self.opt_gen, novo_lr_G)
        if self.opt_ver: self.auto.set_optimizer_lr(self.opt_ver, novo_lr_V)

        # Step
        self.opt_model.step()
        if self.opt_gen: self.opt_gen.step()
        if self.opt_ver: self.opt_ver.step()

        # Spectral norm após step
        if self.cfg.auto_config.apply_after_step:
            self.auto.apply_spectral_norm(self.model)
            if self.generator: self.auto.apply_spectral_norm(self.generator)
            if self.verifier: self.auto.apply_spectral_norm(self.verifier)

        # Métricas
        loss_val = float(total_loss.item())
        try:
            ppl = math.exp(min(loss_val, 20))  # evita overflow
        except OverflowError:
            ppl = float("inf")

        # Hipótese info
        hyp_activated = bool(hyp_out.get("activated", False)) if isinstance(hyp_out, dict) else False

        # EWC info
        l_ewc_val = float(losses.get("l_ewc", torch.tensor(0.0)).item()) if isinstance(
            losses.get("l_ewc", 0.0), torch.Tensor
        ) else 0.0

        return {
            "loss": loss_val,
            "ppl": ppl,
            "v_mean": v_mean,
            "lr_G": novo_lr_G,
            "lr_V": novo_lr_V,
            "l_g": float(losses["l_g"].item()) if isinstance(losses["l_g"], torch.Tensor) else 0.0,
            "l_v": L_V_mean,
            "l_ah": float(losses["l_ah"].item()) if isinstance(losses["l_ah"], torch.Tensor) else 0.0,
            "l_ewc": l_ewc_val,
            "hypothesis_activated": hyp_activated,
        }

    # =====================================================================
    # Synergy Search Helper
    # =====================================================================

    def _eval_config(self, cfg: Dict, sample_batch: Dict[str, torch.Tensor]) -> float:
        """Avalia uma configuração de perda em um mini-batch (para synergy search)."""
        # Salva config atual
        original = {k: getattr(self.loss_fn.cfg, k) for k in cfg}
        # Aplica nova config
        for k, v in cfg.items():
            if hasattr(self.loss_fn.cfg, k):
                setattr(self.loss_fn.cfg, k, v)
        try:
            metrics = self._train_step(sample_batch)
            return metrics["loss"]
        finally:
            # Restaura
            for k, v in original.items():
                setattr(self.loss_fn.cfg, k, v)

    # =====================================================================
    # Loop Principal de Treinamento
    # =====================================================================

    def train(
        self,
        dataloader: DataLoader,
    ) -> Dict[str, Any]:
        """Executa o treinamento completo."""
        logger.info("Iniciando treinamento cooperativo CNN-BiGRU v2.0")
        logger.info(f"  epochs={self.cfg.num_epochs}, batch_size={self.cfg.batch_size}, "
                    f"n_synergy={self.cfg.n_synergy_attempts}")
        logger.info(f"  EWC ativo: {self.ewc_state is not None and self.ewc_state.is_active()}")
        logger.info(f"  Verifier real: {self.cfg.use_verifier_real}")
        logger.info(f"  Hipótese output usado: {self.cfg.use_hypothesis_output}")

        start_time = time.time()

        for epoch in range(self.cfg.num_epochs):
            epoch_loss = 0.0
            epoch_ppl = 0.0
            n_batches = 0
            hypothesis_activations = 0

            # Synergy search no início da época (apenas época 0 para economizar tempo)
            if self.cfg.use_synergy_search and epoch == 0:
                try:
                    # Pega um batch de amostra para avaliação
                    sample_batch = next(iter(dataloader))
                    sample_batch = {k: (v.to(self.device) if isinstance(v, torch.Tensor) else v)
                                    for k, v in sample_batch.items()}

                    base_cfg = {
                        "alpha": self.loss_fn.cfg.alpha,
                        "beta": self.loss_fn.cfg.beta,
                        "gamma_loss": self.loss_fn.cfg.gamma_loss,
                        "delta": self.loss_fn.cfg.delta,
                        "lambda_penal": self.loss_fn.cfg.lambda_penal,
                        "mu_exp_penal": self.loss_fn.cfg.mu_exp_penal,
                    }

                    logger.info(f"Synergy search: tentando {self.cfg.n_synergy_attempts} configurações...")
                    best_cfg = self.synergy.search(
                        base_config=base_cfg,
                        eval_fn=lambda c: self._eval_config(c, sample_batch),
                    )

                    # Aplica a melhor configuração
                    for k, v in best_cfg.items():
                        if hasattr(self.loss_fn.cfg, k):
                            setattr(self.loss_fn.cfg, k, v)
                    logger.info("Synergy search concluído. Melhor config aplicada.")
                except Exception as e:
                    logger.warning(f"Synergy search falhou: {e}")

            # Loop de batches
            self.model.train()
            if self.generator: self.generator.train()
            if self.verifier: self.verifier.train()
            if self.anti_hall: self.anti_hall.train()
            self.hypothesis.train()

            for batch_idx, batch in enumerate(dataloader):
                if batch_idx >= self.cfg.max_batches_per_epoch:
                    break

                try:
                    metrics = self._train_step(batch)
                    epoch_loss += metrics["loss"]
                    epoch_ppl += metrics["ppl"]
                    n_batches += 1
                    if metrics["hypothesis_activated"]:
                        hypothesis_activations += 1

                    # Armazena todas as métricas
                    for k in self.metrics:
                        if k in metrics:
                            self.metrics[k].append(metrics[k])

                    if batch_idx % self.cfg.log_every == 0:
                        logger.info(
                            f"Época {epoch+1} | Batch {batch_idx} | "
                            f"loss={metrics['loss']:.4f} | ppl={metrics['ppl']:.2f} | "
                            f"v_mean={metrics['v_mean']:.3f} | lr_G={metrics['lr_G']:.2e} | "
                            f"l_ewc={metrics['l_ewc']:.4f} | "
                            f"hyp_act={'Y' if metrics['hypothesis_activated'] else 'N'}"
                        )
                except Exception as e:
                    logger.error(f"Erro no batch {batch_idx}: {e}", exc_info=True)
                    continue

            # Cleanup de memória ao final da época
            self.mem.cleanup()

            avg_loss = epoch_loss / max(n_batches, 1)
            avg_ppl = epoch_ppl / max(n_batches, 1)
            logger.info(
                f"Época {epoch+1} concluída | avg_loss={avg_loss:.4f} | "
                f"avg_ppl={avg_ppl:.2f} | hyp_activations={hypothesis_activations} | "
                f"batches_ok={n_batches}/{self.cfg.max_batches_per_epoch}"
            )

            # Detecta falha sistêmica
            if n_batches == 0:
                logger.error(f"Época {epoch+1}: NENHUM batch teve sucesso — possível bug no pipeline")

            self.history.append({
                "epoch": epoch + 1,
                "avg_loss": avg_loss,
                "avg_ppl": avg_ppl,
                "n_batches": n_batches,
                "hypothesis_activations": hypothesis_activations,
                "epoch_loss_total": epoch_loss,
            })

        elapsed = time.time() - start_time
        logger.info(f"Treinamento concluído em {elapsed:.1f}s")

        return {
            "history": self.history,
            "metrics": self.metrics,
            "final_loss": self.history[-1]["avg_loss"] if self.history else 0.0,
            "final_ppl": self.history[-1]["avg_ppl"] if self.history else 0.0,
            "elapsed_s": elapsed,
            "synergy_history": self.synergy.history,
            "auto_history": self.auto.history,
            "ewc_active": self.ewc_state is not None and self.ewc_state.is_active(),
            "ewc_num_tasks": self.ewc_state.num_tasks() if self.ewc_state else 0,
        }

    # =====================================================================
    # Consolidar tarefa para EWC (NOVO v2.0)
    # =====================================================================

    def consolidate_task(
        self,
        forward_fn,
        compute_fisher: bool = True,
    ) -> None:
        """Consolida a tarefa atual no estado EWC.

        Args:
            forward_fn: callable que retorna (logits, target) ou logits
                        para cálculo da Fisher Information.
            compute_fisher: se True, recalcula a Fisher; se False, apenas
                            armazena o theta_star atual.
        """
        if self.ewc_state is None:
            logger.warning("EWC não configurado — consolidate_task ignorado")
            return
        try:
            self.ewc_state.consolidate(self.model, forward_fn=forward_fn if compute_fisher else None)
            logger.info(f"Tarefa consolidada no EWC (total: {self.ewc_state.num_tasks()})")
        except Exception as e:
            logger.error(f"Falha ao consolidar tarefa no EWC: {e}")


__all__ = ["TrainerConfig", "CooperativeTrainer"]