from typing import Dict, List, Tuple, Union, Any

from .abstract import Engine, Aggregation
from ..utils import logger, load_singleton

from pathlib import Path

import pickle
from dataclasses import dataclass

class LSR(Engine):

    def __init__(self, name: str = 'LSR', config: Union[str, Path, Dict[str, Any]] = None, **kwargs):
        super().__init__(name, config, **kwargs)

        from bsparse.anserini import Anserini
        from bsparse.models import SpladeModel

        index = self.config['index_path']
        self.anserini = Anserini(index)

        @dataclass
        class Args:
            model_name: str = self.config['model_name']
            max_length: int = self.config['max_length']
            batch_size: int = self.config['batch_size']
            max_terms: int = self.config['max_terms']
            tokenizer_name: str = None


        config = Args()
        # we use SpladeV3 as the query encoder when MultiLSR is the document encoder
        self.model = SpladeModel(config)

        self.subset_mapper: Dict[str, str] = None
        if 'id_to_subset_mapping' in self.config:
            if self.config['id_to_subset_mapping'].endswith('.pkl'):
                self.subset_mapper = load_singleton(self.config['id_to_subset_mapping'])
            else:
                logger.warning(f"Unable to load subset mapping file {self.config['id_to_subset_mapping']}")



    def filter_subset(self, scores: Dict[str, float], only_subset: str = None):
        if only_subset is None or self.subset_mapper is None:
            return scores
        return {
            doc_id: score
            for doc_id, score in scores.items()
            if self.subset_mapper[doc_id] == only_subset
        }


    async def search_batch(self, queries: List[str], limit: Union[int, List[int]] = 20, subsets: List[str] = None, maxp: bool = True) -> List[Dict[str, float]]:
        if isinstance(limit, int):
            limit = [int(limit)]*len(queries)

        if subsets is None:
            subsets = [None]*len(queries)

        #anserini only takes 1 k, send max and remove any extra results.
        results = self.anserini.query_from_raw_text(queries, model=self.model, k=max(limit)*self.config.get('k_scale',20))
        return [
            dict(sorted(self.filter_subset(scores, subset).items(), key=lambda x: -x[1])[:l])
            for subset, l, scores in zip(subsets, limit, results)
        ]
