import os
from typing import List, Dict

import torch
import faiss
import numpy as np
from tqdm import trange, tqdm

from beir.retrieval.search import BaseSearch
from beir.retrieval.search.dense.util import save_dict_to_tsv, load_tsv_to_dict
from beir.retrieval.search.dense.faiss_index import  FaissIndex

from accelerate import PartialState
from transformers import BertTokenizer, BertModel

class FlatIPFaissSearch(BaseSearch):
    def __init__(
        self, 
        model, 
        batch_size: int = 128,  
        output_dir: str = None,
        **kwargs
        ):
        self.model = model
        self.batch_size = batch_size
        self.score_functions = ['cos_sim','dot']
        self.mapping_tsv_keys = ["beir-docid", "faiss-docid"]
        self.dim_size = 768
        self.output_dir = output_dir

    def _create_mapping_ids(self, corpus_ids):
        self.mapping = {}
        self.rev_mapping = {}
        if not all(isinstance(doc_id, int) for doc_id in corpus_ids):
            for idx in range(len(corpus_ids)):
                self.mapping[corpus_ids[idx]] = idx
                self.rev_mapping[idx] = corpus_ids[idx]
    
    def save(self, output_dir: str):
        """
        save embedding and mapping to disk
        """
        # get current process id and concat to prefix
        distributed_state = PartialState()
        pid = distributed_state.process_index
        # Save Faiss Index to disk
        save_embed_path = os.path.join(output_dir, "{}_{}".format('embed', pid))
        print("Saving Embeddings to path: {}".format(save_embed_path))
        np.save(save_embed_path, self.corpus_embeddings)
        if pid == 0:
            print("Index size: {:.2f}MB".format(os.path.getsize(save_embed_path + ".npy")*0.000001*4))

        # Save Mapping to disk
        if pid == 0:
            save_dict_to_tsv(self.mapping, os.path.join(output_dir, "mapping_dic.tsv"), self.mapping_tsv_keys)
            print("Mapping saved to path: {}".format(os.path.join(output_dir, "mapping_dic.tsv")))

    def _embed(self, corpus: Dict[str, Dict[str, str]], score_function: str = None, **kwargs):
        corpus_ids = [key for key, value in corpus.items()]
        self._create_mapping_ids(corpus_ids)
        corpus = [corpus[cid] for cid in corpus_ids]

        print("Encoding Corpus in batches... Warning: This might take a while!")

        distributed_state = PartialState()
        device = distributed_state.device
        self.model.ctx_encoder.to(device)
        with distributed_state.split_between_processes(corpus) as sharded_corpus:
            shared_corpus = [sharded_corpus[idx:idx+self.batch_size] for idx in range(0, len(sharded_corpus), self.batch_size)]
            encoding_progress = tqdm(total=len(shared_corpus), desc="Encoding Passages: ", ncols=100, disable=not distributed_state.is_main_process)
            corpus_encoding = []
            for data in shared_corpus:
                sub_encoding = self.model.encode_corpus(corpus_batch=data, device=device)
                corpus_encoding.append(sub_encoding)
                encoding_progress.update(1)
            self.corpus_embeddings = np.concatenate(corpus_encoding, axis=0)


    def embed_and_save(self, corpus: Dict[str, Dict[str, str]], score_function: str = None, **kwargs):
        self._embed(corpus, score_function, **kwargs)
        self.save(self.output_dir)

    def load_and_index(self, embed_dir: str = "embed/", mapping_dict_dir:str = "mapping_dic.tsv"):
        # load and get the faiss ids from mapping(msmarco-docid to faiss-docid)
        sub_corpus_embeddings = [np.load(os.path.join(embed_dir, f)) for f in os.listdir(embed_dir) if f.endswith(".npy")] # it should be 4 np files which are generated by 4 shards
        self.corpus_embeddings = np.concatenate(sub_corpus_embeddings, axis=0)
        self.mapping = load_tsv_to_dict(mapping_dict_dir, header=True)
        self.faiss_ids = list(self.mapping.values())
        self.rev_mapping = {v: k for k, v in self.mapping.items()}

        # index the corpus
        base_index = faiss.IndexFlatIP(self.dim_size)
        self.faiss_index = FaissIndex.build(self.faiss_ids, self.corpus_embeddings, base_index)
    
    def search(self, 
               corpus: Dict[str, Dict[str, str]],
               queries: Dict[str, str], 
               top_k: int,
               score_function = str, **kwargs) -> Dict[str, Dict[str, float]]:
        self.results = {}
        
        assert score_function in self.score_functions
        normalize_embeddings = True if score_function == "cos_sim" else False

        query_ids = list(queries.keys())
        queries = [queries[qid] for qid in queries]
        print("Computing Query Embeddings. Normalize: {}...".format(normalize_embeddings))
        query_embeddings = self.model.encode_queries(
            queries, show_progress_bar=True, 
            batch_size=self.batch_size, 
            normalize_embeddings=normalize_embeddings)
        
        faiss_scores, faiss_doc_ids = self.faiss_index.search(query_embeddings, top_k, **kwargs)
        
        for idx in range(len(query_ids)):
            scores = [float(score) for score in faiss_scores[idx]]
            if len(self.rev_mapping) != 0:
                doc_ids = [self.rev_mapping[doc_id] for doc_id in faiss_doc_ids[idx]]
            else:
                doc_ids = [str(doc_id) for doc_id in faiss_doc_ids[idx]]
            self.results[query_ids[idx]] = dict(zip(doc_ids, scores))
        
        return self.results
    
    def get_index_name(self):
        return "flat_faiss_index"


class CustomBiEncoder :
    def __init__(self, model_save_dir=None, ) :
        self.query_encoder = BertModel.from_pretrained(os.path.join(model_save_dir, 'query_encoder'))
        self.ctx_encoder = BertModel.from_pretrained(os.path.join(model_save_dir, 'doc_encoder'))

        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

    def encode_queries(self, queries: List[str], batch_size: int, **kwargs) -> np.ndarray:
        CLS_POS=0
        query_embeddings = []
        self.query_encoder.cuda()
        with torch.no_grad():
            for start_idx in trange(0, len(queries), batch_size, desc = '>>> Encoding Queries : '):
                encoded = self.tokenizer(queries[start_idx:start_idx+batch_size], max_length=256,padding='max_length',return_tensors='pt',truncation=True)
                query_repr = self.query_encoder(
                        input_ids = encoded['input_ids'].cuda(),
                        attention_mask = encoded['attention_mask'].cuda(),
                        token_type_ids = encoded['token_type_ids'].cuda(),
                    ).last_hidden_state[:,CLS_POS,:].detach().cpu()
                query_embeddings += query_repr.detach().cpu()

        return torch.stack(query_embeddings)

    def encode_corpus(self, corpus_batch: List[Dict[str, str]], device, **kwargs) -> np.ndarray:
        CLS_POS=0
        with torch.no_grad():
            titles = [row['title'] for row in corpus_batch]
            texts = [row['text']  for row in corpus_batch]
            encoded = self.tokenizer(titles, texts, truncation='longest_first', padding=True, return_tensors='pt')
            corpus_repr = self.ctx_encoder(
                    input_ids = encoded['input_ids'].to(device),
                    attention_mask = encoded['attention_mask'].to(device),
                    token_type_ids = encoded['token_type_ids'].to(device)
                ).last_hidden_state[:,CLS_POS,:].detach().cpu()
        return corpus_repr