"""nlp.py — Natural Language Processing (NLP) para CNN-BiGRU.

Módulo de processamento de linguagem natural que fornece tarefas clássicas
de NLP em cima do backbone cooperativo CNN-BiGRU:

    - Sequence Classification (sentiment, topic, NLI)
    - Token Classification (NER, POS tagging)
    - Span Detection (QA extractive)
    - Sentence Pair Classification (paraphrase, entailment)
    - Embeddings para retrieval semântico

==============================================================================
ANÁLISE MATEMÁTICA E LÓGICA
==============================================================================

Dado o output do CooperativeCNNBiGRU (representações fused [B, 256] ou
sequências temporais [B, T, 128] por stream), cada tarefa NLP aplica uma
cabeça específica:

1. Sequence Classification:
        h_cls = SelfAttentionSummary(seq)  [B, D]
        logits = Linear(D, num_classes)    [B, C]
        loss = CrossEntropy(logits, y)

2. Token Classification:
        logits = Linear(D, num_labels)     [B, T, L]
        loss = CrossEntropy(logits, y)  (ignore padding)

3. Span Detection (QA):
        start_logits = Linear(D, 1)        [B, T]
        end_logits = Linear(D, 1)          [B, T]
        loss = CE(start_logits, start_pos) + CE(end_logits, end_pos)

4. Sentence Pair Classification:
        Usa o dual-stream A/B do CooperativeCNNBiGRU diretamente
        (já que o modelo é naturalmente dual-stream).
        logits = Linear(D_fused, num_classes)

5. Embeddings para retrieval:
        h = mean_pool(seq)                 [B, D]
        h = normalize(h, p=2)              [B, D]
        similarity = cosine(h_query, h_doc)
        loss = InfoNCE (contrastive)

==============================================================================
INTEGRAÇÃO
==============================================================================

- CooperativeCNNBiGRU (cnn_bigru.models.cooperative_bigru): backbone
- SelfAttentionSummary: já incluído no cooperative_bigru
- BBPETokenizer: tokenização

Autor: CNN-BiGRU Project
"""
from __future__ import annotations

import logging
import math
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple, Union

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

logger = logging.getLogger(__name__)


# ============================================================================
# Configuração
# ============================================================================

@dataclass
class NLPConfig:
    """Configuração do módulo NLP."""
    # Dimensões (precisam bater com o backbone)
    embed_dim: int = 64           # embedding de tokens
    cnn_filters: int = 64
    gru_hidden: int = 64          # 64 por direção -> 128 por stream
    feat_per_stream: int = 128    # = 2 * gru_hidden
    feat_fused: int = 256         # = 2 * feat_per_stream
    n_heads: int = 4
    dropout: float = 0.1
    pad_idx: int = 1
    # Tarefas
    num_classes_seq: int = 3      # sequence classification
    num_labels_tok: int = 7       # token classification (NER-style)
    # Pooling
    pooling: str = "self_attn"    # "self_attn" | "mean" | "max" | "cls"
    # Loss
    label_smoothing: float = 0.0
    # Device
    device: str = "cpu"


# ============================================================================
# Sequence Classification Head
# ============================================================================

