from typing import Any, Dict, List, Union
from torch import nn
import torch
from transformers import AutoTokenizer, AutoModel, CLIPTokenizer, CLIPTextModel
import transformers

# Silence the HuggingFace "UNEXPECTED keys" load report and progress bars
transformers.logging.set_verbosity_error()

class TextEncoder(nn.Module):
    def __init__(self, config: Dict[str, Any]):
        super().__init__()
        encoder_type = config["model"]["text_encoder_type"]
        pretrained_path = config["model"]["text_encoder_pretrained"][encoder_type]
        
        self.encoder_type = encoder_type
        
        if encoder_type == "clip":
            self.tokenizer = CLIPTokenizer.from_pretrained(pretrained_path)
            self.model = CLIPTextModel.from_pretrained(pretrained_path)
        elif encoder_type in ["biobert", "sbert"]:
            self.tokenizer = AutoTokenizer.from_pretrained(pretrained_path)
            self.model = AutoModel.from_pretrained(pretrained_path)
        else:
            raise ValueError(f"Unknown text_encoder_type: {encoder_type}")
            
        for param in self.model.parameters():
            param.requires_grad = False
        self.model.eval()

        # In-memory embedding cache active ONLY during training mode
        self._cache: Dict[str, torch.Tensor] = {}

    def clear_cache(self):
        self._cache.clear()

    def _encode(self, text_list: List[str], device: torch.device) -> torch.Tensor:
        inputs = self.tokenizer(text_list, padding=True, truncation=True, return_tensors="pt").to(device)
        outputs = self.model(**inputs)
            
        if self.encoder_type == "clip":
            # CLIP returns a dedicated pooled_output for the whole sentence
            embeddings = outputs.pooler_output
            
        elif self.encoder_type == "biobert":
            # BERT models typically use the [CLS] token (index 0) for sentence-level tasks
            embeddings = outputs.last_hidden_state[:, 0, :]
            
        elif self.encoder_type == "sbert":
            # Sentence-BERT using Mean Pooling across all tokens
            attention_mask = inputs['attention_mask']
            token_embeddings = outputs.last_hidden_state
            
            # Expand attention mask to match embedding dimensions
            input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
            
            # Sum the embeddings and divide by the number of non-padded tokens
            sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
            sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
            embeddings = sum_embeddings / sum_mask
            
        return embeddings

    @torch.inference_mode()
    def forward(self, text: Union[str, List[str]]) -> torch.Tensor:
        device = next(self.model.parameters()).device
        is_single = isinstance(text, str)
        text_list = [text] if is_single else text

        uncached = list({t for t in text_list if t not in self._cache or self._cache[t].device != device})
        if uncached:
            new_embeddings = self._encode(uncached, device)
            for t, emb in zip(uncached, new_embeddings):
                self._cache[t] = emb.detach()

        embeddings = torch.stack([self._cache[t] for t in text_list], dim=0)

        return embeddings[0] if is_single else embeddings