src/bigru_t/training/data_augmentation.py

"""data_augmentation.py — V7: data augmentation for BiGRU_T training.

ESTRATÉGIAS:
  A. DataAugmenter (clássica — token-level):
     1. Token dropout: substitui aleatoriamente tokens por [MASK]
     2. Token shuffle: embaralha tokens não-adjacentes (preserva ordem local)
     3. Mixup de hidden states: combina hidden states de duas amostras
     4. Cutout de sequência: remove um segmento aleatório da sequência
     5. Back-translation simulada: adiciona ruído controlado aos tokens
  B. AdvancedSOMAugmenter (V7 — manifold-aware):
     6. SOM 4D canônico (4,4,4,4)=256 neurônios treina densidade BMU
     7. Geração sintética via interpolação convexa (0.7*real + 0.3*BMU)
     8. Ruído gaussiano adaptativo proporcional ao Quantization Error (QE)
        — garante que amostras sintéticas pertençam ao manifold dos dados
        reais, evitando outliers e distorções.

============================================================================
V7 — APRIMORAMENTO LÓGICO/MATEMÁTICO: AdvancedSOMAugmenter
============================================================================
O SOM tradicional apenas MAPEIA dados. Esta implementação utiliza a densidade
topológica dos neurônios vencedores (BMU) combinada com um ruído gaussiano
dinâmico proporcional à quantização do erro.

Justificativa matemática:
- Amostras sintéticas devem pertencer à mesma variedade (manifold) dos
  dados reais. O vetor de peso do BMU é o ponto do grid SOM mais próximo
  da amostra real, logo a interpolação convexa 0.7*x + 0.3*W_BMU(x)
  permanece na vizinhança topológica do dado real.
- O ruído gaussiano é proporcional ao QE local (||x - W_BMU(x)||):
    noise ~ N(0, σ_adaptive), σ_adaptive = noise_scale * (1 + QE)
  Quando QE é alto (região mal representada), o ruído é maior — permitindo
  explorar a vizinhança. Quando QE é baixo (região bem representada), o
  ruído é pequeno — evitando distorcer amostras já fieis.
- Inicialização dos pesos baseada em variância para convergência rápida:
    W ~ N(0, 0.1) inicialmente; após fit(), W reflete a estrutura PCA-like
    do corpus (similar à inicialização PCA do Kohonen clássico).

CANÔNICO V6.5-V4:
- Grid SOM: (4, 4, 4, 4) = 256 neurônios (NÃO 10x10 como em referências)
- input_dim: 4 (vetor 4D = [x, y, z, w] da projeção SVD + tempo linear)
- α₀: 0.5 (∈ [0.5, 1.0] Kohonen rough training)
- σ₀: 2.0 (= max(4,4,4,4)/2 = metade da maior dimensão da grade)
- Decaimento exponencial clássico: α_t = α₀ · exp(-t/τ_α), σ_t = σ₀ · exp(-t/τ_σ)

REFERÊNCIA: AdvancedSOMAugmenter (user-provided pattern, adaptado para
PyTorch + grid 4D canônico V6.5-V4).
"""
from __future__ import annotations
import math
import random
import logging
from typing import Optional, Tuple, List, Dict, Any

import torch
import torch.nn as nn

logger = logging.getLogger(__name__)


# ============================================================================
# A. DataAugmenter clássico (token-level) — preservado de V6
# ============================================================================
from dataclasses import dataclass


@dataclass
class AugmentationConfig:
    """Configuração do data augmentation V7 (token-level)."""
    p_token_dropout: float = 0.1
    p_token_shuffle: float = 0.05
    p_mixup: float = 0.0          # desativado por padrão (precisa de 2 amostras)
    p_cutout: float = 0.05
    p_noise: float = 0.05
    mask_token_id: int = 0        # ID do token [MASK]
    vocab_size: int = 16384       # CANÔNICO V6.5-V4
    max_cutout_length: int = 4
    mixup_alpha: float = 0.2