class SequenceClassificationHead(nn.Module):
    """Cabeça para classificação de sequência inteira.

    Aplica pooling (self-attention, mean, max ou cls) sobre a sequência
    temporal e projeta para num_classes logits.

    Args:
        embed_dim: dimensão de entrada (e.g., 256 para fused, 128 por stream)
        num_classes: número de classes de saída
        dropout: prob de dropout
        pooling: estratégia de pooling
        n_heads: cabeças para self-attention pooling
    """

    def __init__(
        self,
        embed_dim: int,
        num_classes: int,
        dropout: float = 0.1,
        pooling: str = "self_attn",
        n_heads: int = 4,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.pooling = pooling
        self.num_classes = num_classes

        # Pooling
        if pooling == "self_attn":
            # Self-attention pooling simplificado (sem CLS token)
            self.attn_query = nn.Linear(embed_dim, 1)
        elif pooling == "cls":
            # CLS token learnable
            self.cls = nn.Parameter(torch.zeros(1, 1, embed_dim))
            nn.init.normal_(self.cls, std=0.02)
        # mean / max: não precisa de params

        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(embed_dim, num_classes)
        nn.init.xavier_uniform_(self.classifier.weight)
        nn.init.zeros_(self.classifier.bias)

    def forward(
        self,
        seq: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            seq: [B, T, D] sequência de hidden states
            mask: [B, T] máscara de padding (1=real, 0=pad)

        Returns:
            logits: [B, C]
            pooled: [B, D] representação pooled
        """
        if seq.dim() == 2:
            # [B, D] — já é pooled (e.g., fused)
            pooled = seq
        else:
            B, T, D = seq.shape

            if self.pooling == "self_attn":
                # Attention pooling: softmax over T
                scores = self.attn_query(seq).squeeze(-1)  # [B, T]
                if mask is not None:
                    scores = scores.masked_fill(mask == 0, -1e9)
                # Lidar com linhas totalmente mascaradas
                all_neg = (scores <= -1e9).all(dim=-1, keepdim=True)
                scores = torch.where(all_neg, torch.zeros_like(scores), scores)
                weights = F.softmax(scores, dim=-1).unsqueeze(-1)  # [B, T, 1]
                pooled = (seq * weights).sum(dim=1)  # [B, D]

            elif self.pooling == "mean":
                if mask is not None:
                    m = mask.unsqueeze(-1).float()
                    pooled = (seq * m).sum(dim=1) / m.sum(dim=1).clamp(min=1.0)
                else:
                    pooled = seq.mean(dim=1)

            elif self.pooling == "max":
                if mask is not None:
                    seq_masked = seq.masked_fill(mask.unsqueeze(-1) == 0, -1e9)
                    pooled, _ = seq_masked.max(dim=1)
                    # Lidar com sequências totalmente mascaradas (retorna 0)
                    all_neg = (pooled <= -1e9).all(dim=-1, keepdim=True)
                    pooled = torch.where(all_neg, torch.zeros_like(pooled), pooled)
                else:
                    pooled, _ = seq.max(dim=1)

            elif self.pooling == "cls":
                cls = self.cls.expand(B, -1, -1)  # [B, 1, D]
                seq_ext = torch.cat([cls, seq], dim=1)  # [B, T+1, D]
                pooled = seq_ext[:, 0, :]  # pega CLS

            else:
                raise ValueError(f"pooling desconhecido: {self.pooling}")

        pooled = self.dropout(pooled)
        logits = self.classifier(pooled)
        return logits, pooled


# ============================================================================
# Token Classification Head
# ============================================================================

class TokenClassificationHead(nn.Module):
    """Cabeça para classificação de tokens (NER, POS tagging).

    Aplica Linear(D, num_labels) sobre cada posição da sequência.

    Args:
        embed_dim: dimensão de entrada
        num_labels: número de rótulos (e.g., 7 para BIO tagging)
        dropout: prob de dropout
    """

    def __init__(self, embed_dim: int, num_labels: int, dropout: float = 0.1):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_labels = num_labels
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(embed_dim, num_labels)
        nn.init.xavier_uniform_(self.classifier.weight)
        nn.init.zeros_(self.classifier.bias)

    def forward(self, seq: torch.Tensor) -> torch.Tensor:
        """
        Args:
            seq: [B, T, D] sequência de hidden states
        Returns:
            logits: [B, T, L]
        """
        return self.classifier(self.dropout(seq))


# ============================================================================
# Span Detection Head (QA Extractive)
# ============================================================================

class SpanDetectionHead(nn.Module):
    """Cabeça para detecção de span (QA extractive).

    Prediz posições de início e fim do span de resposta.

    Args:
        embed_dim: dimensão de entrada
        dropout: prob de dropout
    """

    def __init__(self, embed_dim: int, dropout: float = 0.1):
        super().__init__()
        self.embed_dim = embed_dim
        self.dropout = nn.Dropout(dropout)
        self.start_classifier = nn.Linear(embed_dim, 1)
        self.end_classifier = nn.Linear(embed_dim, 1)
        for layer in (self.start_classifier, self.end_classifier):
            nn.init.xavier_uniform_(layer.weight)
            nn.init.zeros_(layer.bias)

    def forward(self, seq: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            seq: [B, T, D]
        Returns:
            start_logits: [B, T]
            end_logits: [B, T]
        """
        seq = self.dropout(seq)
        start_logits = self.start_classifier(seq).squeeze(-1)  # [B, T]
        end_logits = self.end_classifier(seq).squeeze(-1)
        return start_logits, end_logits


# ============================================================================
# Embedding Head (para retrieval semântico)
# ============================================================================

class EmbeddingHead(nn.Module):
    """Cabeça para gerar embeddings normalizados (retrieval semântico).

    Aplica mean-pool com máscara + LayerNorm + normalização L2.

    Args:
        embed_dim: dimensão de entrada e saída
        use_layer_norm: se True, aplica LayerNorm antes da normalização L2
    """

    def __init__(self, embed_dim: int, use_layer_norm: bool = True):
        super().__init__()
        self.embed_dim = embed_dim
        self.use_layer_norm = use_layer_norm
        if use_layer_norm:
            self.ln = nn.LayerNorm(embed_dim)

    def forward(
        self,
        seq: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
        normalize: bool = True,
    ) -> torch.Tensor:
        """
        Args:
            seq: [B, T, D]
            mask: [B, T] (1=real, 0=pad)
            normalize: se True, normaliza com L2
        Returns:
            embedding: [B, D]
        """
        if mask is not None:
            m = mask.unsqueeze(-1).float()
            pooled = (seq * m).sum(dim=1) / m.sum(dim=1).clamp(min=1.0)
        else:
            pooled = seq.mean(dim=1)

        if self.use_layer_norm:
            pooled = self.ln(pooled)

        if normalize:
            pooled = F.normalize(pooled, p=2, dim=-1)

        return pooled


# ============================================================================
# NLP Module (wrapper integrando todas as cabeças)
# ============================================================================

class NLPModule(nn.Module):
    """Módulo NLP integrando múltiplas tarefas no backbone CNN-BiGRU.

    Args:
        config: configuração NLP
        backbone: modelo CooperativeCNNBiGRU (opcional — se None, cria um novo)
    """

    def __init__(
        self,
        config: NLPConfig,
        backbone: Optional[nn.Module] = None,
    ):
        super().__init__()
        self.config = config

        # Backbone (cria se não fornecido)
        if backbone is None:
            from .cooperative_bigru import CooperativeCNNBiGRU
            backbone = CooperativeCNNBiGRU(
                vocab_size=32000,  # placeholder, sobrescrito via set_vocab_size
                embedding_dim=config.embed_dim,
                cnn_filters=config.cnn_filters,
                gru_hidden=config.gru_hidden,
                n_heads=config.n_heads,
                dropout=config.dropout,
                pad_idx=config.pad_idx,
            )
        self.backbone = backbone

        # Cabeças (uma por tarefa)
        self.seq_head = SequenceClassificationHead(
            embed_dim=config.feat_fused,
            num_classes=config.num_classes_seq,
            dropout=config.dropout,
            pooling=config.pooling,
            n_heads=config.n_heads,
        )
        self.tok_head = TokenClassificationHead(
            embed_dim=config.feat_per_stream,  # por stream (usa seq_A ou seq_B)
            num_labels=config.num_labels_tok,
            dropout=config.dropout,
        )
        self.span_head = SpanDetectionHead(
            embed_dim=config.feat_per_stream,
            dropout=config.dropout,
        )
        self.embed_head = EmbeddingHead(
            embed_dim=config.feat_per_stream,
            use_layer_norm=True,
        )

    # ----------------------------------------------------------------------
    # Tarefas
    # ----------------------------------------------------------------------

    def sequence_classification(
        self,
        input_ids_a: torch.Tensor,
        input_ids_b: torch.Tensor,
    ) -> Dict[str, torch.Tensor]:
        """Classificação de sequência (usando representação fused).

        Args:
            input_ids_a, input_ids_b: [B, T]

        Returns:
            dict com logits [B, C] e pooled [B, D]
        """
        out = self.backbone(input_ids_a, input_ids_b, return_sequences=False)
        fused = out["fused"]  # [B, 256]
        logits, pooled = self.seq_head(fused, mask=None)
        return {"logits": logits, "pooled": pooled, "fused": fused}

    def token_classification(
        self,
        input_ids_a: torch.Tensor,
        input_ids_b: torch.Tensor,
        stream: str = "A",
    ) -> torch.Tensor:
        """Classificação de tokens (NER/POS).

        Args:
            input_ids_a, input_ids_b: [B, T]
            stream: "A" ou "B" — qual stream usar para token classification

        Returns:
            logits: [B, T, L]
        """
        out = self.backbone(input_ids_a, input_ids_b, return_sequences=True)
        seq = out.get("seq_A") if stream == "A" else out.get("seq_B")
        if seq is None:
            # Fallback: usar fused (broadcast para T)
            fused = out["fused"]  # [B, 256]
            T = input_ids_a.size(1)
            seq = fused.unsqueeze(1).expand(-1, T, -1)  # [B, T, 256]
            # Projeta para feat_per_stream se necessário
            if fused.size(-1) != self.config.feat_per_stream:
                if not hasattr(self, "_proj_tok"):
                    self._proj_tok = nn.Linear(
                        fused.size(-1), self.config.feat_per_stream, bias=False
                    ).to(fused.device)
                seq = self._proj_tok(seq)
        return self.tok_head(seq)

    def span_detection(
        self,
        input_ids_a: torch.Tensor,
        input_ids_b: torch.Tensor,
        stream: str = "A",
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Detecção de span (QA extractive).

        Returns:
            (start_logits, end_logits) — ambos [B, T]
        """
        out = self.backbone(input_ids_a, input_ids_b, return_sequences=True)
        seq = out.get("seq_A") if stream == "A" else out.get("seq_B")
        if seq is None:
            fused = out["fused"]
            T = input_ids_a.size(1)
            seq = fused.unsqueeze(1).expand(-1, T, -1)
            if fused.size(-1) != self.config.feat_per_stream:
                if not hasattr(self, "_proj_span"):
                    self._proj_span = nn.Linear(
                        fused.size(-1), self.config.feat_per_stream, bias=False
                    ).to(fused.device)
                seq = self._proj_span(seq)
        return self.span_head(seq)

    def embed(
        self,
        input_ids_a: torch.Tensor,
        input_ids_b: torch.Tensor,
        stream: str = "A",
        normalize: bool = True,
    ) -> torch.Tensor:
        """Gera embedding normalizado para retrieval semântico.

        Returns:
            embedding: [B, D]
        """
        out = self.backbone(input_ids_a, input_ids_b, return_sequences=True)
        seq = out.get("seq_A") if stream == "A" else out.get("seq_B")
        mask = out.get("mask_A") if stream == "A" else out.get("mask_B")
        if seq is None:
            # Fallback: fused
            return self.embed_head(
                out["fused"].unsqueeze(1), mask=None, normalize=normalize
            )
        if seq.size(-1) != self.config.feat_per_stream:
            if not hasattr(self, "_proj_emb"):
                self._proj_emb = nn.Linear(
                    seq.size(-1), self.config.feat_per_stream, bias=False
                ).to(seq.device)
            seq = self._proj_emb(seq)
        return self.embed_head(seq, mask=mask, normalize=normalize)

    # ----------------------------------------------------------------------
    # Loss functions para cada tarefa
    # ----------------------------------------------------------------------

    def sequence_classification_loss(
        self,
        input_ids_a: torch.Tensor,
        input_ids_b: torch.Tensor,
        labels: torch.Tensor,
    ) -> torch.Tensor:
        """Computa CE loss para sequence classification."""
        out = self.sequence_classification(input_ids_a, input_ids_b)
        return F.cross_entropy(
            out["logits"], labels,
            label_smoothing=self.config.label_smoothing,
        )

    def token_classification_loss(
        self,
        input_ids_a: torch.Tensor,
        input_ids_b: torch.Tensor,
        labels: torch.Tensor,
        stream: str = "A",
    ) -> torch.Tensor:
        """Computa CE loss para token classification.

        Args:
            labels: [B, T] rótulos por token (use -100 ou pad_idx para ignorar)
        """
        logits = self.token_classification(input_ids_a, input_ids_b, stream=stream)
        # Ignora posições com label = -100 ou mask = 0
        mask = (input_ids_a != self.config.pad_idx).float()
        ignore_mask = (labels == -100) | (labels == self.config.pad_idx)
        labels_clamped = labels.clone()
        labels_clamped[ignore_mask] = -100
        return F.cross_entropy(
            logits.view(-1, self.config.num_labels_tok),
            labels_clamped.view(-1),
            ignore_index=-100,
        )

    def span_detection_loss(
        self,
        input_ids_a: torch.Tensor,
        input_ids_b: torch.Tensor,
        start_positions: torch.Tensor,
        end_positions: torch.Tensor,
        stream: str = "A",
    ) -> torch.Tensor:
        """Computa loss para span detection.

        Args:
            start_positions, end_positions: [B] índices do span (long)
        """
        start_logits, end_logits = self.span_detection(
            input_ids_a, input_ids_b, stream=stream
        )
        # Ignora posições fora do span (mask)
        loss_start = F.cross_entropy(start_logits, start_positions, ignore_index=-100)
        loss_end = F.cross_entropy(end_logits, end_positions, ignore_index=-100)
        return (loss_start + loss_end) / 2.0

    def contrastive_loss(
        self,
        query_ids_a: torch.Tensor,
        query_ids_b: torch.Tensor,
        pos_ids_a: torch.Tensor,
        pos_ids_b: torch.Tensor,
        neg_ids_a: torch.Tensor,
        neg_ids_b: torch.Tensor,
        temperature: float = 0.07,
    ) -> torch.Tensor:
        """InfoNCE loss para retrieval semântico.

        Args:
            query, pos, neg: tensores [B, T] para query, positivo e negativo
        """
        q = self.embed(query_ids_a, query_ids_b, normalize=True)
        p = self.embed(pos_ids_a, pos_ids_b, normalize=True)
        n = self.embed(neg_ids_a, neg_ids_b, normalize=True)

        # Similaridade coseno (já normalizado -> produto escalar)
        sim_pos = (q * p).sum(dim=-1, keepdim=True) / temperature  # [B, 1]
        sim_neg = (q * n).sum(dim=-1, keepdim=True) / temperature  # [B, 1]

        # InfoNCE: maximizar sim_pos - sim_neg
        logits = torch.cat([sim_pos, sim_neg], dim=-1)  # [B, 2]
        labels = torch.zeros(q.size(0), dtype=torch.long, device=q.device)  # 0 = pos
        return F.cross_entropy(logits, labels)


# ============================================================================
# Self-test
# ============================================================================

def _self_test():
    """Teste rápido do módulo NLP."""
    torch.manual_seed(42)
    from .cooperative_bigru import CooperativeCNNBiGRU
    config = NLPConfig(
        embed_dim=32,
        cnn_filters=32,
        gru_hidden=32,
        feat_per_stream=64,
        feat_fused=128,
        n_heads=4,
        dropout=0.1,
        pad_idx=1,
        num_classes_seq=3,
        num_labels_tok=5,
        pooling="self_attn",
    )
    backbone = CooperativeCNNBiGRU(
        vocab_size=100,
        embedding_dim=config.embed_dim,
        cnn_filters=config.cnn_filters,
        gru_hidden=config.gru_hidden,
        n_heads=config.n_heads,
        dropout=config.dropout,
        pad_idx=config.pad_idx,
    )
    nlp = NLPModule(config, backbone=backbone)

    B, T = 2, 8
    ids_a = torch.randint(2, 100, (B, T))
    ids_b = torch.randint(2, 100, (B, T))

    # Sequence classification
    out = nlp.sequence_classification(ids_a, ids_b)
    print(f"Seq cls: logits {out['logits'].shape}, pooled {out['pooled'].shape}")
    labels = torch.randint(0, 3, (B,))
    loss_seq = nlp.sequence_classification_loss(ids_a, ids_b, labels)
    print(f"  loss={loss_seq.item():.4f}")

    # Token classification
    tok_logits = nlp.token_classification(ids_a, ids_b)
    print(f"Tok cls: {tok_logits.shape}")
    tok_labels = torch.randint(0, 5, (B, T))
    loss_tok = nlp.token_classification_loss(ids_a, ids_b, tok_labels)
    print(f"  loss={loss_tok.item():.4f}")

    # Span detection
    start_logits, end_logits = nlp.span_detection(ids_a, ids_b)
    print(f"Span: start {start_logits.shape}, end {end_logits.shape}")
    sp = torch.randint(0, T, (B,))
    ep = torch.randint(0, T, (B,))
    loss_span = nlp.span_detection_loss(ids_a, ids_b, sp, ep)
    print(f"  loss={loss_span.item():.4f}")

    # Embedding
    emb = nlp.embed(ids_a, ids_b)
    print(f"Embed: {emb.shape}, norm: {emb.norm(dim=-1).mean().item():.4f}")

    print("NLP self-test OK")


if __name__ == "__main__":
    _self_test()


__all__ = [
    "NLPConfig",
    "SequenceClassificationHead",
    "TokenClassificationHead",
    "SpanDetectionHead",
    "EmbeddingHead",
    "NLPModule",
]