"""
Vector embedding for RAG - uses model's own hidden states as embeddings (no external API)
Flaw fix: embedding drift -> L2 normalize, mean pool last hidden.
"""
import torch
import numpy as np
from typing import List

class AresEmbedder:
    def __init__(self, model, tokenizer, device="cpu"):
        self.model = model
        self.tokenizer = tokenizer
        self.device = device
        self.model.eval()

    @torch.no_grad()
    def embed(self, texts: List[str], batch_size=8, max_len=512) -> np.ndarray:
        embeddings = []
        for i in range(0, len(texts), batch_size):
            batch = texts[i:i+batch_size]
            ids_batch = [self.tokenizer.encode(t)[:max_len] for t in batch]
            max_l = max(len(x) for x in ids_batch)
            padded = [x + [self.tokenizer.pad_token_id]*(max_l-len(x)) for x in ids_batch]
            input_ids = torch.tensor(padded, dtype=torch.long, device=self.device)
            out = self.model(input_ids=input_ids)
            hidden = out["hidden_states"]  # [b,s,hidden]
            # Mean pooling ignoring pad
            mask = (input_ids != self.tokenizer.pad_token_id).float().unsqueeze(-1)  # [b,s,1]
            summed = (hidden * mask).sum(dim=1)
            counts = mask.sum(dim=1).clamp(min=1)
            mean = summed / counts
            # L2 normalize
            norm = torch.nn.functional.normalize(mean, p=2, dim=1)
            embeddings.append(norm.cpu().numpy())
        if embeddings:
            return np.concatenate(embeddings, axis=0)
        return np.zeros((0, self.model.config.hidden_size))

    @staticmethod
    def cosine_similarity(query_emb: np.ndarray, doc_embs: np.ndarray) -> np.ndarray:
        # query_emb: [hidden] or [1, hidden], doc_embs: [n, hidden]
        if query_emb.ndim == 1:
            query_emb = query_emb[None, :]
        # Assume normalized
        return np.dot(doc_embs, query_emb.T).squeeze()