import torch
from transformers import AutoTokenizer, AutoModel


def batch_sequences_by_length(sequences, max_tokens=8000):
    """Groups sequences into batches whose padded token budget (n_seqs x
    max_len_in_batch) stays under max_tokens. Sequences are sorted
    longest-first so the first entry in each batch sets the padding length."""
    lengths = torch.tensor([len(s) for s in sequences])
    sort_idx = torch.argsort(lengths, descending=True)

    sorted_seqs = [sequences[i] for i in sort_idx]
    sorted_lengths = lengths[sort_idx]

    batches, reorder_map = [], []
    current_batch, current_indices, batch_max_len = [], [], 0

    for seq, length, orig_idx in zip(sorted_seqs, sorted_lengths, sort_idx):
        l_val = length.item()
        new_max = max(batch_max_len, l_val)
        if (len(current_batch) + 1) * new_max <= max_tokens:
            current_batch.append(seq)
            current_indices.append(orig_idx.item())
            batch_max_len = new_max
        else:
            if current_batch:
                batches.append(current_batch)
                reorder_map.extend(current_indices)
            current_batch = [seq]
            current_indices = [orig_idx.item()]
            batch_max_len = l_val

    if current_batch:
        batches.append(current_batch)
        reorder_map.extend(current_indices)

    restore_indices = torch.argsort(torch.tensor(reorder_map))
    return batches, restore_indices


class ESM2Embedder:
    def __init__(self, model_name="facebook/esm2_t12_35M_UR50D", device="cpu", max_aa=8000):
        self.model_name = model_name
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name).to(device)
        self.model.eval()
        self.device = device
        self.max_aa = max_aa

    def _pool(self, last_hidden_state, attention_mask):
        # residue-only mean: drop CLS (position 0) and EOS (last non-pad
        # position of each sequence) from the mask before averaging.
        mask = attention_mask.clone()
        mask[:, 0] = 0
        eos_pos = attention_mask.sum(dim=1) - 1
        mask[torch.arange(mask.size(0)), eos_pos] = 0
        mask = mask.unsqueeze(-1).float()
        return (last_hidden_state * mask).sum(1) / mask.sum(1)

    def _embed_batch(self, batch):
        if len(batch) == 0:
            return []
        inputs = self.tokenizer(
            batch, return_tensors="pt", padding=True, truncation=False,
        ).to(self.device)
        with torch.no_grad():
            outputs = self.model(**inputs)
        emb = self._pool(outputs.last_hidden_state, inputs["attention_mask"]).float().cpu()
        return [emb]

    def get_embeddings(self, sequences, max_aa=None):
        limit = max_aa if max_aa is not None else self.max_aa
        if len(sequences) == 1:
            batches = [sequences]
            restore_idx = torch.arange(1)
        else:
            batches, restore_idx = batch_sequences_by_length(sequences, limit)

        all_results = []
        for batch in batches:
            all_results.extend(self._embed_batch(batch))

        combined = torch.cat(all_results, dim=0)
        return combined[restore_idx]