src/bigru_t/reasoning/cyclic_reasoning.py

"""
cyclic_reasoning — Mecanismo de raciocínio cíclico melhorado.

Baseado no código GLMCyclicMultimodalAgent fornecido, com as seguintes melhorias:

1. CyclicStateController com Kalman filter + fixed-point convergence
2. CyclicMemoryManager com compressão atencional para 2M tokens
3. GLMCyclicMultimodalAgent integrando LLM + memória + controller + tool agent
4. Detecção de falhas: não-convergência, overflow, ciclos oscilantes

Melhorias sobre o código original:
- Convergence check robusto (norma relativa, não absoluta)
- Max cycle limit com fallback graceful
- Memory compression com attention real (não apenas mean pooling)
- State controller com GRUCell + Kalman gain aprendido
- Tool action learning via gradient (não random walk)
- Detecção de oscilação (estado divergente)
- Logging estruturado para análise pós-treino
"""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple, Any

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

# v12.9: RMSNorm importado do vqvae2_hierarchical (reutilização)
from .vqvae2_hierarchical import RMSNorm


@dataclass
class CyclicReasoningConfig:
    """Configuração do sistema de raciocínio cíclico."""
    hidden_size: int = 256
    tool_dim: int = 64
    max_cycles: int = 10
    convergence_threshold: float = 0.01
    # Detecção de falhas
    divergence_threshold: float = 10.0  # Se ||error|| > 10× inicial, aborta
    oscillation_window: int = 3         # Janela para detectar oscilação
    oscillation_tol: float = 0.05       # Tolerância para considerar oscilante
    # Memória
    memory_compress_threshold: int = 512  # Tokens antes de comprimir
    memory_window: int = 256              # Tamanho do contexto recente mantido
    # Transformer (configurável — conforme solicitado v12.5)
    num_transformer_layers: int = 8       # Era hard-coded 2; agora configurável
    transformer_nhead: int = 8            # Número de cabeças da attention
    transformer_dim_ff: int = 2048        # dim_feedforward (default 4× hidden)
    # v12.9: Normalização configurável
    # "layernorm": LayerNorm no transformer + LayerNorm entre camadas (v12.6)
    # "rmsnorm": RMSNorm no transformer + RMSNorm entre camadas
    # "hybrid": RMSNorm nos MLPs (tool_generator, kalman) + LayerNorm no transformer
    transformer_norm_type: str = "layernorm"  # "layernorm" | "rmsnorm"
    mlp_norm_type: str = "none"  # "none" | "rmsnorm" (LayerNorm nos MLPs não testado)
    # Logging
    log_every_cycle: bool = True


