import torch
import numpy as np
from transformers import AutoTokenizer, AutoModel
import faiss
from tqdm import tqdm

class DenseEncoder:
    def __init__(self, model_name="intfloat/simlm-base-msmarco-finetuned", device="cuda"):
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name).to(self.device)
        self.model.eval()
        
    def _mean_pool(self, token_embeddings, attention_mask):
        mask = attention_mask.unsqueeze(-1).float()
        return (token_embeddings * mask).sum(1) / mask.sum(1).clamp(min=1e-9)

    def encode(self, texts, batch_size=64, show_progress=True):
        all_embs = []
        iterator = range(0, len(texts), batch_size)
        if show_progress:
            iterator = tqdm(iterator, desc="Encoding")
            
        for i in iterator:
            batch = texts[i:i+batch_size]
            enc = self.tokenizer(batch, padding=True, truncation=True, max_length=128, return_tensors="pt").to(self.device)
            with torch.no_grad():
                out = self.model(**enc)
            emb = self._mean_pool(out.last_hidden_state, enc["attention_mask"])
            all_embs.append(emb.cpu().float().numpy())
            
        embs = np.vstack(all_embs).astype("float32")
        faiss.normalize_L2(embs)
        return embs