from abc import ABC, abstractmethod
from typing import List, Dict
from tqdm import tqdm


class LLMUncertaintyEstimator(ABC):
    @abstractmethod
    def compute_uncertainty(self, *args, **kwargs) -> Dict:
        pass

    @property
    @abstractmethod
    def get_required_fields(self) -> List[str]:
        # get the fields necessary for the compute uncertainty fn
        pass

    @abstractmethod
    def prepare_records(self, records) -> List:
        # select ms or bs or combine somehow (unlikely)
        pass

    def compute_for_records(self, records, use_tqdm=True, include_ids=True) -> List[Dict]:
        records = self.prepare_records(records)
        
        buffer = []
        for rec in tqdm(records) if use_tqdm else records:
            # only leave necessary entries
            rec_exclusive = {k: v for k, v in rec.items() if k in self.get_required_fields}
            uncerts = self.compute_uncertainty(**rec_exclusive)
            if include_ids:
                uncerts['dataset_idx'] = rec['dataset_idx']
            buffer.append(uncerts)
        
        return buffer
