"""generator_verifier.py — Gerador + Verificador + Anti-Alucinação.

Implementa o pseudocódigo da seção 5-7 de dados.txt:

    - Generator: CNN-BiGRU encoder + GRU decoder com atenção Bahdanau
    - Verifier: CNN-BiGRU + classificador sigmoid (v ∈ [0,1])
    - Anti-hallucination layer: ativadores lógicos fuzzy de Łukasiewicz
      NOT(x)  = 1 - sigmoid(x)
      AND(x,y) = relu(sigmoid(x) + sigmoid(y) - 1)
      OR(x,y)  = min(1, sigmoid(x) + sigmoid(y))
      IMP(x,y) = min(1, 1 - sigmoid(x) + sigmoid(y))
"""
from __future__ import annotations

import logging
import math
from typing import Optional, Tuple

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

from .cooperative_bigru import CooperativeCNNBiGRU

logger = logging.getLogger(__name__)


# ============================================================================
# Funções fuzzy de Łukasiewicz
# ============================================================================

def fuzzy_not(x: torch.Tensor) -> torch.Tensor:
    """NOT(x) = 1 - sigmoid(x)."""
    return 1.0 - torch.sigmoid(x)


def fuzzy_and(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """AND(x, y) = relu(sigmoid(x) + sigmoid(y) - 1)."""
    s = torch.sigmoid(x) + torch.sigmoid(y)
    return F.relu(s - 1.0)


def fuzzy_or(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """OR(x, y) = min(1, sigmoid(x) + sigmoid(y))."""
    s = torch.sigmoid(x) + torch.sigmoid(y)
    return torch.clamp(s, 0.0, 1.0)


def fuzzy_imp(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """IMP(x, y) = min(1, 1 - sigmoid(x) + sigmoid(y))."""
    s = 1.0 - torch.sigmoid(x) + torch.sigmoid(y)
    return torch.clamp(s, 0.0, 1.0)


# ============================================================================
# Camada Anti-Alucinação
# ============================================================================

class AntiHallucinationLayer(nn.Module):
    """Camada anti-alucinação com lógica fuzzy de Łukasiewicz.

    Recebe os logits [B, V] de um passo gerado e calcula um valor de verdade
    suave τ ∈ [0, 1] que indica o quão "lógico" é o passo.

    Implementação diferenciável:
        - Calcula sigmoid(logit) -> probabilidade soft
        - Aplica operadores fuzzy em pares de tokens mais prováveis
        - Combina resultados via IMP (implicação) e AND
    """

    def __init__(self, vocab_size: int, embed_dim: int = 32):
        super().__init__()
        self.proj = nn.Linear(vocab_size, embed_dim)
        self.combine = nn.Linear(embed_dim * 2, 1)

    def forward(self, logits: torch.Tensor) -> torch.Tensor:
        """
        Args:
            logits: [B, V] logits do passo gerado

        Returns:
            tau: [B, 1] valor de verdade suave em [0, 1]
        """
        # Probabilidades soft via sigmoid
        probs = torch.sigmoid(logits)  # [B, V]

        # Top-2 tokens mais prováveis (diferenciável via topk)
        top2_vals, top2_idx = torch.topk(probs, k=min(2, probs.size(-1)), dim=-1)  # [B, 2]

        # Logits dos top-2 tokens via gather
        top1_logit = logits.gather(1, top2_idx[:, 0:1])  # [B, 1]
        top2_logit = logits.gather(1, top2_idx[:, 1:2])  # [B, 1]

        # Operadores fuzzy de Łukasiewicz sobre os top-2 logits
        imp_val = fuzzy_imp(top1_logit, top2_logit)  # [B, 1]
        and_val = fuzzy_and(top1_logit, top2_logit)  # [B, 1]
        or_val = fuzzy_or(top1_logit, top2_logit)    # [B, 1]

        # Projeção do vetor de logits completo para embedding
        proj = self.proj(logits)  # [B, embed_dim]

        # Modulariza a projeção pelos valores fuzzy (broadcasting)
        fuzzy_factor = (imp_val + and_val + or_val) / 3.0  # [B, 1]
        modulated = proj * fuzzy_factor  # [B, embed_dim]

        # Combina: [modulated, original_proj] -> [B, 2*embed_dim]
        combined = torch.cat((modulated, proj), dim=-1)  # [B, 2*embed_dim]
        tau = torch.sigmoid(self.combine(combined))  # [B, 1]
        return tau


# ============================================================================
# Gerador (CNN-BiGRU encoder + GRU decoder com atenção)
# ============================================================================

class GeneratorCNNBiGRU(nn.Module):
    """Gerador: encoder CNN-BiGRU + decoder GRU com atenção Bahdanau.

    Args:
        vocab_size, embedding_dim, cnn_filters, gru_hidden: herdados do núcleo.
        max_proof_len: comprimento máximo da sequência gerada.
    """

    def __init__(
        self,
        vocab_size: int,
        embedding_dim: int = 64,
        cnn_filters: int = 64,
        gru_hidden: int = 64,
        n_heads: int = 4,
        dropout: float = 0.1,
        pad_idx: int = 1,
        max_proof_len: int = 32,
    ):
        super().__init__()
        self.pad_idx = pad_idx
        self.vocab_size = vocab_size
        self.gru_hidden = gru_hidden
        self.max_proof_len = max_proof_len

        # Encoder cooperativo
        self.encoder = CooperativeCNNBiGRU(
            vocab_size=vocab_size,
            embedding_dim=embedding_dim,
            cnn_filters=cnn_filters,
            gru_hidden=gru_hidden,
            n_heads=n_heads,
            dropout=dropout,
            pad_idx=pad_idx,
        )

        # Decoder GRU (unidirecional)
        self.dec_embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=pad_idx)
        self.dec_gru = nn.GRUCell(embedding_dim + gru_hidden * 2, gru_hidden)

        # Atenção Bahdanau
        self.attn_W = nn.Linear(gru_hidden * 2, gru_hidden, bias=False)
        self.attn_U = nn.Linear(gru_hidden, gru_hidden, bias=False)
        self.attn_v = nn.Linear(gru_hidden, 1, bias=False)

        # Projeção final
        self.out_proj = nn.Linear(gru_hidden + gru_hidden * 2 + embedding_dim, vocab_size)

        self.bos_id = 0  # placeholder, será sobrescrito
        self.eos_id = 2

    def _attention(
        self,
        dec_state: torch.Tensor,  # [B, H]
        enc_out: torch.Tensor,    # [B, T, 2H]
        enc_mask: torch.Tensor,   # [B, T]
    ) -> torch.Tensor:
        """Atenção Bahdanau: retorna contexto [B, 2H]."""
        # score = v^T tanh(W enc + U dec)
        proj_enc = self.attn_W(enc_out)              # [B, T, H]
        proj_dec = self.attn_U(dec_state).unsqueeze(1)  # [B, 1, H]
        scores = self.attn_v(torch.tanh(proj_enc + proj_dec)).squeeze(-1)  # [B, T]
        scores = scores.masked_fill(enc_mask == 0, -1e9)
        attn = F.softmax(scores, dim=-1)             # [B, T]
        ctx = torch.bmm(attn.unsqueeze(1), enc_out).squeeze(1)  # [B, 2H]
        return ctx

    def forward(
        self,
        input_ids_a: torch.Tensor,
        input_ids_b: torch.Tensor,
        target_proof: Optional[torch.Tensor] = None,
        bos_id: int = 0,
        eos_id: int = 2,
        max_len: Optional[int] = None,
    ) -> dict:
        """
        Args:
            input_ids_a, input_ids_b: [B, T] entradas (axiomas + conjectura)
            target_proof: [B, T_proof] para teacher forcing (treino)
            bos_id, eos_id: IDs dos tokens especiais

        Returns:
            dict com:
                logits: [B, T_proof, V]
                enc_out: [B, T, 2H] — sequência REAL do encoder (não repetida)
                dec_state: [B, H]

        CORREÇÃO: O bug original era `enc_out = fused.unsqueeze(1).expand(-1, T_enc, -1)`
        que repetia o mesmo vetor a cada timestep, tornando a Bahdanau attention
        meaningless (todas as chaves idênticas → atenção uniforme).
        Agora usamos a sequência temporal REAL produzida pelo encoder.
        """
        # 1. Encoder — retornar sequências temporais para a atenção
        enc = self.encoder(input_ids_a, input_ids_b, return_sequences=True)

        # Sequência temporal real do stream A: [B, T_a, 2H]
        # (resultado do self-attention summary que preserva a dimensão 2H)
        # Mas self_attn_A retorna [B, 2H] (sumarizado), então usamos seq_A
        # que contém os estados ocultos em cada timestep
        if "seq_A" in enc:
            # seq_A: [B, T, 2H] — estados BiGRU concatenados por timestep
            enc_out = enc["seq_A"]  # [B, T, 2H]
        else:
            # Fallback: usar out_A repetido (mas agora com aviso)
            logger.warning("GeneratorCNNBiGRU: seq_A não disponível, usando fallback")
            out_A = enc["out_A"]  # [B, 2H]
            T_enc = enc["mask_A"].size(1)
            enc_out = out_A.unsqueeze(1).expand(-1, T_enc, -1)

        mask_A = enc["mask_A"]
        enc_mask = mask_A

        # Fused para inicializar o decoder
        fused = enc["fused"]  # [B, 256] = [B, 4H] (2 streams * 2H)
        B = input_ids_a.size(0)
        device = input_ids_a.device

        # 2. Inicialização do decoder
        # Estado inicial: projeção linear do fused para gru_hidden
        # (em vez de cortar a primeira metade, que perde metade da informação)
        if not hasattr(self, "_init_proj"):
            # Lazy init — evita quebrar a API existente
            gru_in = fused.size(-1)  # 256 = 4H
            self._init_proj = nn.Linear(gru_in, self.gru_hidden).to(device)
            nn.init.xavier_uniform_(self._init_proj.weight)
            nn.init.zeros_(self._init_proj.bias)
        h = torch.tanh(self._init_proj(fused))  # [B, H]

        # 3. Token inicial = BOS
        if target_proof is not None:
            # Teacher forcing
            T_dec = target_proof.size(1)
        else:
            T_dec = max_len or self.max_proof_len

        logits_list = []
        cur_token = torch.full((B, 1), bos_id, dtype=torch.long, device=device)

        for t in range(T_dec):
            if target_proof is not None:
                tok = target_proof[:, t]  # [B]
            else:
                tok = cur_token.squeeze(1)

            emb = self.dec_embedding(tok)  # [B, D]
            ctx = self._attention(h, enc_out, enc_mask)  # [B, 2H]
            gru_input = torch.cat((emb, ctx), dim=-1)    # [B, D+2H]
            h = self.dec_gru(gru_input, h)               # [B, H]

            # Logits
            logit_input = torch.cat((h, ctx, emb), dim=-1)  # [B, H+2H+D]
            logit = self.out_proj(logit_input)              # [B, V]
            logits_list.append(logit)

            if target_proof is None:
                cur_token = logit.argmax(dim=-1, keepdim=True)

        logits = torch.stack(logits_list, dim=1)  # [B, T_dec, V]

        return {
            "logits": logits,
            "enc_out": enc_out,
            "dec_state": h,
            "fused": fused,
        }


# ============================================================================
# Verificador (CNN-BiGRU + classificador sigmoid)
# ============================================================================

class VerifierCNNBiGRU(nn.Module):
    """Verificador: CNN-BiGRU + classificador sigmoid (v ∈ [0,1]).

    Recebe premissas (concat de axiomas + conjectura) e um passo hipotético,
    e retorna uma probabilidade de o passo estar correto.
    """

    def __init__(
        self,
        vocab_size: int,
        embedding_dim: int = 64,
        cnn_filters: int = 64,
        gru_hidden: int = 64,
        n_heads: int = 4,
        dropout: float = 0.1,
        pad_idx: int = 1,
    ):
        super().__init__()
        self.pad_idx = pad_idx
        self.gru_hidden = gru_hidden

        self.encoder = CooperativeCNNBiGRU(
            vocab_size=vocab_size,
            embedding_dim=embedding_dim,
            cnn_filters=cnn_filters,
            gru_hidden=gru_hidden,
            n_heads=n_heads,
            dropout=dropout,
            pad_idx=pad_idx,
        )

        # Classificador sigmoid
        # fused = out_A || out_B = (2*gru_hidden) || (2*gru_hidden) = 4*gru_hidden
        # combined = feat_pre || feat_passo = 8*gru_hidden
        self.classifier = nn.Sequential(
            nn.Linear(gru_hidden * 4 * 2, gru_hidden),  # 8*gru_hidden -> gru_hidden
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(gru_hidden, 1),
        )

    def forward(
        self,
        premissas_a: torch.Tensor,
        premissas_b: torch.Tensor,
        passo_a: torch.Tensor,
        passo_b: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            premissas_a, premissas_b: [B, T] (contexto da prova até o momento)
            passo_a, passo_b: [B, T_p] (passo gerado a ser validado)

        Returns:
            v: [B, 1] probabilidade de o passo estar correto
        """
        # Encode premissas
        enc_pre = self.encoder(premissas_a, premissas_b)
        feat_pre = enc_pre["fused"]  # [B, 256]

        # Encode passo
        enc_passo = self.encoder(passo_a, passo_b)
        feat_passo = enc_passo["fused"]  # [B, 256]

        # Concatena e classifica
        combined = torch.cat((feat_pre, feat_passo), dim=-1)  # [B, 512]
        logit = self.classifier(combined)  # [B, 1]
        v = torch.sigmoid(logit)  # [B, 1]
        return v


__all__ = [
    "fuzzy_not", "fuzzy_and", "fuzzy_or", "fuzzy_imp",
    "AntiHallucinationLayer",
    "GeneratorCNNBiGRU",
    "VerifierCNNBiGRU",
]