import pickle
from pathlib import Path
import joblib
import json
import numpy as np

import warnings

from typing import Union, List

from xmen.linkers import EntityLinker, logger
from xmen.kb import CompositeKnowledgebase

from scispacy.candidate_generation import (
    CandidateGenerator,
    load_approximate_nearest_neighbours_index,
    LinkerPaths,
    create_tfidf_ann_index,
)
from scispacy.linking import EntityLinker as ScispacyLinker
from scispacy.linking_utils import KnowledgeBase


class TFIDFNGramLinker(EntityLinker):
    """
    A class for linking named entities to concepts in a knowledge base using the TF-IDF and n-gram techniques.

    Args:
    - candidate_generator (CandidateGenerator, optional): The candidate generator object used to retrieve candidate concepts for entity linking. Defaults to None.
    - index_base_path (Union[Path, str], optional): The path to the folder containing the index for the knowledge base. Defaults to None.
    - kb (KnowledgeBase, optional): The knowledge base object used for entity linking. Defaults to None.
    - kb_name (str, optional): The name of the knowledge base. Defaults to "UMLS".
    - expand_abbreviations (bool, optional): Whether to expand abbreviations in entity text. Defaults to True.
    - k (int, optional): The number of candidate concepts to retrieve. Defaults to 30.
    - threshold (float, optional): The similarity threshold for a candidate concept to be considered a match. Defaults to 0.0.
    - no_definition_threshold (float, optional): The similarity threshold for a candidate concept without a definition to be considered a match. Defaults to 0.95.
    - filter_for_definitions (bool, optional): Whether to filter out candidate concepts without a definition. Defaults to False.
    - filter_types (bool, optional): Whether to filter out candidate concepts that are not of type "Entity". Defaults to False.
    """

    @staticmethod
    def write_index(index_base_path: Union[str, Path], jsonl_files: List[Union[str, Path]]):
        """
        Writes the knowledge base to disk in a pickle format and creates an index using the TF-IDF and n-gram techniques.
        - index_base_path: A string or Path object representing the path to the directory where the index will be stored.
        - jsonl_files: A list of string or Path objects representing the paths to the JSONL files containing the knowledge base.
        """
        kb = CompositeKnowledgebase(jsonl_files)

        index_base_path.mkdir(exist_ok=True, parents=True)

        pickle.dump(kb, open(Path(index_base_path) / "kb.pickle", "wb"))

        create_tfidf_ann_index(index_base_path, kb)

    @staticmethod
    def default_scispacy():
        """
        Returns a ScispacyLinker object with default parameters for the candidate generator.
        """
        return TFIDFNGramLinker.scispacy()

    @staticmethod
    def scispacy_no_thresholds():
        """
        Returns a TFIDFNGramLinker object with ScispacyLinker as candidate generator and the given parameters.
        """
        return TFIDFNGramLinker.scispacy(
            k=100,
            threshold=0.0,
            filter_for_definitions=False,
        )

    @staticmethod
    def load_candidate_generator(index_base_path: Union[str, Path]):
        """
        Loads the candidate generator object from the specified index base path.

        Args:
        - index_base_path (Union[str, Path]): A string or Path object representing the base path of the index.

        Returns:
        - CandidateGenerator: A CandidateGenerator object.
        """
        index_base_path = Path(index_base_path)
        lp = LinkerPaths(
            ann_index=index_base_path / "nmslib_index.bin",
            tfidf_vectorizer=index_base_path / "tfidf_vectorizer.joblib",
            tfidf_vectors=index_base_path / "tfidf_vectors_sparse.npz",
            concept_aliases_list=index_base_path / "concept_aliases.json",
        )
        ann_index = load_approximate_nearest_neighbours_index(lp)
        tfidf_vectorizer = joblib.load(lp.tfidf_vectorizer)
        ann_concept_aliases_list = json.load(open(lp.concept_aliases_list, encoding="utf-8"))
        kb = pickle.load(open(index_base_path / "kb.pickle", "rb"))

        cg = CandidateGenerator(ann_index, tfidf_vectorizer, ann_concept_aliases_list, kb)
        return cg

    @staticmethod
    def scispacy(
        expand_abbreviations: bool = None,
        k: int = None,
        threshold: float = None,
        no_definition_threshold: float = None,
        filter_for_definitions: bool = None,
    ):
        """
        Initializes a TFIDFNGramLinker object with a ScispacyLinker candidate generator.

        Args:
        - expand_abbreviations (bool): A boolean indicating whether to expand abbreviations or not. Defaults to None.
        - k (int): An integer indicating the number of candidates to generate for each mention. Defaults to None.
        - threshold (float): A float indicating the similarity threshold required to consider a candidate as a valid match. Defaults to None.
        - no_definition_threshold (float): A float indicating the similarity threshold required to consider a candidate without definition as a valid match. Defaults to None.
        - filter_for_definitions (bool): A boolean indicating whether to filter out candidates without definition or not. Defaults to None.

        Returns:
        - TFIDFNGramLinker: A TFIDFNGramLinker object.
        """
        logger.info("Initializing scispaCy")
        default_linker = ScispacyLinker()
        return TFIDFNGramLinker(
            default_linker.candidate_generator,
            kb=None,
            kb_name="UMLS",
            expand_abbreviations=expand_abbreviations if expand_abbreviations else default_linker.resolve_abbreviations,
            k=k if k else default_linker.k,
            threshold=threshold if threshold else default_linker.threshold,
            no_definition_threshold=no_definition_threshold
            if no_definition_threshold
            else default_linker.no_definition_threshold,
            filter_for_definitions=filter_for_definitions
            if filter_for_definitions
            else default_linker.filter_for_definitions,
            filter_types=False,
        )

    # Ignore caching when using dataset.map
    def __getstate__(self):
        return {}

    def __init__(
        self,
        candidate_generator: CandidateGenerator = None,
        index_base_path: Union[Path, str] = None,
        kb: KnowledgeBase = None,
        kb_name: str = "UMLS",
        expand_abbreviations: bool = True,
        k: int = 30,
        threshold: float = 0.0,
        no_definition_threshold: float = 0.95,
        filter_for_definitions: bool = False,
        filter_types: bool = False,
    ):
        if candidate_generator:
            self.candidate_generator = candidate_generator
        elif index_base_path:
            self.candidate_generator = TFIDFNGramLinker.load_candidate_generator(index_base_path)
        else:
            raise Exception("Either candidate_generator or index_base_path have to be provided")
        self.expand_abbreviations = expand_abbreviations
        self.k = k
        self.kb = kb if kb else self.candidate_generator.kb
        self.kb_name = kb_name
        self.threshold = threshold
        self.no_definition_threshold = no_definition_threshold
        self.filter_for_definitions = filter_for_definitions
        self.filter_types = filter_types

    def predict(self, passages: list, entities: list) -> dict:
        """
        Predicts the normalized form of the entities based on the provided passages and returns the updated entities object.

        Args:
        - passages (list): A list of strings representing the passages where the entities are found.
        - entities (list): A list of dictionaries representing the entities. Each dictionary must contain at least the "text" key with the entity text.

        Returns:
        - entities: A dictionary representing the updated entities object.
        """
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", np.VisibleDeprecationWarning)

            def get_m_string(e):
                if self.expand_abbreviations and e.get("long_form", None):
                    return e["long_form"]
                else:
                    return " ".join(e["text"])

            mention_strings = [get_m_string(e) for e in entities]
            batch_candidates = self.candidate_generator(mention_strings, self.k)

            for _, candidates, e in zip(mention_strings, batch_candidates, entities):
                predicted = []
                for cand in candidates:
                    score = max(cand.similarities)
                    if (
                        self.filter_for_definitions
                        and self.kb.cui_to_entity[cand.concept_id].definition is None
                        and score < self.no_definition_threshold
                    ):
                        continue
                    if score > self.threshold:
                        predicted.append((cand.concept_id, score))
                sorted_predicted = sorted(predicted, reverse=True, key=lambda x: x[1])
                e["normalized"] = [
                    {"db_id": cui, "score": score, "db_name": self.kb_name, "predicted_by": ["ngram"]}
                    for cui, score in sorted_predicted
                ]
        return entities
