src/bigru_t/training/ewc.py

"""ewc.py — V6/V7: Elastic Weight Consolidation with Kohonen 4D per-neuron support
+ Fisher accumulation (online EWC with EMA decay).

═══════════════════════════════════════════════════════════════════════════════
V7 — FISHER ACCUMULATION (ONLINE EWC)
═══════════════════════════════════════════════════════════════════════════════

V6 mantinha `online_ewc: bool = True` e `online_decay: float = 0.99` na
configuração, mas NUNCA os utilizava — a cada chamada de `compute_fisher()`,
a Fisher anterior era SOBRESCRITA pela nova (perda total do conhecimento
de tarefas anteriores). Isto explica o sintoma observado na análise V7:

  Fisher_w_mean = 0.0  (EWC sem acumulação Fisher — possível bug)

A correção V7 implementa a acumulação online clássica (Schwarz et al. 2018):

  F_new = decay · F_old + (1 - decay) · F_computed

onde:
  - F_old     : Fisher acumulada de tarefas anteriores (persistente)
  - F_computed: Fisher computada nesta tarefa (a partir de gradients²)
  - decay     : fator EMA (default 0.99 — preserva 99% do passado)
  - F_new     : Fisher acumulada após esta tarefa

Analogamente, os pesos ótimos (anchors) seguem EMA:
  W*_new = decay · W*_old + (1 - decay) · W_current

Isto garante que:
  1. Tarefas antigas continuem a ser "ancoradas" mesmo após muitas tarefas
     novas (evita catastrophic forgetting multi-tarefa).
  2. A Fisher seja suave (sem saltos bruscos) — estável sob gradient noise.
  3. A penalty V6 (clássica + Kohonen per-neuron + topológica) é aplicada
     sobre a Fisher ACUMULADA, não sobre a Fisher da última tarefa.

V6 — EWC + KOHONEN 4D PER-NEURON WEIGHTS
═══════════════════════════════════════════════════════════════════════════════

O usuário observou que o SOM 4D Kohonen (orq_cell.py V6) exige ajustes
matemáticos no EWC porque:

  "cada neurônio da grade tem peso próprio"

No EWC clássico (Kirkpatrick 2017), a Fisher Information Matrix é uma
diagonal que associa cada parâmetro θ_i a uma "importância" F_i. Os parâmetros
importantes para tarefas anteriores são penalizados quando se movem.

No V6, o OrqCell tem 8 neurônios SOM, cada um com:
  - Um vetor de cache w_n ∈ R^d_cache (peso próprio do neurônio)
  - Acesso compartilhado a proj_in, cross_attn, FFN

A novidade matemática é:

  1. **Fisher por neurônio**: cada neurônio n tem sua própria Fisher F_n (um
     vetor d_cache-dimensional). A penalidade EWC é:
        L_ewc_kohonen = λ * Σ_n Σ_i F_{n,i} (w_{n,i} - w*_{n,i})^2
     Isto penaliza a mudança individual de cada neurônio — preservando a
     topologia aprendida pelo SOM.

  2. **Anchoring do BMU**: quando uma nova tarefa chega, o BMU atual é
     identificado para cada sample. A Fisher do BMU é reforçada (multiplicada
     por γ_bmu > 1), porque o BMU é o neurônio mais "responsável" pela tarefa
     anterior e portanto o mais sensível a mudanças.

  3. **Topological regularization**: além da penalidade por neurônio, EWC V6
     adiciona uma penalidade que preserva a distância entre neurônios
     adjacentes no grid 4D:
        L_topo = λ_topo * Σ_(i,j)∈edges ||d_grid(i,j) - d_orig(i,j)||^2
     onde d_grid é a distância no grid 4D (fixa) e d_orig é a distância
     atual entre os pesos dos neurônios. Isto garante que a topologia do SOM
     não colapse com a tarefa nova.

  4. **Compatibilidade W8A8 em eval**: quando o modelo está quantizado W8A8
     e em eval mode, o EWC não computa Fisher (não há backward em eval).
     Apenas a penalidade forward é aplicada nos pesos do mapa. Isto permite
     avaliar o quão bem o modelo retém conhecimento de tarefas anteriores
     sem re-quantizar.

Matemática
──────────
  L_total = L_task + λ * Σ_n Σ_i F_{n,i} (w_{n,i} - w*_{n,i})^2
              + λ_topo * Σ_(i,j)∈edges (||w_i - w_j|| - d_grid(i,j))^2

  λ adaptativo: λ_t = λ_0 * ||∇L_task|| / (||∇L_ewc|| + ε)
"""
from __future__ import annotations
import gc
import math
from typing import Dict, List, Optional, Tuple, Any
from dataclasses import dataclass, field

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


