"""
Xavante - embedding_reconfig.py
Responsabilidade: Embedding reconfiguravel (Teorema 12.1/12.2).
Permite expansao dinamica do vocabulario e reconfiguracao de dim do modelo.
"""
from __future__ import annotations

import logging
import math
from typing import Optional

import torch
import torch.nn as nn

logger = logging.getLogger(__name__)


class ReconfigurableEmbedding(nn.Module):
    """
    Embedding que suporta:
    - Expansao de vocabulario sem reset
    - Reconfiguracao de dim (projecao)
    - Quantizacao para memoria
    """

    def __init__(self, vocab_size: int, d_model: int, padding_idx: Optional[int] = None):
        super().__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.padding_idx = padding_idx
        self.weight = nn.Parameter(torch.empty(vocab_size, d_model))
        nn.init.normal_(self.weight, mean=0.0, std=1.0 / math.sqrt(d_model) if False else 0.02)
        if padding_idx is not None:
            with torch.no_grad():
                self.weight[padding_idx].fill_(0)

    def forward(self, idx: torch.Tensor) -> torch.Tensor:
        return torch.nn.functional.embedding(idx, self.weight, padding_idx=self.padding_idx)

    def expand_vocab(self, new_size: int) -> None:
        """Expande o vocabulario preservando pesos antigos."""
        if new_size <= self.vocab_size:
            return
        old = self.weight.data
        new = torch.empty(new_size - self.vocab_size, self.d_model, device=old.device, dtype=old.dtype)
        nn.init.normal_(new, mean=0.0, std=0.02)
        self.weight = nn.Parameter(torch.cat([old, new], dim=0))
        self.vocab_size = new_size
        logger.info("Vocab expandido para %d", new_size)

    def reconfigure_dim(self, new_d: int) -> None:
        """Reconfigura a dimensao via projecao linear (random init para a parte nova)."""
        if new_d == self.d_model:
            return
        old = self.weight.data  # [V, d_model]
        if new_d > self.d_model:
            pad = torch.empty(old.shape[0], new_d - self.d_model, device=old.device, dtype=old.dtype)
            nn.init.normal_(pad, mean=0.0, std=0.02)
            new_w = torch.cat([old, pad], dim=1)
        else:
            new_w = old[:, :new_d]
        self.weight = nn.Parameter(new_w)
        self.d_model = new_d
        logger.info("Dim reconfigurada para %d", new_d)


__all__ = ["ReconfigurableEmbedding"]