from pathlib import Path
from typing import Any, Callable, Iterable, List, Optional, Union, cast

from spacy.language import Language
from spacy.tokens import Doc
from spacy.training import Example
from spacy.util import SimpleFrozenList
import srsly


try:
    from transformers import AutoTokenizer, PreTrainedTokenizerBase
    from wsd_torch_models.bem import BEM
except ImportError:
    pass

from pymusas.file_utils import ensure_path
from pymusas.rankers.lexicon_entry import LexiconEntryRanker
from pymusas.spacy_api.taggers.neural import NeuralTagger
from pymusas.spacy_api.taggers.rule_based import RuleBasedTagger
from pymusas.spacy_api.utils import remove_custom_token_extension
from pymusas.taggers.rules.rule import Rule
from pymusas.utils import neural_extra_installed


class HybridTagger(RuleBasedTagger, NeuralTagger):
    '''
    [spaCy pipeline component](https://spacy.io/usage/processing-pipelines)
    of the :class:`pymusas.taggers.hybrid.HybridTagger`.

    This is a hybrid tagger which uses both the
    :class:`pymusas.spacy_api.taggers.rule_based.RuleBasedTagger`
    and the :class:`pymusas.spacy_api.taggers.neural.NeuralTagger` taggers. This tagger
    inherits from the `RuleBasedTagger` and `NeuralTagger`.
    The difference between this and the
    `RuleBasedTagger` is that this tagger will use the `NeuralTagger` to
    tag tokens that the `RuleBasedTagger` cannot tag, these are the tokens that
    will be tagged with the `Z99` default tag using the `RuleBasedTagger`.

    The tagger when called, through :func:`__call__`, and given a sequence of
    tokens and their associated linguistic data (lemma, Part Of Speech (POS))
    will apply one or more :class:`pymusas.taggers.rules.rule.Rule`s
    to create a list of possible candidate tags for each token in the sequence.
    Each candidate, represented as a
    :class:`pymusas.rankers.ranking_meta_data.RankingMetaData` object, for each
    token is then Ranked using a
    :class:`pymusas.rankers.lexicon_entry.LexiconEntryRanker` ranker. The best
    candidate and it's associated tag(s) for each token are then returned along
    with a `List` of token indexes indicating if the token is part of a Multi
    Word Expression (MWE).

    If we cannot tag a token then the following process will happen:
    1. If the token's POS tag is in `default_punctuation_tags` then it will assign the
    tag `PUNCT`.
    2. If the token's POS tag is in `default_number_tags` then it will assign the tag
    `N1`.
    3. Use the `NeuralTagger` to tag the token. The tags generated by the `NeuralTagger`
    are determined by how you have initialised the `NeuralTagger`.

    # Assigned Attributes

    <table>
        <tr>
            <th> Location </th>
            <th> Type </th>
            <th> Value </th>
        </tr>
        <tr>
            <td> Token._.pymusas_tags </td>
            <td> `List[str]` </td>
            <td> Predicted tags, the first tag in the List of tags is the
            most likely tag.</td>
        </tr>
        <tr>
            <td> Token._.pymusas_mwe_indexes </td>
            <td> `List[Tuple[int, int]]` </td>
            <td> Each `Tuple` indicates the start and end token index of the
            associated Multi Word Expression (MWE). If the `List` contains
            more than one `Tuple` then the MWE is discontinuous. For single word
            expressions the `List` will only contain 1 `Tuple` which will be
            (token_start_index, token_start_index + 1).</td>
        </tr>
    </table>
    
    # Config and implementation

    The default config is defined by the pipeline component factory and describes
    how the component should be configured. You can override its settings via the `config`
    argument on [nlp.add_pipe](https://spacy.io/api/language#add_pipe) or in your
    [config.cfg for training](https://spacy.io/usage/training#config).

    | Setting                  | Description                  |
    |--------------------------|------------------------------|
    | pymusas_tags_token_attr  | See parameters section below |
    | pymusas_mwe_indexes_attr | See parameters section below |
    | pos_attribute            | See parameters section below |
    | lemma_attribute          | See parameters section below |
    | top_n                    | See parameters section below |
    | device                   | See parameters section below |
    | tokenizer_kwargs         | See parameters section below |

    # Parameters

    name : `str`, optional (default = `pymusas_hybrid_tagger`)
        The component name. Defaults to the same name as the class variable
        `COMPONENT_NAME`.
    pymusas_tags_token_attr : `str`, optional (default = `pymusas_tags`)
        The name of the attribute to assign the predicted tags too under
        the `Token._` class.
    pymusas_mwe_indexes_attr : `str`, optional (default = `pymusas_mwe_indexes`)
        The name of the attribute to assign the start and end token index of the
        associated MWE too under the `Token._` class.
    pos_attribute : `str`, optional (default = `pos_`)
        The name of the attribute that the Part Of Speech (POS) tag is assigned too
        within the `Token` class. The POS tag value that comes from this attribute
        has to be of type `str`. With the current default we take the POS tag
        from `Token.pos_`. The POS tag can be an empty string if you do not require
        POS information or if you do not have a POS tagger. **NOTE** that if you
        do not have a POS tagger the default value for `Token.pos_` is an empty
        string.
    lemma_attribute : `str`, optional (default = `lemma_`)
        The name of the attribute that the lemma is assigned too within the `Token`
        class. The lemma value that comes from this attribute has to be of
        type `str`. With the current default we take the lemma from `Token.lemma_`.
        The lemma can be an empty string if you do not require
        lemma information or if you do not have a lemmatiser. **NOTE** that if you
        do not have a lemmatiser the default value for `Token.lemma_` is an empty
        string.
    top_n : `int`, optional (default = `5`)
        The number of tags the NeuralTagger will predict.
        If -1 all tags will be predicted.
        If 0 or less than 0 will raise a ValueError.
    device : `str`, optional (default = `'cpu'`)
        The device to load the NeuralTagger model, `wsd_model`, on. e.g.
        `'cpu'`, it has to be a string that can be passed to
        [`torch.device`](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device).
    tokenizer_kwargs : `dict[str, Any] | None` (default = `None`)
        Keyword arguments to pass to the NeuralTagger's sub-word tokenizer's
        `transformers.AutoTokenizer.from_pretrained` method.
        These keyword arguments are only passed to the tokenizer on initialization.

    # Instance Attributes

    name : `str`
        The component name.
    pymusas_tags_token_attr : `str`, optional (default = `pymusas_tags`)
        The given `pymusas_tags_token_attr`
    pymusas_mwe_indexes_attr : `str`, optional (default = `pymusas_mwe_indexes`)
        The given `pymusas_mwe_indexes_attr`
    rules : `List[pymusas.taggers.rules.rule.Rule]`, optional (default = `None`)
        For the RuleBasedTagger.
        The `rules` is set through the :func:`initialize` method. Before it is
        set by the :func:`initialize` method the value of this attribute is `None`.
    ranker : `pymusas.rankers.lexicon_entry.LexiconEntryRanker`, optional (default = `None`)
        For the RuleBasedTagger.
        The `ranker` is set through the :func:`initialize` method. Before it is
        set by the :func:`initialize` method the value of this attribute is `None`.
    default_punctuation_tags : `Set[str]`
        For the RuleBasedTagger.
        The `default_punctuation_tags` is set through the :func:`initialize` method.
    default_number_tags : `Set[str]`
        For the RuleBasedTagger.
        The `default_number_tags` is set through the :func:`initialize` method.
    pos_attribute : `str`, optional (default = `pos_`)
        For the RuleBasedTagger.
        The given `pos_attribute`
    lemma_attribute : `str`, optional (default = `lemma_`)
        For the RuleBasedTagger.
        The given `lemma_attribute`
    top_n : `int`, optional (default = `5`)
        For the NeuralTagger.
        The number of tags to predict. If -1 all tags will be predicted.
        If 0 or less than 0 will raise a ValueError.
    device : `torch.device`
        For the NeuralTagger.
        The device that the `wsd_model` will be loaded on. e.g. `torch.device`
    wsd_model : `wsd_torch_models.bem.BEM | None` (default = `None`)
        For the NeuralTagger.
        The neural Word Sense Disambiguation (WSD) model. This is `None` until
        the component is initialized or has been loaded from disk or bytes.
    tokenizer : `transformers.PreTrainedTokenizerBase | None` (default = `None`)
        For the NeuralTagger.
        The sub-word tokenizer that the `wsd_model` uses. This tokenizer
        further tokenizes the tokens from the spaCy tokenizer, hence it being a
        sub-word tokenizer. This is `None` until the component is initialized
        or has been loaded from disk or bytes.
    _tokenizer_kwargs : `dict[str, Any] | None` (default = `None`)
        For the NeuralTagger.
        The keyword arguments that have
        or will be passed to the tokenizer's `transformers.AutoTokenizer.from_pretrained`
        method. These keyword arguments are only passed to the tokenizer on
        initialization.

    # Class Attributes

    COMPONENT_NAME : `str`
        Name of component factory that this component is registered under. This
        is used as the first argument to
        [`Language.add_pipe`](https://spacy.io/api/language#add_pipe)
        if you want to add this component to your spaCy pipeline.

    # Raises
    
    `ValueError`
        If `top_n` is 0 or less than -1.

    # Examples

    ``` python
    >>> import spacy
    >>> from pymusas.taggers.rules.single_word import SingleWordRule
    >>> from pymusas.rankers.lexicon_entry import ContextualRuleBasedRanker
    >>> from pymusas.lexicon_collection import LexiconCollection
    >>> from pymusas.rankers.lexicon_entry import ContextualRuleBasedRanker
    >>> english_lexicon_url = 'https://raw.githubusercontent.com/UCREL/Multilingual-USAS/e5cef7be2aa6182e300152f4f55152310007f051/English/semantic_lexicon_en.tsv'
    >>> lexicon_lookup = LexiconCollection.from_tsv(english_lexicon_url, include_pos=True)
    >>> lemma_lexicon_lookup = LexiconCollection.from_tsv(english_lexicon_url, include_pos=False)
    >>> single_word_rule = SingleWordRule(lexicon_lookup, lemma_lexicon_lookup)
    >>> ranker = ContextualRuleBasedRanker(1, 0)
    >>> # Construction via spaCy pipeline
    >>> nlp = spacy.blank('en')
    >>> # Using default config
    >>> tagger = nlp.add_pipe('pymusas_hybrid_tagger')
    >>> tagger.initialize(rules=[single_word_rule],
    ...                   ranker=ranker,
    ...                   pretrained_model_name_or_path="ucrelnlp/PyMUSAS-Neural-English-Small-BEM")
    >>> tokens = nlp('The river full of creaturez')
    >>> all_tags = [token._.pymusas_tags for token in tokens]
    >>> all_indexes = [token._.pymusas_mwe_indexes for token in tokens]
    >>> assert all_tags == [['Z5'], ['W3/M4', 'N5+'], ['N5.1+'], ['Z5'], ['Z1', 'S2', 'S2.2', 'S3.2', 'S2.1']]
    >>> assert all_indexes == [[(0, 1)], [(1, 2)], [(2, 3)], [(3, 4)], [(4, 5)]]
    >>> # Custom config
    >>> custom_config = {'pymusas_tags_token_attr': 'semantic_tags',
    ...                  'pymusas_mwe_indexes_attr': 'mwe_indexes',
    ...                  'top_n': 2,
    ...                  'tokenizer_kwargs': {'add_prefix_space': True}}
    >>> nlp = spacy.blank('en')
    >>> tagger = nlp.add_pipe('pymusas_hybrid_tagger', config=custom_config)
    >>> tagger.initialize(rules=[single_word_rule],
    ...                   ranker=ranker,
    ...                   pretrained_model_name_or_path="ucrelnlp/PyMUSAS-Neural-English-Small-BEM")
    >>> tokens = nlp('The river full of creaturez')
    >>> all_tags = [token._.semantic_tags for token in tokens]
    >>> all_indexes = [token._.mwe_indexes for token in tokens]
    >>> assert all_tags == [['Z5'], ['W3/M4', 'N5+'], ['N5.1+'], ['Z5'], ['Z1', 'S2']]
    >>> assert all_indexes == [[(0, 1)], [(1, 2)], [(2, 3)], [(3, 4)], [(4, 5)]]

    ```
    '''
    COMPONENT_NAME = 'pymusas_hybrid_tagger'

    def __init__(self,
                 name: str = 'pymusas_hybrid_tagger',
                 pymusas_tags_token_attr: str = 'pymusas_tags',
                 pymusas_mwe_indexes_attr: str = 'pymusas_mwe_indexes',
                 pos_attribute: str = 'pos_',
                 lemma_attribute: str = 'lemma_',
                 top_n: int = 5,
                 device: str = 'cpu',
                 tokenizer_kwargs: dict[str, Any] | None = None,
                 ) -> None:
        RuleBasedTagger.__init__(self, name, pymusas_tags_token_attr, pymusas_mwe_indexes_attr, pos_attribute, lemma_attribute)
        # These custom token extension/attributes are also set by the NeuralTagger
        # if they are set more than once they generate a User Warning, thus
        # this is to mitigate this.
        remove_custom_token_extension(pymusas_tags_token_attr)
        remove_custom_token_extension(pymusas_mwe_indexes_attr)
        NeuralTagger.__init__(self, name, pymusas_tags_token_attr, pymusas_mwe_indexes_attr, top_n, device, tokenizer_kwargs)

    def _validate(self) -> None:
        '''
        Checks that: `self.rules`, `self.ranker`, `self.wsd_model`,
        and `self.tokenizer` are not `None`.

        In addition if the `self.wsd_model` is not loaded onto `self.device`,
        the model is loaded onto `self.device`.

        # Raises

        `ValueError`
            If `self.rules`, `self.ranker`, `self.wsd_model` or `self.tokenizer` are `None`
        '''
        error_msg = ('The `{}` attribute cannot be `None`, this '
                     'attribute can be set through the `initialize` method.')
        if self.rules is None:
            raise ValueError(error_msg.format('rules'))
        
        if self.ranker is None:
            raise ValueError(error_msg.format('ranker'))
        
        if self.wsd_model is None:
            raise ValueError(error_msg.format('wsd_model'))
        
        if self.tokenizer is None:
            raise ValueError(error_msg.format('tokenizer'))

        if self.wsd_model.base_model.device != self.device:
            self.wsd_model.to(self.device)

        self._validated = True
    
    def initialize(self,
                   get_examples: Optional[Callable[[], Iterable[Example]]] = None,
                   *,
                   nlp: Optional[Language] = None,
                   rules: Optional[List[Rule]] = None,
                   ranker: Optional[LexiconEntryRanker] = None,
                   default_punctuation_tags: Optional[List[str]] = None,
                   default_number_tags: Optional[List[str]] = None,
                   pretrained_model_name_or_path: Optional[str | Path] = None,
                   ) -> None:
        '''
        Initialize the tagger and load any of the resources given. The method is
        typically called by
        [`Language.initialize`](https://spacy.io/api/language#initialize)
        and lets you customize arguments it receives via the
        [`initialize.components`](https://spacy.io/api/data-formats#config-initialize)
        block in the config. The loading only happens during initialization,
        typically before training. At runtime, all data is load from disk.

        # Parameters

        rules : `List[pymusas.taggers.rules.rule.Rule]`
            A list of rules to apply to the sequence of tokens in the
            :func:`__call__`. The output from each rule is concatenated and given
            to the `ranker`.
        ranker : `pymusas.rankers.lexicon_entry.LexiconEntryRanker`
            A ranker to rank the output from all of the `rules`.
        default_punctuation_tags : `List[str]`, optional (default = `None`)
            The POS tags that represent punctuation. If `None` then we will use
            `['punc']`. The list will be converted into a `Set` before assigning
            to the `default_punctuation_tags` attribute.
        default_number_tags : `List[str]`, optional (default = `None`)
            The POS tags that represent numbers. If `None` then we will use
            `['num']`. The list will be converted into a `Set` before assigning
            to the `default_number_tags` attribute.
        pretrained_model_name_or_path : `str | Path`
            The string ID or path of the pretrained neural
            Word Sense Disambiguation (WSD) model to load.

            **NOTE:** currently we only support the
            [wsd_torch_models.bem.BEM model](https://github.com/UCREL/WSD-Torch-Models/blob/main/src/wsd_torch_models/bem.py#L29)
        
            * A string, the model id of a pretrained
            [wsd-torch-models](https://github.com/UCREL/WSD-Torch-Models/tree/main)
            that is hosted on the HuggingFace Hub.
            * A `Path` or `str` that is a directory that can be loaded
            through `from_pretrained` method from a
            [wsd-torch-models model](https://github.com/UCREL/WSD-Torch-Models/tree/main)

            **NOTE:** this model name or path has to also be able to load the tokenizer
            using the function `transformers.AutoTokenizer.from_pretrained(pretrained_model_name_or_path)`
        '''
        # Copied the code from the relevant methods in RuleBasedTagger and NeuralTagger
        # so that self._validate() is only called once and at the appropriate time.

        # Taken from RuleBasedTagger
        if rules is not None:
            self.rules = rules
        
        if ranker is not None:
            self.ranker = ranker
        
        if default_punctuation_tags is not None:
            self.default_punctuation_tags = set(default_punctuation_tags)
        
        if default_number_tags is not None:
            self.default_number_tags = set(default_number_tags)

        # Taken from NeuralTagger
        neural_extra_installed()
        if pretrained_model_name_or_path is not None:
            self.wsd_model = BEM.from_pretrained(pretrained_model_name_or_path)
            tokenizer_kwargs = {}
            if self._tokenizer_kwargs is not None:
                tokenizer_kwargs = self._tokenizer_kwargs
            tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path,  # type: ignore
                                                      **tokenizer_kwargs)
            assert isinstance(tokenizer, PreTrainedTokenizerBase)
            self.tokenizer = tokenizer
        
        self._validate()
    
    def __call__(self, doc: Doc) -> Doc:
        '''
        Applies the tagger to the spaCy document, modifies it in place, and
        returns it. This usually happens under the hood when the `nlp` object is
        called on a text and all pipeline components are applied to the `Doc` in
        order.
        
        # Parameters

        doc : `Doc`
            A [spaCy `Doc`](https://spacy.io/api/doc)

        # Returns

        `Doc`
        '''
        if not self._validated:
            self._validate()
        RuleBasedTagger.__call__(self, doc)

        self.tokenizer = cast(PreTrainedTokenizerBase, self.tokenizer)
        self.wsd_model = cast(BEM, self.wsd_model)
        
        # Try, catch error handling reference:
        # https://github.com/explosion/spaCy/blob/6af6c2e86cc7b08573b261563786bd1ab87d45e9/spacy/pipeline/lemmatizer.py#L131
        error_handler = self.get_error_handler()

        try:
            # We have made the assumption that all unknown tokens are single word
            # expressions, i.e. they only affect one token.
            unknown_token_indexes: list[int] = []
            token_texts: list[str] = []
            for token in doc:
                token_texts.append(token.text)
                rule_based_tags = getattr(token._, self.pymusas_tags_token_attr)
                if rule_based_tags == ["Z99"]:
                    tag_indexes = getattr(token._, self.pymusas_mwe_indexes_attr)
                    unknown_token_index = tag_indexes[0][0]
                    unknown_token_indexes.append(unknown_token_index)
            if unknown_token_indexes:
                predicted_tags_candidates = self.wsd_model.predict(token_texts,
                                                                   sub_word_tokenizer=self.tokenizer,
                                                                   top_n=self.top_n)
                for token_index in unknown_token_indexes:
                    token_text = token_texts[token_index]
                    neural_tags = predicted_tags_candidates[token_index]
                    if token_text.strip() == "":
                        neural_tags = ["Z9"]
                    token = doc[token_index]
                    setattr(token._, self.pymusas_tags_token_attr, neural_tags)
        except Exception as e:
            error_handler(self.name, self, [doc], e)
        
        return doc
    
    def to_bytes(self, *, exclude: Iterable[str] = SimpleFrozenList()) -> bytes:
        """
        **Not Implemented**

        Even though the HyBridTagger inherits from RuleBased tagger which
        has implemented this method, NeuralTagger has not therefore
        it is not implemented for the HybridTagger.
        """
        raise NotImplementedError("to_bytes is not implemented for HybridTagger")
    
    def from_bytes(self, bytes_data: bytes, *,
                   exclude: Iterable[str] = SimpleFrozenList()
                   ) -> "HybridTagger":
        """
        **Not Implemented**

        Even though the HyBridTagger inherits from RuleBased tagger which
        has implemented this method, NeuralTagger has not therefore
        it is not implemented for the HybridTagger.
        """
        raise NotImplementedError("from_bytes is not implemented for HybridTagger")

    def to_disk(self,
                path: Union[str, Path],
                *,
                exclude: Iterable[str] = SimpleFrozenList()
                ) -> None:
        '''
        Serialises the tagger to the given `path`.

        # Parameters

        path : `Union[str, Path]`
            Path to a directory. Path may be either string or `Path`-like
            object. If the directory does not exist it attempts to create a
            directory at the given `path`.
        
        exclude : `Iterable[str]`, optional (default = `SimpleFrozenList()`)
            This currently does not do anything, please ignore it.

        # Returns

        `None`

        # Examples

        ```python
        >>> from tempfile import TemporaryDirectory
        >>> from pymusas.spacy_api.taggers.hybrid import HybridTagger
        >>> from pymusas.taggers.rules.single_word import SingleWordRule
        >>> from pymusas.rankers.lexicon_entry import ContextualRuleBasedRanker
        >>> from pymusas.lexicon_collection import LexiconCollection
        >>> from pymusas.rankers.lexicon_entry import ContextualRuleBasedRanker
        >>> english_lexicon_url = 'https://raw.githubusercontent.com/UCREL/Multilingual-USAS/e5cef7be2aa6182e300152f4f55152310007f051/English/semantic_lexicon_en.tsv'
        >>> lexicon_lookup = LexiconCollection.from_tsv(english_lexicon_url, include_pos=True)
        >>> lemma_lexicon_lookup = LexiconCollection.from_tsv(english_lexicon_url, include_pos=False)
        >>> single_word_rule = SingleWordRule(lexicon_lookup, lemma_lexicon_lookup)
        >>> ranker = ContextualRuleBasedRanker(1, 0)
        >>> tagger = HybridTagger()
        >>> tagger.initialize(rules=[single_word_rule],
        ...                   ranker=ranker,
        ...                   pretrained_model_name_or_path="ucrelnlp/PyMUSAS-Neural-English-Small-BEM")
        >>> with TemporaryDirectory() as temp_dir:
        ...     _ = tagger.to_disk(temp_dir)
        ...

        ```
        '''
        if not self._validated:
            self._validate()
        RuleBasedTagger.to_disk(self, path, exclude=exclude)
        NeuralTagger.to_disk(self, path, exclude=exclude)

    def from_disk(self,
                  path: Union[str, Path],
                  *,
                  exclude: Iterable[str] = SimpleFrozenList()
                  ) -> "HybridTagger":
        '''
        Loads the tagger from the given `path` in place and returns it.

        # Parameters

        path : `Union[str, Path]`
            Path to an existing directory. Path may be either string or
            `Path`-like object.
        
        exclude : `Iterable[str]`, optional (default = `SimpleFrozenList()`)
            This currently does not do anything, please ignore it.

        # Returns

        :class:`HybridTagger`

        # Examples

        ```python
        >>> from tempfile import TemporaryDirectory
        >>> from pymusas.spacy_api.taggers.hybrid import HybridTagger
        >>> from pymusas.taggers.rules.single_word import SingleWordRule
        >>> from pymusas.rankers.lexicon_entry import ContextualRuleBasedRanker
        >>> from pymusas.lexicon_collection import LexiconCollection
        >>> from pymusas.rankers.lexicon_entry import ContextualRuleBasedRanker
        >>> english_lexicon_url = 'https://raw.githubusercontent.com/UCREL/Multilingual-USAS/e5cef7be2aa6182e300152f4f55152310007f051/English/semantic_lexicon_en.tsv'
        >>> lexicon_lookup = LexiconCollection.from_tsv(english_lexicon_url, include_pos=True)
        >>> lemma_lexicon_lookup = LexiconCollection.from_tsv(english_lexicon_url, include_pos=False)
        >>> single_word_rule = SingleWordRule(lexicon_lookup, lemma_lexicon_lookup)
        >>> ranker = ContextualRuleBasedRanker(1, 0)
        >>> tagger = HybridTagger()
        >>> tagger.initialize(rules=[single_word_rule],
        ...                   ranker=ranker,
        ...                   pretrained_model_name_or_path="ucrelnlp/PyMUSAS-Neural-English-Small-BEM")
        >>> tagger_2 = HybridTagger()
        >>> assert tagger_2.wsd_model is None
        >>> assert tagger_2.ranker is None
        >>> assert tagger_2.rules is None
        >>> with TemporaryDirectory() as temp_dir:
        ...     _ = tagger.to_disk(temp_dir)
        ...     _ = tagger_2.from_disk(temp_dir)
        ...
        >>> assert tagger_2.wsd_model.base_model_name == tagger.wsd_model.base_model_name
        >>> assert tagger_2.ranker == ranker
        >>> assert tagger_2.rules == [single_word_rule]

        ```

        '''
        # Copied the code from the relevant methods in RuleBasedTagger and NeuralTagger
        # so that self._validate() is only called once and at the appropriate time.

        neural_extra_installed()
        # Taken from RuleBasedTagger
        component_folder = ensure_path(path)
        
        ranker_path = Path(component_folder, 'ranker.bin')
        serialised_ranker = srsly.read_msgpack(ranker_path)
        ranker = LexiconEntryRanker.serialise_object_from_bytes(serialised_ranker)
        self.ranker = cast(LexiconEntryRanker, ranker)

        rules_path = Path(component_folder, 'rules.bin')
        serialised_rules = srsly.read_msgpack(rules_path)
        rules = Rule.serialise_object_list_from_bytes(serialised_rules)
        self.rules = cast(List[Rule], rules)

        default_punctuation_tags_path = Path(component_folder,
                                             'default_punctuation_tags.bin')
        default_punctuation_tags = srsly.read_msgpack(default_punctuation_tags_path)
        self.default_punctuation_tags = set(default_punctuation_tags)
        
        default_number_tags_path = Path(component_folder, 'default_number_tags.bin')
        default_number_tags = srsly.read_msgpack(default_number_tags_path)
        self.default_number_tags = set(default_number_tags)

        # Taken from NeuralTagger
        model_path = component_folder / "model"
        self.wsd_model = BEM.from_pretrained(model_path)

        tokenizer_path = component_folder / "tokenizer"
        self.tokenizer = cast(PreTrainedTokenizerBase,
                              AutoTokenizer.from_pretrained(tokenizer_path))  # type: ignore[no-untyped-call]

        self._validate()
        return self
        
    @property
    def pymusas_tags_token_attr(self) -> str:
        return self._pymusas_tags_token_attr

    @property
    def pymusas_mwe_indexes_attr(self) -> str:
        return self._pymusas_mwe_indexes_attr

    @property
    def pos_attribute(self) -> str:
        return self._pos_attribute

    @property
    def lemma_attribute(self) -> str:
        return self._lemma_attribute


@Language.factory(HybridTagger.COMPONENT_NAME,
                  assigns=['token._.pymusas_tags', 'token._.pymusas_mwe_indexes'],
                  default_config={'pymusas_tags_token_attr': 'pymusas_tags',
                                  'pymusas_mwe_indexes_attr': 'pymusas_mwe_indexes',
                                  'pos_attribute': 'pos_',
                                  'lemma_attribute': 'lemma_',
                                  'top_n': 5,
                                  'device': 'cpu',
                                  'tokenizer_kwargs': None})
def make_usas_hybrid_tagger(nlp: Language,
                            name: str,
                            pymusas_tags_token_attr: str,
                            pymusas_mwe_indexes_attr: str,
                            pos_attribute: str,
                            lemma_attribute: str,
                            top_n: int,
                            device: str,
                            tokenizer_kwargs: None | dict[str, Any]
                            ) -> HybridTagger:
    return HybridTagger(name,
                        pymusas_tags_token_attr,
                        pymusas_mwe_indexes_attr,
                        pos_attribute,
                        lemma_attribute,
                        top_n,
                        device,
                        tokenizer_kwargs)