class CyclicStateController(nn.Module):
    """Controlador cíclico com filtro de Kalman e convergência por ponto fixo.

    Matemática:
        S_hat_{t+1} = f(S_t, u_t)         # Predição (GRUCell)
        e_t = y_t - H * S_hat             # Inovação (erro)
        K_t = MLP([S_hat, e_t])            # Ganho de Kalman aprendido
        S_{t+1} = S_hat + K_t * e_t       # Atualização
        converge se ||e_t|| < epsilon     # Ponto fixo
    """

    def __init__(self, hidden_size: int, tool_dim: int, mlp_norm_type: str = "none"):
        super().__init__()
        self.hidden_size = hidden_size
        self.tool_dim = tool_dim
        self.mlp_norm_type = mlp_norm_type
        # Preditor do próximo estado (GRU recorrente).
        self.state_transition = nn.GRUCell(hidden_size, hidden_size)
        # Projeção da ação da ferramenta para o espaço do estado.
        self.action_proj = nn.Linear(tool_dim, hidden_size)
        # Projeção do estado predito para o espaço de observação.
        self.obs_proj = nn.Linear(hidden_size, hidden_size)
        # v12.9: MLP do Kalman gain com normalização opcional.
        if mlp_norm_type == "rmsnorm":
            self.kalman_gain_net = nn.Sequential(
                nn.Linear(hidden_size * 2, hidden_size),
                RMSNorm(hidden_size),
                nn.ReLU(),
                nn.Linear(hidden_size, hidden_size),
                nn.Sigmoid(),  # K_t ∈ [0, 1]
            )
        else:
            self.kalman_gain_net = nn.Sequential(
                nn.Linear(hidden_size * 2, hidden_size),
                nn.ReLU(),
                nn.Linear(hidden_size, hidden_size),
                nn.Sigmoid(),  # K_t ∈ [0, 1]
            )

    def forward(
        self,
        current_state: torch.Tensor,
        tool_action: torch.Tensor,
        observed_state: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, float]]:
        """
        Args:
            current_state: (B, H) estado atual.
            tool_action: (B, T) ação da ferramenta.
            observed_state: (B, H) estado observado.

        Retorna:
            adjusted_state: (B, H) próximo estado ajustado.
            convergence_metric: (B, 1) norma L2 do erro.
            diagnostics: dict com métricas para logging.
        """
        # 1. Predição do próximo estado baseado na ação.
        action_embed = self.action_proj(tool_action)
        predicted_state = self.state_transition(current_state, action_embed)

        # 2. Inovação (erro): e_t = y_t - H * S_hat.
        H_predicted = self.obs_proj(predicted_state)
        error = observed_state - H_predicted

        # 3. Ganho de Kalman K_t ∈ [0, 1]^H (sigmoid).
        K = self.kalman_gain_net(torch.cat([predicted_state, error], dim=-1))
        # v13.9: Fisher Clipping no Kalman gain (reduz mal-conditioning).
        # K pode ter dimensões ≈ 0 (baixa confiança) e ≈ 1 (alta confiança),
        # criando condition number extremo. Fisher Clipping comprime o range.
        from .mathematical_optimizations import fisher_clip_log
        K = fisher_clip_log(K, eps=1e-8)

        # 4. Atualização do estado.
        adjusted_state = predicted_state + K * error

        # Norma do erro (para checar convergência).
        convergence_metric = torch.norm(error, p=2, dim=-1, keepdim=True)

        # Diagnósticos.
        with torch.no_grad():
            diagnostics = {
                'error_norm_mean': float(convergence_metric.mean().item()),
                'error_norm_max': float(convergence_metric.max().item()),
                'kalman_gain_mean': float(K.mean().item()),
                'state_norm': float(torch.norm(adjusted_state, p=2, dim=-1).mean().item()),
                'innovation_cosine': float(F.cosine_similarity(
                    observed_state, predicted_state, dim=-1
                ).mean().item()),
            }

        return adjusted_state, convergence_metric, diagnostics


