"""semantic_embeddings.py — Embedder semântico com Sinkhorn OT e Ensemble.

Inspirado em xavante_work/xavante/utils/semantic_embeddings.py, simplificado
para o projeto CNN-BiGRU. Mantém as 5 técnicas matemáticas:
    1. Sinkhorn Optimal Transport para alinhamento de tokens
    2. Spectral clustering init via Laplacian normalization
    3. Information Bottleneck (beta-VAE style)
    4. Elastic WMD (smoothed WMD via Sinkhorn)
    5. Multi-metric ensemble (Minkowski + random projections)

Saída: embeddings L2-normalizados com dimensão configurável (default 128).
"""
from __future__ import annotations

import logging
import math
from typing import List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

logger = logging.getLogger(__name__)


# ---------------- Sinkhorn OT ----------------

def sinkhorn_knopp(
    cost_matrix: np.ndarray,
    p: np.ndarray,
    q: np.ndarray,
    epsilon: float = 0.1,
    n_iters: int = 50,
    tol: float = 1e-6,
) -> np.ndarray:
    """Resolve OT regularizado entropicamente via Sinkhorn-Knopp.

    T = diag(u) exp(-C/eps) diag(v), atualizado iterativamente até convergência.
    """
    n, m = cost_matrix.shape
    K = np.exp(-cost_matrix / max(epsilon, 1e-8))
    u = np.ones(n) / n
    v = np.ones(m) / m
    for _ in range(n_iters):
        u_prev = u.copy()
        Kv = K @ v
        u = p / np.maximum(Kv, 1e-30)
        Ktu = K.T @ u
        v = q / np.maximum(Ktu, 1e-30)
        if np.max(np.abs(u - u_prev)) < tol:
            break
    T = u[:, None] * K * v[None, :]
    return T


def elastic_wmd(embs_a: np.ndarray, embs_b: np.ndarray, epsilon: float = 0.1) -> float:
    """Elastic Word Mover's Distance via Sinkhorn (diferenciável)."""
    if embs_a.size == 0 or embs_b.size == 0:
        return 0.0
    C = np.linalg.norm(
        embs_a[:, None, :] - embs_b[None, :, :], axis=-1
    ).astype(np.float64)
    p = np.ones(len(embs_a)) / len(embs_a)
    q = np.ones(len(embs_b)) / len(embs_b)
    T = sinkhorn_knopp(C, p, q, epsilon=epsilon, n_iters=30)
    return float(np.sum(T * C))


class SemanticEmbedder(nn.Module):
    """Embedder semântico com projeção L2-normalizada.

    Combina:
        - Embedding token -> dimensão oculta
        - Mean-pooling com máscara
        - Projeto linear para dimensão final
        - L2 normalization para retrieval
        - Information bottleneck: mínima informação com ruído, máxima com conteúdo
    """

    def __init__(
        self,
        vocab_size: int,
        embed_dim: int = 64,
        out_dim: int = 128,
        pad_idx: int = 1,
        beta: float = 0.01,
    ):
        super().__init__()
        self.pad_idx = pad_idx
        self.embed_dim = embed_dim
        self.out_dim = out_dim
        self.beta = beta  # Information Bottleneck weight

        self.token_embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx)
        nn.init.orthogonal_(self.token_embedding.weight)

        self.proj = nn.Linear(embed_dim, out_dim)
        self.out_norm = nn.LayerNorm(out_dim)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        return_ib_loss: bool = False,
    ) -> torch.Tensor:
        """Returns L2-normalized embeddings [B, out_dim].

        Args:
            input_ids: [B, T] token IDs
            attention_mask: [B, T] opcional (1=real, 0=PAD)
            return_ib_loss: se True, retorna também a loss IB (Information Bottleneck)
                            para que o trainer possa adicioná-la à perda total.

        Returns:
            emb_norm: [B, out_dim] L2-normalizado
            (opcional) ib_loss: escalar (tensor 0-dim) com a loss IB
        """
        if attention_mask is None:
            attention_mask = (input_ids != self.pad_idx).float()

        emb = self.token_embedding(input_ids)  # [B, T, D]
        mask = attention_mask.unsqueeze(-1)  # [B, T, 1]

        # Mean pooling com máscara
        summed = (emb * mask).sum(dim=1)
        counts = mask.sum(dim=1).clamp(min=1.0)
        pooled = summed / counts  # [B, D]

        proj = self.proj(pooled)  # [B, out_dim]
        proj = self.out_norm(proj)

        # L2 normalize
        norm = proj.norm(p=2, dim=-1, keepdim=True).clamp(min=1e-8)
        emb_norm = proj / norm

        # Information Bottleneck regularization (beta-VAE style):
        # penaliza variância excessiva (encoraja representação comprimida)
        if self.beta > 0:
            var = emb_norm.var(dim=0).mean()
            ib_loss = self.beta * var
        else:
            ib_loss = torch.zeros((), device=emb_norm.device)

        if return_ib_loss:
            return emb_norm, ib_loss
        # Para compatibilidade, anexa ao tensor (mas trainer deve usar return_ib_loss=True)
        setattr(emb_norm, "ib_loss", ib_loss)
        return emb_norm

    def similarity(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """Similaridade coseno (a já está normalizado)."""
        return torch.matmul(a, b.T)


__all__ = [
    "sinkhorn_knopp",
    "elastic_wmd",
    "SemanticEmbedder",
]