@dataclass
class EWCConfigV6:
    """Configuração EWC V6 com suporte a Kohonen 4D per-neuron.

    V6.1 adiciona `skip_som_filled_neurons`: quando True (default), o EWC
    NÃO computa Fisher nem aplica penalidade para neurônios do SOM que
    já estão "filled" (neuron_filled_mask=True). Isto evita a dupla
    penalização — o SOM já está consolidando esses pesos via regra de
    Kohonen, então o EWC não precisa "ancorar" aquela região do mapa.

    O usuário observou: "Elastic Weight Consolidation que deve ser
    matematicamente aprimorado para não usar pesos já existentes em SOM".
    """
    lambda_ewc: float = 100.0
    lambda_topo: float = 10.0           # penalidade topológica
    fisher_n_samples: int = 32
    fisher_batch_size: int = 4
    online_ewc: bool = True
    online_decay: float = 0.99          # V7 — agora EFETIVAMENTE usado (era decorativo)
    fisher_smoothing_alpha: float = 0.1
    adaptive_lambda: bool = True
    fisher_clip_eps: float = 1e-8
    bmu_reinforcement_gamma: float = 2.0  # γ para reforço do BMU na Fisher
    eval_mode_penalty: bool = True        # aplicar penalidade em eval mode (sem Fisher)
    # ───────────────────────────────────────────────────────────────────────
    # V6.1 — Skip SOM-filled neurons (não reutilizar pesos do SOM)
    # ───────────────────────────────────────────────────────────────────────
    skip_som_filled_neurons: bool = True  # default True (user requirement)
    # Skip_mask é lido do OrqCell.neuron_filled_mask em compute_fisher()
    # e aplicado em compute_penalty() para zerar a penalidade dos filled.
    # ───────────────────────────────────────────────────────────────────────
    # V7 — Fisher accumulation diagnostics (apenas leitura, não-configurável)
    # ───────────────────────────────────────────────────────────────────────
    # fisher_accumulation_active: se True, compute_fisher() faz EMA entre
    #   F_old e F_computed (V7 — ativação do online_ewc).
    # fisher_ema_initialized: True após a 1ª compute_fisher (não há F_old
    #   na 1ª chamada — apenas inicializa).
    fisher_accumulation_active: bool = True  # V7 — ativa por padrão