class CyclicMemoryManager(nn.Module):
    """Gestor de memória cíclica com compressão atencional.

    Estratégia:
    - Mantém últimos `memory_window` tokens sem compressão (contexto recente).
    - Quando o histórico excede `memory_compress_threshold`, comprime o histórico
      antigo em um único vetor latente via cross-attention.
    - Isto permite suportar contextos longos (até 2M tokens teoricamente)
      sem estourar a memória.
    """

    def __init__(self, config: CyclicReasoningConfig):
        super().__init__()
        self.config = config
        self.attention_compressor = nn.MultiheadAttention(
            embed_dim=config.hidden_size,
            num_heads=8,  # fixo para simplicidade
            batch_first=True,
        )
        # LayerNorm para estabilizar após compressão.
        self.norm = nn.LayerNorm(config.hidden_size)

    def forward(
        self,
        history_embeddings: torch.Tensor,
        current_cycle_embeddings: torch.Tensor,
    ) -> Tuple[torch.Tensor, Dict[str, int]]:
        """
        Args:
            history_embeddings: (B, T_hist, H) histórico acumulado.
            current_cycle_embeddings: (B, T_curr, H) embeddings do ciclo atual.

        Retorna:
            context_embeds: (B, T_out, H) contexto comprimido.
            stats: dict com tamanhos antes/depois.
        """
        # Concatena histórico com ciclo atual.
        full_context = torch.cat([history_embeddings, current_cycle_embeddings], dim=1)
        T_full = full_context.size(1)
        stats = {'tokens_before': T_full}

        if T_full > self.config.memory_compress_threshold:
            # Comprime histórico antigo em um único vetor latente.
            # Query: últimos `memory_window` tokens.
            # Key/Value: todos os tokens anteriores ao `memory_window`.
            window = self.config.memory_window
            query = full_context[:, -window:, :]
            kv = full_context[:, :-window, :]
            compressed, _ = self.attention_compressor(query, kv, kv)
            # Pool do comprimido para 1 token + manter janela recente.
            compressed_pooled = compressed.mean(dim=1, keepdim=True)
            compressed_pooled = self.norm(compressed_pooled)
            output = torch.cat([compressed_pooled, full_context[:, -window:, :]], dim=1)
            stats['tokens_after'] = output.size(1)
            stats['compression_ratio'] = stats['tokens_before'] / max(stats['tokens_after'], 1)
        else:
            output = full_context
            stats['tokens_after'] = T_full
            stats['compression_ratio'] = 1.0

        return output, stats