class DataAugmenter(nn.Module):
    """V7: data augmentation token-level para BiGRU_T.

    Forward:
        input_ids: (B, T) → augmented_ids: (B, T)
        targets: (B, T) → targets (inalterados)

    Em eval mode, retorna input_ids sem modificação.
    """
    def __init__(self, config: Optional[AugmentationConfig] = None):
        super().__init__()
        self.cfg = config or AugmentationConfig()

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        if not self.training:
            return input_ids
        ids = input_ids.clone()
        B, T = ids.shape

        # 1. Token dropout
        if self.cfg.p_token_dropout > 0:
            mask = torch.rand_like(ids.float()) < self.cfg.p_token_dropout
            ids[mask] = self.cfg.mask_token_id

        # 2. Token shuffle (não-adjacente)
        if self.cfg.p_token_shuffle > 0 and T > 4:
            for b in range(B):
                if random.random() < self.cfg.p_token_shuffle:
                    even_idx = list(range(0, T, 2))
                    if len(even_idx) > 1:
                        perm = even_idx.copy()
                        random.shuffle(perm)
                        for orig, new in zip(even_idx, perm):
                            ids[b, orig] = input_ids[b, new]

        # 3. Cutout de sequência
        if self.cfg.p_cutout > 0:
            for b in range(B):
                if random.random() < self.cfg.p_cutout:
                    cut_len = random.randint(1, self.cfg.max_cutout_length)
                    cut_start = random.randint(0, max(0, T - cut_len))
                    ids[b, cut_start:cut_start + cut_len] = self.cfg.mask_token_id

        # 4. Noise (substitui por token aleatório)
        if self.cfg.p_noise > 0:
            mask = torch.rand_like(ids.float()) < self.cfg.p_noise
            random_tokens = torch.randint_like(ids, 0, self.cfg.vocab_size)
            ids[mask] = random_tokens[mask]

        return ids

    def mixup(
        self,
        x1: torch.Tensor,
        x2: torch.Tensor,
        targets1: torch.Tensor,
        targets2: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Mixup de hidden states: combina duas amostras.

        λ ~ Beta(α, α)
        x_mix = λ * x1 + (1-λ) * x2
        targets_mix = λ * targets1 + (1-λ) * targets2
        """
        if not self.training or self.cfg.p_mixup == 0:
            return x1, targets1, torch.ones(x1.size(0), device=x1.device)
        lam = float(torch.distributions.Beta(self.cfg.mixup_alpha, self.cfg.mixup_alpha).sample())
        x_mix = lam * x1 + (1 - lam) * x2
        t_mix = lam * targets1 + (1 - lam) * targets2
        return x_mix, t_mix, torch.tensor(lam)


# ============================================================================
# B. AdvancedSOMAugmenter (V7 — manifold-aware SOM augmentation)
# ============================================================================
class AdvancedSOMAugmenter(nn.Module):
    """V7: SOM 4D canônico para data augmentation manifold-aware.

    Diferente do SOM tradicional (que apenas mapeia dados), esta implementação
    utiliza a densidade topológica dos neurônios vencedores (BMU) combinada
    com um ruído gaussiano dinâmico proporcional ao Quantization Error (QE).
    Isso garante que os dados sintéticos gerados pertençam à mesma variedade
    (manifold) dos dados reais, evitando distorções (outliers).

    Pipeline de augmentação:
        1. Treina SOM 4D sobre o buffer (fit)
        2. Para cada amostra sintética desejada:
           a. Seleciona amostra real aleatória como âncora (x)
           b. Encontra BMU(x) — neurônio mais próximo no grid 4D
           c. Recupera vetor de peso W_BMU(x)
           d. Computa QE local: ||x - W_BMU(x)||
           e. Gera ruído gaussiano adaptativo: N(0, σ_adaptive)
              onde σ_adaptive = noise_scale * (1 + QE)
           f. Amostra sintética = 0.7*x + 0.3*W_BMU(x) + ruído

    Justificativa matemática:
    - 0.7*x + 0.3*W_BMU: interpolação convexa que preserva 70% da estrutura
      original e puxa 30% em direção ao protótipo do BMU (manifold).
    - Ruído adaptativo por QE: em regiões mal representadas (QE alto),
      permite-se mais variação; em regiões bem representadas (QE baixo),
      mantém-se fiel ao manifold. Isto evita gerar outliers em regiões
      densas e permite explorar regiões esparsas.
    - Inicialização N(0, 0.1) dos pesos garante convergência rápida
      (pequena variância inicial → estável nos primeiros passos).

    Args:
        dims: tupla (I, J, K, L) — default canônico V6.5-V4 = (4, 4, 4, 4) = 256
            neurônios. NÃO usar 10x10 (2D) — código BiGRU-T opera em 4D.
        input_dim: dimensão do vetor de entrada. Default 4 = [x, y, z, w]
            (projeção SVD 3D + coordenada temporal linear).
        lr: taxa de aprendizado inicial α₀. Default 0.5 (∈ [0.5, 1.0]
            Kohonen rough training, conforme especificação do usuário).
        sigma: raio inicial de vizinhança σ₀. Default = max(dims)/2 = 2.0
            (metade da maior dimensão da grade, conforme Kohonen clássico).
        device: dispositivo PyTorch ('cpu' ou 'cuda').
    """

    def __init__(
        self,
        dims: Tuple[int, int, int, int] = (4, 4, 4, 4),
        input_dim: int = 4,
        lr: float = 0.5,
        sigma: Optional[float] = None,
        device: Optional[str] = None,
    ):
        super().__init__()
        assert len(dims) == 4, f"AdvancedSOMAugmenter requer grid 4D, recebeu {len(dims)}D"
        assert input_dim == 4, (
            f"AdvancedSOMAugmenter requer input_dim=4 (vetor 4D [x,y,z,w]), "
            f"recebeu input_dim={input_dim}"
        )
        self.dims = dims
        self.input_dim = input_dim
        self.lr_init = float(lr)
        # σ₀ = metade da maior dimensão da grade (Kohonen clássico)
        self.sigma_init = float(sigma) if sigma is not None else max(dims) / 2.0
        self.device = torch.device(device) if device else torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )

        # Inicialização dos pesos baseada em variância para convergência rápida
        # W ~ N(0, 0.1) — pequena variância inicial garante estabilidade
        # nos primeiros passos (Kohonen rough training).
        self.weights = nn.Parameter(
            torch.randn((*self.dims, self.input_dim), dtype=torch.float32,
                        device=self.device) * 0.1
        )

        # Pré-computa coordenadas do grid 4D (para cálculo de vizinhança)
        grid_coords = torch.stack(torch.meshgrid(
            torch.arange(dims[0], device=self.device),
            torch.arange(dims[1], device=self.device),
            torch.arange(dims[2], device=self.device),
            torch.arange(dims[3], device=self.device),
            indexing='ij'
        ), dim=-1).float()  # (I, J, K, L, 4)
        # Guarda como buffer (não é parâmetro treinável, mas migra com .to())
        self.register_buffer("grid_coords", grid_coords.view(-1, 4))  # (I*J*K*L, 4)

        # Hit map para tracking de densidade BMU (para diagnóstico)
        self.register_buffer(
            "hit_map", torch.zeros(self._n_neurons(), dtype=torch.float32,
                                   device=self.device)
        )

        # Métricas do último fit (para inspeção)
        self.last_qe: float = 0.0
        self.last_te: float = 0.0
        self._fitted: bool = False

    def _n_neurons(self) -> int:
        return int(self.dims[0] * self.dims[1] * self.dims[2] * self.dims[3])

    def _flat_weights(self) -> torch.Tensor:
        """Retorna pesos como (N, input_dim) para operações matriciais."""
        return self.weights.view(-1, self.input_dim)

    def find_bmu(self, x: torch.Tensor) -> torch.Tensor:
        """Encontra BMU para cada vetor de entrada (busca matricial).

        Args:
            x: (B, input_dim) — batch de vetores 4D.

        Returns:
            flat_idx: (B,) — índice flatten do BMU no grid.
        """
        flat_w = self._flat_weights()  # (N, D)
        # Distância euclidiana quadrada: (B, N)
        # ||x - w||² = ||x||² - 2·x·wᵀ + ||w||²
        # (mas a forma direta é mais estável numericamente para 4D)
        diff = x.unsqueeze(1) - flat_w.unsqueeze(0)  # (B, N, D)
        dist_sq = torch.sum(diff * diff, dim=-1)  # (B, N)
        flat_idx = torch.argmin(dist_sq, dim=-1)  # (B,)
        return flat_idx

    def fit(
        self,
        data: torch.Tensor,
        epochs: int = 50,
        batch_size: int = 64,
    ) -> Dict[str, float]:
        """Treina o SOM 4D sobre o buffer de dados.

        Args:
            data: (N, 4) — buffer de vetores 4D.
            epochs: número de épocas de treinamento SOM.
            batch_size: tamanho do batch (vetorizado para velocidade).

        Returns:
            Dict com métricas finais: {'qe': float, 'te': float, 'final_lr': float,
            'final_sigma': float}

        V7-fix: TODO o loop de treino é executado dentro de `torch.no_grad()`.
        Isto é CRÍTICO porque:
          (1) `self.weights` é nn.Parameter (requires_grad=True);
          (2) Sem no_grad, cada iteração `flat_w = flat_w + curr_lr * update_mean`
              constrói um nó no grafo computacional;
          (3) Após 15 epochs × 8 batches = 120 iterações, o grafo acumula
              ~120 níveis de backward, consumindo centenas de MB de RAM
              (causa OOM-killed em ambiente com 4GB cgroup);
          (4) O grafo também faz `.item()` falhar em alguns pontos porque
              `mean()` de um tensor com grad_fn pode retornar um tensor
              com 0-dim mas `torch.sqrt(...).mean()` de um tensor que
              "requer grad" produz um tensor diferente (com grad_fn),
              e chamadas subsequentes de `.item()` podem disparar o erro
              "only one element tensors can be converted to Python scalars"
              em versões antigas do PyTorch.
        """
        if data.numel() == 0:
            logger.warning("[AdvancedSOMAugmenter.fit] data vazia — skip")
            return {"qe": 0.0, "te": 0.0, "final_lr": 0.0, "final_sigma": 0.0}

        # V7-fix: TODO o cálculo é no_grad. O SOM é um algoritmo NÃO-diferenciável
        # (usa argmin, scatter_add_) — não há sentido em manter grafo.
        with torch.no_grad():
            data = data.to(self.device).float()
            N = data.size(0)
            flat_w = self._flat_weights().detach().clone()  # (n_neurons, 4)
            n_neurons = flat_w.size(0)

            # Reset hit_map para este fit
            self.hit_map.zero_()

            curr_lr = float(self.lr_init)
            curr_sigma = float(self.sigma_init)

            for epoch in range(epochs):
                # Decaimento exponencial dos hiperparâmetros (Kohonen clássico)
                # α_t = α₀ · exp(-t/τ_α), τ_α = epochs (decai ao longo do treino)
                # σ_t = σ₀ · exp(-t/τ_σ), τ_σ = epochs/2 (decai mais rápido)
                curr_lr = self.lr_init * math.exp(-epoch / max(1, epochs))
                curr_sigma = self.sigma_init * math.exp(-epoch / max(1, epochs / 2.0))
                curr_sigma_sq = 2.0 * (curr_sigma ** 2) if curr_sigma > 0 else 1e-5

                # Shuffle dos dados (estocasticidade SGD)
                perm = torch.randperm(N, device=self.device)
                data_shuffled = data[perm]

                # Treina em batches (vetorizado)
                for start in range(0, N, batch_size):
                    end = min(start + batch_size, N)
                    batch = data_shuffled[start:end]  # (B, 4)
                    B = batch.size(0)
                    if B == 0:
                        continue

                    # 1. BMU search matricial: (B, n_neurons)
                    diff = batch.unsqueeze(1) - flat_w.unsqueeze(0)  # (B, n_neurons, 4)
                    dist_sq = torch.sum(diff * diff, dim=-1)  # (B, n_neurons)
                    bmu_flat = torch.argmin(dist_sq, dim=-1)  # (B,)

                    # Atualiza hit_map (tracking densidade)
                    self.hit_map.scatter_add_(
                        0, bmu_flat, torch.ones_like(bmu_flat, dtype=torch.float32)
                    )

                    # 2. Distância topológica no grid 4D: (B, n_neurons)
                    bmu_coords = self.grid_coords[bmu_flat]  # (B, 4)
                    grid_diff = self.grid_coords.unsqueeze(0) - bmu_coords.unsqueeze(1)
                    grid_dist_sq = torch.sum(grid_diff * grid_diff, dim=-1)

                    # 3. Vizinhança gaussiana 4D: (B, n_neurons)
                    neighborhood = torch.exp(-grid_dist_sq / curr_sigma_sq)

                    # 4. Atualização de pesos (Kohonen update):
                    # ΔW = α · Λ · (x - W), média sobre o batch
                    update = neighborhood.unsqueeze(-1) * (
                        batch.unsqueeze(1) - flat_w.unsqueeze(0)
                    )  # (B, n_neurons, 4)
                    update_mean = update.sum(dim=0) / float(B)  # (n_neurons, 4)
                    flat_w = flat_w + curr_lr * update_mean

                # Sincroniza de volta para o parâmetro nn.Parameter
                self.weights.copy_(flat_w.view(*self.dims, self.input_dim))

            # Computa métricas finais (já dentro do no_grad)
            diff = data.unsqueeze(1) - flat_w.unsqueeze(0)  # (N, n_neurons, 4)
            dist_sq = torch.sum(diff * diff, dim=-1)  # (N, n_neurons)
            bmu_flat = torch.argmin(dist_sq, dim=-1)  # (N,)
            bmu_dist = dist_sq.gather(1, bmu_flat.unsqueeze(1)).squeeze(1)
            # V7-fix: .mean() → float() direto. bmu_dist é (N,) e mean() é 0-dim.
            # Usar .item() diretamente evita qualquer ambiguidade de tipo.
            qe_val_tensor = torch.sqrt(bmu_dist + 1e-12).mean()
            qe = float(qe_val_tensor.item()) if qe_val_tensor.numel() == 1 else 0.0

            # Topological Error: % de amostras onde BMU e 2nd-BMU não são vizinhos
            if N >= 2 and n_neurons >= 2:
                top2 = torch.topk(dist_sq, k=2, largest=False, dim=-1)  # (N, 2)
                bmu1 = top2.indices[:, 0]
                bmu2 = top2.indices[:, 1]
                bmu1_coords = self.grid_coords[bmu1]  # (N, 4)
                bmu2_coords = self.grid_coords[bmu2]  # (N, 4)
                grid_diff = bmu2_coords - bmu1_coords  # (N, 4)
                manhattan = grid_diff.abs().sum(dim=-1)  # (N,)
                are_neighbors = (manhattan == 1).float()  # 6-connectivity em 4D
                te_val_tensor = 1.0 - are_neighbors.mean()
                te = float(te_val_tensor.item()) if te_val_tensor.numel() == 1 else 0.0
            else:
                te = 0.0

        self.last_qe = qe
        self.last_te = te
        self._fitted = True

        return {
            "qe": qe,
            "te": te,
            "final_lr": float(curr_lr),
            "final_sigma": float(curr_sigma),
            "n_neurons": int(n_neurons),
            "n_samples_trained": int(N),
            "epochs": int(epochs),
        }

    def augment(
        self,
        data: torch.Tensor,
        num_samples: int = 10,
        noise_scale: float = 0.02,
        interp_real: float = 0.7,
        interp_bmu: float = 0.3,
    ) -> torch.Tensor:
        """Gera amostras sintéticas manifold-aware via interpolação convexa
        + ruído gaussiano adaptativo por QE local.

        Args:
            data: (N, 4) — buffer de vetores 4D reais.
            num_samples: número de amostras sintéticas a gerar.
            noise_scale: σ_base do ruído. σ_adaptive = noise_scale * (1 + QE_local).
            interp_real: peso da amostra real na interpolação (default 0.7).
            interp_bmu: peso do vetor BMU na interpolação (default 0.3).
                interp_real + interp_bmu deve ser ≈ 1.0.

        Returns:
            synthetic: (num_samples, 4) — amostras sintéticas.
        """
        if not self._fitted:
            logger.warning(
                "[AdvancedSOMAugmenter.augment] SOM não treinado — "
                "retornando ruído gaussiano puro (sem garantia de manifold)."
            )
            return torch.randn(num_samples, self.input_dim, device=self.device) * 0.1

        if data.numel() == 0:
            logger.warning("[AdvancedSOMAugmenter.augment] data vazia — skip")
            return torch.empty(0, self.input_dim, device=self.device)

        data = data.to(self.device).float()
        N = data.size(0)
        flat_w = self._flat_weights()  # (n_neurons, 4)

        # Normaliza pesos de interpolação
        total_interp = interp_real + interp_bmu
        if total_interp <= 0:
            interp_real, interp_bmu = 0.7, 0.3
        else:
            interp_real /= total_interp
            interp_bmu /= total_interp

        synthetic_samples = []
        for _ in range(num_samples):
            # 1. Seleciona amostra real aleatória como âncora
            anchor_idx = torch.randint(0, N, (1,), device=self.device)
            target_sample = data[anchor_idx[0]]  # (4,)

            # 2. Encontra BMU(target_sample) — busca matricial
            diff = flat_w - target_sample.unsqueeze(0)  # (n_neurons, 4)
            dist_sq = torch.sum(diff * diff, dim=-1)  # (n_neurons,)
            bmu_flat = torch.argmin(dist_sq).item()
            bmu_weight = flat_w[bmu_flat]  # (4,)

            # 3. Computa Quantization Error local
            quantization_error = float(
                torch.sqrt(dist_sq[bmu_flat] + 1e-12).item()
            )

            # 4. Ruído gaussiano adaptativo proporcional ao QE
            # σ_adaptive = noise_scale * (1 + QE)
            # — quando QE alto (região mal representada): mais ruído (exploração)
            # — quando QE baixo (região bem representada): menos ruído (fidelidade)
            sigma_adaptive = noise_scale * (1.0 + quantization_error)
            adaptive_noise = torch.randn(self.input_dim, device=self.device) * sigma_adaptive

            # 5. Amostra sintética = interpolação convexa + ruído adaptativo
            synthetic_sample = (
                interp_real * target_sample
                + interp_bmu * bmu_weight
                + adaptive_noise
            )
            synthetic_samples.append(synthetic_sample)

        return torch.stack(synthetic_samples, dim=0)  # (num_samples, 4)

    def get_hit_map_density(self) -> torch.Tensor:
        """Retorna densidade do hit_map normalizada (para diagnóstico).

        Returns:
            density: (n_neurons,) — densidade BMU por neurônio (soma 1).
        """
        total = self.hit_map.sum()
        if total > 0:
            return self.hit_map / total
        return self.hit_map

    def get_dead_neuron_rate(self) -> float:
        """Calcula taxa de neurônios mortos (nunca foram BMU no último fit).

        Returns:
            dead_rate: float ∈ [0, 1] — fração de neurônios sem hits.
        """
        if self.hit_map.sum() == 0:
            return 1.0
        n_dead = int((self.hit_map == 0).sum().item())
        return float(n_dead / self._n_neurons())


__all__ = [
    "DataAugmenter",
    "AugmentationConfig",
    "AdvancedSOMAugmenter",
]