class EWCV6:
    """EWC V6 com suporte a SOM 4D Kohonen (orq_cell.OrqCell).

    Detecta automaticamente a presença de OrqCell no modelo e ativa o modo
    per-neuron. Caso contrário, degrada para EWC clássico (compatibilidade).
    """

    def __init__(self, model: nn.Module, config: Optional[EWCConfigV6] = None):
        self.model = model
        self.cfg = config or EWCConfigV6()
        # Fisher clássico (parâmetros gerais)
        self.fisher: Dict[str, torch.Tensor] = {}
        self.optimal_weights: Dict[str, torch.Tensor] = {}
        # Fisher por neurônio (orq_cell)
        self.neuron_fisher: Dict[str, torch.Tensor] = {}  # {orq_name: (n_neurons, d_cache)}
        self.neuron_optimal: Dict[str, torch.Tensor] = {}  # {orq_name: (n_neurons, d_cache)}
        # Topologia original (para regularização)
        self.neuron_grid_dist: Dict[str, torch.Tensor] = {}  # {orq_name: (n_neurons, n_neurons)}
        # ───────────────────────────────────────────────────────────────────
        # V6.1 — fisher_neuron_mask: marca quais neurônios têm Fisher > 0
        # (apenas os não-filled no SOM, para evitar dupla penalização).
        # Se skip_som_filled_neurons=True (default), neuron_fisher é zerado
        # para neurônios que estão em OrqCell.neuron_filled_mask=True.
        # ───────────────────────────────────────────────────────────────────
        self.fisher_neuron_mask: Dict[str, torch.Tensor] = {}  # {orq_name: (n_neurons,) bool}
        # Stats
        self.n_tasks = 0
        self._last_task_grad_norm: float = 0.0
        self._last_ewc_grad_norm: float = 0.0
        self._adaptive_lambda: float = self.cfg.lambda_ewc
        # ───────────────────────────────────────────────────────────────────
        # V7 — Fisher accumulation state (online EWC with EMA)
        # ───────────────────────────────────────────────────────────────────
        # _fisher_ema_initialized: False até a 1ª compute_fisher() concluir.
        #   Antes disso, compute_fisher() apenas inicializa F = F_computed.
        #   Após isso, F = decay·F_old + (1-decay)·F_computed.
        self._fisher_ema_initialized: bool = False
        # Diagnóstico: média da Fisher acumulada (para auditoria)
        self._fisher_w_mean: float = 0.0   # média da Fisher clássica (scalar)
        self._fisher_n_mean: float = 0.0   # média da Fisher per-neuron (scalar)
        # Detectar OrqCells no modelo
        self.orq_cells: Dict[str, nn.Module] = self._find_orq_cells()

    # =====================================================================
    # Detecção de OrqCell
    # =====================================================================
    def _find_orq_cells(self) -> Dict[str, nn.Module]:
        """Encontra todas as instâncias de OrqCell no modelo."""
        cells = {}
        try:
            # Tenta importar OrqCell para isinstance check
            from ..model.orq_cell import OrqCell
        except (ImportError, AttributeError):
            return cells
        for name, module in self.model.named_modules():
            if isinstance(module, OrqCell):
                cells[name] = module
        return cells

    # =====================================================================
    # Compute Fisher Information
    # =====================================================================
    def compute_fisher(
        self,
        data_iter,
        n_samples: Optional[int] = None,
        bmu_indices: Optional[List[int]] = None,
    ) -> None:
        """Computa Fisher Information Matrix diagonal.

        Para OrqCell, computa Fisher por neurônio (n_neurons, d_cache).
        Para outros parâmetros, computa Fisher clássico (shape do parâmetro).

        Args:
            data_iter: iterável de (inputs, targets) — função closure que
                       retorna um batch
            n_samples: nº de samples a usar (default: cfg.fisher_n_samples)
            bmu_indices: lista de BMU IDs por sample (para reforço γ_bmu)
        """
        n_samples = n_samples or self.cfg.fisher_n_samples
        bs = self.cfg.fisher_batch_size

        # Inicializa Fisher acumulado
        fisher_accum: Dict[str, torch.Tensor] = {}
        neuron_fisher_accum: Dict[str, torch.Tensor] = {}

        self.model.eval()  # fisher em eval para não atualizar BN/dropout

        n_used = 0
        sample_idx = 0
        for batch in data_iter:
            if n_used >= n_samples:
                break
            inputs, targets = batch
            self.model.zero_grad()
            try:
                # Tenta forward com returns padrão
                outputs = self.model(inputs)
                if isinstance(outputs, tuple):
                    logits = outputs[0]
                else:
                    logits = outputs
                # Loss para Fisher: cross-entropy ou MSE dependendo da shape
                if logits.dim() == targets.dim() + 1:
                    loss = F.cross_entropy(logits.view(-1, logits.size(-1)),
                                            targets.view(-1))
                else:
                    loss = F.mse_loss(logits, targets.float())
                loss.backward()
            except (RuntimeError, ValueError):
                continue

            # Acumula Fisher para cada parâmetro
            for name, p in self.model.named_parameters():
                if p.grad is None:
                    continue
                g2 = p.grad.detach().pow(2)
                if name not in fisher_accum:
                    fisher_accum[name] = g2.clone()
                else:
                    fisher_accum[name] += g2

            # Acumula Fisher por neurônio para OrqCells
            for orq_name, orq_cell in self.orq_cells.items():
                if not hasattr(orq_cell, "neuron_cache"):
                    continue
                # neuron_cache: (n_neurons, d_cache) — parâmetro único
                param_name = f"{orq_name}.neuron_cache"
                if param_name not in fisher_accum:
                    continue
                # Fisher já acumulada acima; vamos separá-la por neurônio
                # (fisher_accum[param_name] tem shape (n_neurons, d_cache))
                # Aplicar reforço do BMU se disponível
                bmu_id = None
                if bmu_indices is not None and sample_idx < len(bmu_indices):
                    bmu_id = bmu_indices[sample_idx]
                fisher_per_neuron = fisher_accum[param_name].clone()
                if bmu_id is not None and 0 <= bmu_id < fisher_per_neuron.size(0):
                    # γ_bmu reforço no BMU
                    fisher_per_neuron[bmu_id] *= self.cfg.bmu_reinforcement_gamma
                if orq_name not in neuron_fisher_accum:
                    neuron_fisher_accum[orq_name] = fisher_per_neuron
                else:
                    neuron_fisher_accum[orq_name] += fisher_per_neuron

            n_used += bs
            sample_idx += 1
            if n_used % (bs * 10) == 0:
                gc.collect(0)

        # Normaliza pelo nº de samples
        if n_used > 0:
            for name in fisher_accum:
                fisher_accum[name] /= max(1, n_used // bs)
            for orq_name in neuron_fisher_accum:
                neuron_fisher_accum[orq_name] /= max(1, n_used // bs)

        # Smooth + clip Fisher (compute F_computed — smoothed/clipped)
        fisher_computed = self._smooth_and_clip(fisher_accum)
        # V6.1 — Aplicar skip mask também no Fisher CLÁSSICO para
        # orq_cell.neuron_cache (evita dupla penalização: o classic Fisher
        # e o Kohonen per-neuron Fisher apontam para o MESMO parâmetro).
        # Sem isto, o classic Fisher penalizaria neuron_cache mesmo com
        # skip_som_filled_neurons=True, anulando o efeito da skip mask.
        if self.cfg.skip_som_filled_neurons:
            for orq_name, orq_cell in self.orq_cells.items():
                if not hasattr(orq_cell, "neuron_filled_mask"):
                    continue
                param_name = f"{orq_name}.neuron_cache"
                if param_name not in fisher_computed:
                    continue
                filled = orq_cell.neuron_filled_mask.bool()
                keep_mask = (~filled).float()  # 1 onde EWC atua, 0 onde skip
                # Zera o Fisher das posições filled no classic Fisher
                fisher_computed[param_name] = fisher_computed[param_name] * keep_mask.unsqueeze(-1)

        # V7 ── Fisher accumulation (online EWC with EMA) ─────────────────
        # F_new = decay · F_old + (1 - decay) · F_computed
        # Se for a 1ª chamada (não inicializado), apenas copia F_computed.
        # Aplicado APÓS smoothing/clip/skip_mask para herdar as propriedades
        # de estabilidade numérica (log-clip evita explosão de Fisher EMA).
        if self.cfg.online_ewc and self.cfg.fisher_accumulation_active \
                and self._fisher_ema_initialized and len(self.fisher) > 0:
            decay = float(self.cfg.online_decay)
            for name, f_new_computed in fisher_computed.items():
                f_old = self.fisher.get(name)
                if f_old is not None and f_old.shape == f_new_computed.shape:
                    fisher_computed[name] = (
                        decay * f_old + (1.0 - decay) * f_new_computed
                    )
                # else: novo parâmetro — apenas usa F_computed (cold start)
        # Aplica Fisher final (acumulada ou cold-start)
        self.fisher = fisher_computed

        # ── Per-neuron Fisher accumulation (V7) ──────────────────────────
        # Mesma lógica EMA para neuron_fisher (orq_cell).
        neuron_fisher_computed: Dict[str, torch.Tensor] = {}
        self.fisher_neuron_mask = {}  # V6.1
        for orq_name, f_tensor in neuron_fisher_accum.items():
            # V6.1 — Aplicar skip mask: neurônios filled no SOM NÃO recebem Fisher
            orq_cell = self.orq_cells.get(orq_name)
            skip_mask = None
            if (self.cfg.skip_som_filled_neurons and
                    orq_cell is not None and
                    hasattr(orq_cell, "neuron_filled_mask")):
                # filled_mask=True → skip (zerar Fisher). Inverter para "keep_mask"
                filled = orq_cell.neuron_filled_mask.bool()
                skip_mask = ~filled  # True onde EWC DEVE atuar (não-filled)
                # Aplicar: zero out Fisher dos filled
                f_tensor = f_tensor * skip_mask.unsqueeze(-1).float()
                self.fisher_neuron_mask[orq_name] = skip_mask.detach().clone()
            else:
                # Sem skip mask: todos os neurônios são elegíveis
                if f_tensor.size(0) > 0:
                    self.fisher_neuron_mask[orq_name] = torch.ones(
                        f_tensor.size(0), dtype=torch.bool
                    )
            f_smoothed = self._smooth_and_clip_tensor(f_tensor)
            # V7 — EMA accumulation sobre neuron_fisher
            if self.cfg.online_ewc and self.cfg.fisher_accumulation_active \
                    and self._fisher_ema_initialized \
                    and orq_name in self.neuron_fisher:
                f_old = self.neuron_fisher[orq_name]
                if f_old.shape == f_smoothed.shape:
                    decay = float(self.cfg.online_decay)
                    f_smoothed = decay * f_old + (1.0 - decay) * f_smoothed
            neuron_fisher_computed[orq_name] = f_smoothed
        self.neuron_fisher = neuron_fisher_computed

        # Salvar pesos ótimos atuais (anchors) — V7 EMA accumulation
        # W*_new = decay · W*_old + (1 - decay) · W_current
        # Se 1ª chamada, apenas W*_new = W_current (cold start).
        new_optimal_weights = {
            name: p.detach().clone()
            for name, p in self.model.named_parameters()
            if name in self.fisher
        }
        if (self.cfg.online_ewc and self.cfg.fisher_accumulation_active
                and self._fisher_ema_initialized and len(self.optimal_weights) > 0):
            decay = float(self.cfg.online_decay)
            for name, w_new in new_optimal_weights.items():
                w_old = self.optimal_weights.get(name)
                if w_old is not None and w_old.shape == w_new.shape:
                    new_optimal_weights[name] = (
                        decay * w_old + (1.0 - decay) * w_new
                    )
        self.optimal_weights = new_optimal_weights

        # V7 — EMA accumulation para neuron_optimal (orq_cell anchors)
        for orq_name, orq_cell in self.orq_cells.items():
            if orq_name in self.neuron_fisher:
                new_w_n = orq_cell.neuron_cache.detach().clone()
                if (self.cfg.online_ewc and self.cfg.fisher_accumulation_active
                        and self._fisher_ema_initialized
                        and orq_name in self.neuron_optimal):
                    w_old = self.neuron_optimal[orq_name]
                    if w_old.shape == new_w_n.shape:
                        decay = float(self.cfg.online_decay)
                        new_w_n = decay * w_old + (1.0 - decay) * new_w_n
                self.neuron_optimal[orq_name] = new_w_n
                # Salvar topologia do grid (fixa — NÃO sofre EMA, é constante)
                if hasattr(orq_cell, "grid_dist"):
                    self.neuron_grid_dist[orq_name] = orq_cell.grid_dist.detach().clone()

        # V7 — Diagnóstico: média das Fishers acumuladas
        if self.fisher:
            try:
                all_means = [float(f.mean().item()) for f in self.fisher.values()]
                self._fisher_w_mean = sum(all_means) / len(all_means) if all_means else 0.0
            except (RuntimeError, ValueError):
                self._fisher_w_mean = 0.0
        if self.neuron_fisher:
            try:
                all_n_means = [float(f.mean().item()) for f in self.neuron_fisher.values()]
                self._fisher_n_mean = sum(all_n_means) / len(all_n_means) if all_n_means else 0.0
            except (RuntimeError, ValueError):
                self._fisher_n_mean = 0.0

        # V7 — Marca que a Fisher foi inicializada (próxima compute_fisher fará EMA)
        self._fisher_ema_initialized = True

        self.n_tasks += 1
        self.model.zero_grad()
        gc.collect(0)

    # =====================================================================
    # Smoothing + Clipping (V5/V13.2数学)
    # =====================================================================
    def _smooth_and_clip(self, fisher: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        out = {}
        for name, f in fisher.items():
            out[name] = self._smooth_and_clip_tensor(f)
        return out

    def _smooth_and_clip_tensor(self, f: torch.Tensor) -> torch.Tensor:
        # Fisher smoothing: (1-α) F + α mean(F)
        if self.cfg.fisher_smoothing_alpha > 0:
            mean_f = f.mean()
            f = (1 - self.cfg.fisher_smoothing_alpha) * f + self.cfg.fisher_smoothing_alpha * mean_f
        # Log clipping: log(1 + F/ε) / log(1 + max(F)/ε)
        eps = self.cfg.fisher_clip_eps
        f_max = f.max().clamp(min=eps)
        f = torch.log1p(f / eps) / torch.log1p(f_max / eps)
        return f

    # =====================================================================
    # Compute EWC penalty
    # =====================================================================
    def compute_penalty(
        self,
        bmu_indices: Optional[torch.Tensor] = None,
        eval_mode: bool = False,
    ) -> Tuple[torch.Tensor, Dict[str, float]]:
        """Computa a penalidade EWC V6 (clássica + Kohonen per-neuron + topológica).

        Args:
            bmu_indices: (batch,) — índices BMU atuais (para reforço)
            eval_mode: se True, apenas forward penalty (sem Fisher backward)

        Returns:
            penalty: tensor escalar
            metrics: dict com breakdown da penalidade
        """
        if not self.fisher and not self.neuron_fisher:
            return torch.tensor(0.0, device=next(self.model.parameters()).device), {}

        # V6.1 — Short-circuit: se TODOS os neurônios do SOM estão filled,
        # a skip mask zera todo o Fisher Kohonen. E se o classic Fisher
        # também foi mascarado para orq_cell.neuron_cache (feito em
        # compute_fisher), resta apenas o classic Fisher dos outros params.
        # Para evitar NaN por instabilidade numérica na topological penalty
        # (sqrt de diferenças de normas), pulamos o topo penalty quando
        # todos os neurônios estão filled — a topologia já está sendo
        # preservada pelo próprio SOM (BMU + vizinhos puxados na direção
        # do input, Gaussian Λ).
        all_filled = all(
            hasattr(oc, "neuron_filled_mask") and bool(oc.neuron_filled_mask.all().item())
            for oc in self.orq_cells.values()
        ) if self.orq_cells else False
        if all_filled and self.cfg.skip_som_filled_neurons:
            # SOM cobre tudo — EWC não tem nada a ancorar
            device = next(self.model.parameters()).device
            metrics = {
                "total_pen": 0.0,
                "n_masked_som_neurons": sum(
                    int(oc.neuron_filled_mask.sum().item())
                    for oc in self.orq_cells.values()
                    if hasattr(oc, "neuron_filled_mask")
                ),
                "skipped_reason": "all_som_neurons_filled",
            }
            return torch.tensor(0.0, device=device), metrics

        device = next(self.model.parameters()).device
        metrics: Dict[str, float] = {}
        total_penalty = torch.tensor(0.0, device=device)

        # 1. Penalidade clássica (parâmetros gerais)
        classic_pen = torch.tensor(0.0, device=device)
        n_classic = 0
        for name, p in self.model.named_parameters():
            if name not in self.fisher or name not in self.optimal_weights:
                continue
            f = self.fisher[name].to(device)
            w_star = self.optimal_weights[name].to(device)
            classic_pen = classic_pen + (f * (p - w_star).pow(2)).sum()
            n_classic += 1
        if n_classic > 0:
            classic_pen = classic_pen / max(1, n_classic)
            metrics["classic_pen"] = float(classic_pen.item())
            total_penalty = total_penalty + self._adaptive_lambda * classic_pen

        # 2. Penalidade Kohonen per-neuron
        kohonen_pen = torch.tensor(0.0, device=device)
        n_kohonen = 0
        n_masked_total = 0  # V6.1: contador de neurônios skipped (já no SOM)
        for orq_name, orq_cell in self.orq_cells.items():
            if orq_name not in self.neuron_fisher or orq_name not in self.neuron_optimal:
                continue
            f_n = self.neuron_fisher[orq_name].to(device)  # (n_neurons, d_cache)
            w_star_n = self.neuron_optimal[orq_name].to(device)
            w_n = orq_cell.neuron_cache  # (n_neurons, d_cache)
            # V6.1 — Aplicar fisher_neuron_mask: zero penalty para filled neurons
            keep_mask = self.fisher_neuron_mask.get(orq_name)
            if keep_mask is not None:
                keep_mask_dev = keep_mask.to(device).float()
                f_n_masked = f_n * keep_mask_dev.unsqueeze(-1)
                n_masked_total += int((1 - keep_mask_dev).sum().item())
            else:
                f_n_masked = f_n
            # Σ_n Σ_i F_{n,i} (w_{n,i} - w*_{n,i})^2
            pen_n = (f_n_masked * (w_n - w_star_n).pow(2)).sum()
            # Reforço do BMU: se bmu_indices fornecido, multiplica a penalidade
            # do neurônio BMU por γ_bmu
            if bmu_indices is not None and len(bmu_indices) > 0:
                bmu_counts = torch.zeros(w_n.size(0), device=device)
                for b in bmu_indices.tolist():
                    if 0 <= b < w_n.size(0):
                        bmu_counts[b] += 1.0
                # weight per neuron: 1 + γ*(bmu_counts normalized)
                bmu_weight = 1.0 + (self.cfg.bmu_reinforcement_gamma - 1.0) * (
                    bmu_counts / (bmu_counts.sum() + 1e-8)
                )
                pen_n = (bmu_weight.unsqueeze(-1) * f_n_masked * (w_n - w_star_n).pow(2)).sum()
            kohonen_pen = kohonen_pen + pen_n
            n_kohonen += 1
        if n_kohonen > 0:
            kohonen_pen = kohonen_pen / max(1, n_kohonen)
            metrics["kohonen_pen"] = float(kohonen_pen.item())
            metrics["n_masked_som_neurons"] = n_masked_total  # V6.1
            total_penalty = total_penalty + self._adaptive_lambda * kohonen_pen

        # 3. Penalidade topológica (preserva distâncias no grid 4D)
        topo_pen = torch.tensor(0.0, device=device)
        n_topo = 0
        for orq_name, orq_cell in self.orq_cells.items():
            if orq_name not in self.neuron_grid_dist:
                continue
            grid_d = self.neuron_grid_dist[orq_name].to(device)  # (n,n) — fixa
            w_n = orq_cell.neuron_cache  # (n, d_cache)
            # Distância atual entre neurônios (espaço de pesos)
            sq = (w_n ** 2).sum(dim=-1)
            actual_d = torch.sqrt(
                sq.unsqueeze(0) + sq.unsqueeze(1) - 2 * w_n @ w_n.T
            ).clamp(min=0.0)
            # Penalidade: (actual - grid)^2 — upper triangular
            diff = (actual_d - grid_d).pow(2)
            mask = torch.triu(torch.ones_like(diff), diagonal=1).bool()
            topo_pen = topo_pen + diff[mask].sum()
            n_topo += 1
        if n_topo > 0:
            topo_pen = topo_pen / max(1, n_topo)
            metrics["topo_pen"] = float(topo_pen.item())
            total_penalty = total_penalty + self.cfg.lambda_topo * topo_pen

        metrics["total_pen"] = float(total_penalty.item())
        metrics["adaptive_lambda"] = float(self._adaptive_lambda)
        return total_penalty, metrics

    # =====================================================================
    # Update adaptive lambda (após cada step)
    # =====================================================================
    def update_adaptive_lambda(self, task_loss: torch.Tensor, ewc_penalty: torch.Tensor) -> None:
        """λ_t = λ_0 * ||∇L_task|| / (||∇L_ewc|| + ε)"""
        if not self.cfg.adaptive_lambda:
            return
        try:
            task_grad_norm = 0.0
            ewc_grad_norm = 0.0
            for name, p in self.model.named_parameters():
                if p.grad is None:
                    continue
                g = p.grad.detach()
                task_grad_norm += float(g.norm().item()) ** 2
            ewc_grad_norm = float(ewc_penalty.item()) + 1e-8
            task_grad_norm = math.sqrt(task_grad_norm) + 1e-8
            self._last_task_grad_norm = task_grad_norm
            self._last_ewc_grad_norm = ewc_grad_norm
            # λ adaptive (clipado em [0.1*λ_0, 10*λ_0])
            new_lambda = self.cfg.lambda_ewc * task_grad_norm / ewc_grad_norm
            self._adaptive_lambda = max(
                0.1 * self.cfg.lambda_ewc,
                min(10.0 * self.cfg.lambda_ewc, new_lambda)
            )
        except (RuntimeError, AttributeError):
            pass

    # =====================================================================
    # Status / monitoramento
    # =====================================================================
    def get_status(self) -> Dict[str, Any]:
        n_masked = sum(int((~m).sum().item()) for m in self.fisher_neuron_mask.values()) if self.fisher_neuron_mask else 0
        n_active = sum(int(m.sum().item()) for m in self.fisher_neuron_mask.values()) if self.fisher_neuron_mask else 0
        return {
            "n_tasks": self.n_tasks,
            "n_classic_params": len(self.fisher),
            "n_kohonen_cells": len(self.neuron_fisher),
            "adaptive_lambda": self._adaptive_lambda,
            "lambda_ewc_base": self.cfg.lambda_ewc,
            "lambda_topo": self.cfg.lambda_topo,
            "bmu_reinforcement_gamma": self.cfg.bmu_reinforcement_gamma,
            "last_task_grad_norm": self._last_task_grad_norm,
            "last_ewc_grad_norm": self._last_ewc_grad_norm,
            # V6.1 — info sobre skip mask
            "skip_som_filled_neurons": self.cfg.skip_som_filled_neurons,
            "n_active_ewc_neurons": n_active,    # EWC atua nestes
            "n_masked_som_neurons": n_masked,    # SOM já cobre estes (não duplica)
            # V7 — Fisher accumulation diagnostics
            "online_ewc_active": bool(self.cfg.online_ewc and self.cfg.fisher_accumulation_active),
            "online_decay": float(self.cfg.online_decay),
            "fisher_ema_initialized": bool(self._fisher_ema_initialized),
            "fisher_w_mean": float(self._fisher_w_mean),  # diagnostic
            "fisher_n_mean": float(self._fisher_n_mean),  # diagnostic
        }


__all__ = ["EWCV6", "EWCConfigV6"]