class GLMCyclicReasoningAgent(nn.Module):
    """Agente de raciocínio cíclico GLM com LLM + memória + controller + tool.

    Melhorias sobre o código original:
    - Convergence check robusto (norma relativa)
    - Detecção de divergência e oscilação
    - Logging estruturado por ciclo
    - Tool action via MLP (não random walk)
    - Fallback graceful quando não converge
    """

    def __init__(
        self,
        config: CyclicReasoningConfig,
        vocab_size: int = 1000,
        image_feature_dim: int = 256,
    ):
        super().__init__()
        self.config = config
        self.vocab_size = vocab_size
        self.image_feature_dim = image_feature_dim

        # v12.9: Helper para criar módulos de normalização conforme config.
        norm_type_tf = config.transformer_norm_type  # "layernorm" | "rmsnorm"
        norm_type_mlp = config.mlp_norm_type          # "none" | "rmsnorm"
        use_rmsnorm_tf = (norm_type_tf == "rmsnorm")
        use_rmsnorm_mlp = (norm_type_mlp == "rmsnorm")

        def make_norm(dim):
            return RMSNorm(dim) if use_rmsnorm_tf else nn.LayerNorm(dim)

        # Embeddings.
        self.text_embed = nn.Embedding(vocab_size, config.hidden_size)
        self.image_proj = nn.Linear(image_feature_dim, config.hidden_size)
        # v12.9: Norm após embedding (LayerNorm ou RMSNorm conforme config).
        self.input_norm = make_norm(config.hidden_size)

        # LLM base (Transformer Encoder com norm entre camadas).
        dim_ff = config.transformer_dim_ff if config.transformer_dim_ff > 0 else config.hidden_size * 4
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=config.hidden_size,
            nhead=config.transformer_nhead,
            dim_feedforward=dim_ff,
            batch_first=True,
            activation='gelu',
            norm_first=True,  # Pre-LN (mais estável com muitas camadas)
        )
        # v12.9: Substituir LayerNorms internos do encoder_layer por RMSNorm se configurado.
        if use_rmsnorm_tf:
            encoder_layer.norm1 = RMSNorm(config.hidden_size)
            encoder_layer.norm2 = RMSNorm(config.hidden_size)
        # v13.10: Aplicar Fisher Clipping em attention weights do transformer.
        # Registra hook no self_attn para clipping dos scores antes do softmax.
        self._clip_attention = True
        original_self_attn = encoder_layer.self_attn
        
        class ClippedAttention(nn.Module):
            """Wrapper que aplica Fisher Clipping nos attention scores."""
            def __init__(self, attn_module):
                super().__init__()
                self.attn = attn_module
                
            def forward(self, query, key, value, **kwargs):
                # Computar scores manualmente para aplicar clipping.
                # nn.MultiheadAttention faz Q@K^T internamente.
                # Em vez de interceptar, usamos abordagem via need_weights.
                kwargs.setdefault('need_weights', False)
                return self.attn(query, key, value, **kwargs)
        
        # Nota: Fisher Clipping em attention scores requer acesso aos scores
        # antes do softmax. nn.MultiheadAttention não expõe isso facilmente.
        # Alternativa: aplicar clipping via post-hoc no output do transformer.
        # Para implementação completa, seria necessário custom MultiheadAttention.
        # Por agora, registramos que o clipping está disponível via clip_attention_weights()
        # e pode ser aplicado em implementações custom de attention.
        # Camadas individuais com norm entre elas.
        layers = []
        for i in range(config.num_transformer_layers):
            layers.append(encoder_layer)
            if i < config.num_transformer_layers - 1:
                layers.append(make_norm(config.hidden_size))
        self.transformer = nn.Sequential(*layers)
        # Norm final após transformer.
        self.output_norm = make_norm(config.hidden_size)

        # Componentes cíclicos.
        self.memory_manager = CyclicMemoryManager(config)
        self.state_controller = CyclicStateController(
            config.hidden_size, config.tool_dim,
            mlp_norm_type=norm_type_mlp,
        )

        # v12.9: Tool action generator com normalização opcional nos MLPs.
        if use_rmsnorm_mlp:
            self.tool_generator = nn.Sequential(
                nn.Linear(config.hidden_size, config.hidden_size),
                RMSNorm(config.hidden_size),
                nn.GELU(),
                nn.Linear(config.hidden_size, config.tool_dim),
            )
        else:
            self.tool_generator = nn.Sequential(
                nn.Linear(config.hidden_size, config.hidden_size),
                nn.GELU(),
                nn.Linear(config.hidden_size, config.tool_dim),
            )

        # LM head.
        self.lm_head = nn.Linear(config.hidden_size, vocab_size)

    def run_inference_loop(
        self,
        initial_input_ids: torch.Tensor,
        image_features: Optional[torch.Tensor] = None,
        initial_tool_action: Optional[torch.Tensor] = None,
        return_history: bool = True,
    ) -> Dict[str, Any]:
        """
        Executa o sistema cíclico até convergência.

        Args:
            initial_input_ids: (B, T) tokens de entrada.
            image_features: (B, T_img, D_img) features de imagem (opcional).
            initial_tool_action: (B, tool_dim) ação inicial (opcional).
            return_history: se True, retorna histórico completo de ciclos.

        Retorna:
            dict com:
                final_logits: (B, T, V) logits finais.
                final_state: (B, H) estado final do agente.
                converged: bool — se convergiu dentro de max_cycles.
                cycles_run: int — número de ciclos executados.
                final_error: float — norma final do erro.
                history: lista de dicts com logs por ciclo.
                failure_mode: str | None — 'diverged', 'oscillating', 'max_cycles', None.
        """
        batch_size = initial_input_ids.size(0)
        device = initial_input_ids.device

        # Inicialização.
        inputs_embeds = self.text_embed(initial_input_ids)
        if image_features is not None:
            img_embeds = self.image_proj(image_features)
            inputs_embeds = torch.cat([img_embeds, inputs_embeds], dim=1)

        # Estado inicial do agente (resumo do problema).
        agent_state = inputs_embeds.mean(dim=1)
        history_embeds = inputs_embeds

        # Tool action inicial.
        if initial_tool_action is None:
            tool_action = torch.zeros(batch_size, self.config.tool_dim, device=device)
        else:
            tool_action = initial_tool_action

        # Tracking para detecção de falhas.
        error_history: List[float] = []
        history: List[Dict[str, Any]] = []
        failure_mode = None
        converged = False
        cycles_run = 0
        initial_error = None

        for cycle in range(self.config.max_cycles):
            cycles_run = cycle + 1

            # 1. Memória (evita estouro de contexto).
            context_embeds, mem_stats = self.memory_manager(
                history_embeds, history_embeds
            )

            # 2. Processamento LLM (com input_norm e output_norm).
            # v13.4: Sanitização NaN/Inf antes do transformer.
            context_embeds = torch.nan_to_num(context_embeds, nan=0.0, posinf=1e4, neginf=-1e4)
            context_embeds = self.input_norm(context_embeds)
            hidden_states = self.transformer(context_embeds)
            # v13.4: Sanitização NaN/Inf após transformer (attention pode produzir extremos).
            hidden_states = torch.nan_to_num(hidden_states, nan=0.0, posinf=1e4, neginf=-1e4)
            hidden_states = self.output_norm(hidden_states)

            # 3. Observação: em sistema real viria da execução da tool.
            # Aqui simulamos convergência gradual para o estado predito.
            noise_scale = max(0.5 / (cycle + 1), 0.01)
            observed_state = hidden_states.mean(dim=1) + noise_scale * torch.randn_like(agent_state)

            # 4. Atualização do estado via Kalman.
            agent_state, convergence_metric, diag = self.state_controller(
                agent_state, tool_action, observed_state
            )

            error_norm = float(convergence_metric.mean().item())
            error_history.append(error_norm)
            if initial_error is None:
                initial_error = max(error_norm, 1e-6)

            # 5. Geração de próxima ação/tool.
            tool_action = self.tool_generator(agent_state)

            # 6. Geração de próximo token.
            logits = self.lm_head(hidden_states)
            next_token_id = torch.argmax(logits[:, -1, :], dim=-1)

            # Log do ciclo.
            cycle_log = {
                'cycle': cycle + 1,
                'token_generated': int(next_token_id[0].item()),
                'error_norm': error_norm,
                'relative_error': error_norm / initial_error,
                'kalman_gain': diag['kalman_gain_mean'],
                'innovation_cosine': diag['innovation_cosine'],
                'mem_tokens_before': mem_stats['tokens_before'],
                'mem_tokens_after': mem_stats['tokens_after'],
                'mem_compression_ratio': mem_stats['compression_ratio'],
            }
            history.append(cycle_log)

            if self.config.log_every_cycle:
                print(f"  Cycle {cycle+1}: error={error_norm:.4f} "
                      f"(rel={error_norm/initial_error:.3f}) "
                      f"K={diag['kalman_gain_mean']:.3f} "
                      f"token={next_token_id[0].item()}")

            # 7. Condição de parada (ponto fixo).
            if error_norm < self.config.convergence_threshold:
                converged = True
                print(f">>> Convergência atingida no ciclo {cycle+1}!")
                break

            # 8. Detecção de divergência (erro crescendo muito).
            if error_norm > self.config.divergence_threshold * initial_error:
                failure_mode = 'diverged'
                print(f"!!! Divergência detectada no ciclo {cycle+1} "
                      f"(error={error_norm:.4f} > {self.config.divergence_threshold * initial_error:.4f})")
                break

            # 9. Detecção de oscilação (variância alta nos últimos N ciclos).
            if len(error_history) >= self.config.oscillation_window:
                window = error_history[-self.config.oscillation_window:]
                window_std = float(torch.tensor(window).std().item())
                window_mean = float(torch.tensor(window).mean().item())
                if window_std > self.config.oscillation_tol and window_std / max(window_mean, 1e-6) > 0.3:
                    failure_mode = 'oscillating'
                    print(f"!!! Oscilação detectada no ciclo {cycle+1} "
                          f"(std={window_std:.4f}, mean={window_mean:.4f})")
                    break

            # Prepara próximo ciclo: adiciona token gerado ao histórico.
            next_token_embed = self.text_embed(next_token_id.unsqueeze(-1))
            history_embeds = torch.cat([history_embeds, next_token_embed], dim=1)
            # v13.1: limpeza leve entre ciclos (libera tensores intermediários).
            # Only gc.collect every 50 cycles, not every cycle (very expensive).
            if (cycle + 1) % 50 == 0:
                import gc as _gc
                _gc.collect(0)

        if not converged and failure_mode is None:
            failure_mode = 'max_cycles'
            print(f">>> Atingiu max_cycles={self.config.max_cycles} sem convergir.")

        return {
            'final_logits': logits,
            'final_state': agent_state,
            'converged': converged,
            'cycles_run': cycles_run,
            'final_error': error_history[-1] if error_history else float('inf'),
            'initial_error': initial_error,
            'history': history if return_history else [],
            'failure_mode': failure_mode,
            'error_history': error_history,
        }


class ReasoningDiagnostics:
    """Análise pós-execução do loop de raciocínio cíclico.

    Detecta padrões de falha e sugere melhorias.
    """

    @staticmethod
    def analyze(result: Dict[str, Any]) -> Dict[str, Any]:
        """Analisa resultado de run_inference_loop e retorna diagnóstico."""
        error_hist = result.get('error_history', [])
        history = result.get('history', [])

        analysis = {
            'converged': result['converged'],
            'cycles_run': result['cycles_run'],
            'failure_mode': result['failure_mode'],
            'initial_error': result.get('initial_error'),
            'final_error': result.get('final_error'),
            'error_reduction_pct': 0.0,
            'avg_kalman_gain': 0.0,
            'avg_innovation_cosine': 0.0,
            'mem_compression_used': False,
            'avg_compression_ratio': 1.0,
            'error_trend': 'unknown',
            'recommendations': [],
        }

        if error_hist and result.get('initial_error'):
            init = result['initial_error']
            final = result['final_error']
            if init > 0:
                analysis['error_reduction_pct'] = (init - final) / init * 100

        if history:
            analysis['avg_kalman_gain'] = sum(h['kalman_gain'] for h in history) / len(history)
            analysis['avg_innovation_cosine'] = sum(h['innovation_cosine'] for h in history) / len(history)
            compressions = [h['mem_compression_ratio'] for h in history if h['mem_compression_ratio'] > 1.0]
            if compressions:
                analysis['mem_compression_used'] = True
                analysis['avg_compression_ratio'] = sum(compressions) / len(compressions)

            # Tendência do erro.
            if len(error_hist) >= 2:
                first_half = error_hist[:len(error_hist)//2]
                second_half = error_hist[len(error_hist)//2:]
                if sum(second_half) < sum(first_half):
                    analysis['error_trend'] = 'decreasing'
                elif sum(second_half) > sum(first_half):
                    analysis['error_trend'] = 'increasing'
                else:
                    analysis['error_trend'] = 'stable'

        # Recomendações.
        if not analysis['converged']:
            if analysis['failure_mode'] == 'diverged':
                analysis['recommendations'].append(
                    "Reduzir learning rate do state_controller ou aumentar divergence_threshold"
                )
            elif analysis['failure_mode'] == 'oscillating':
                analysis['recommendations'].append(
                    "Aumentar patience para oscilação ou suavizar Kalman gain (maior sigmoid)"
                )
            elif analysis['failure_mode'] == 'max_cycles':
                analysis['recommendations'].append(
                    "Aumentar max_cycles ou reduzir convergence_threshold"
                )

        if analysis['avg_kalman_gain'] > 0.8:
            analysis['recommendations'].append(
                "Kalman gain muito alto — estado super-reativo às observações"
            )
        elif analysis['avg_kalman_gain'] < 0.1:
            analysis['recommendations'].append(
                "Kalman gain muito baixo — estado não está incorporando observações"
            )

        if analysis['avg_innovation_cosine'] < 0.0:
            analysis['recommendations'].append(
                "Inovação e predição anti-correlacionadas — revisar obs_proj"
            )

        return analysis