import os
import typing
from pathlib import Path
from typing import Any, Optional, Text, Dict, List, Type

import numpy as np

from bpemb import BPEmb

import rasa.utils.train_utils as train_utils
from rasa.nlu.components import Component
from rasa.nlu.config import RasaNLUModelConfig
from rasa.nlu.tokenizers.tokenizer import Tokenizer
from rasa.nlu.training_data import Message, TrainingData
from rasa.nlu.featurizers.featurizer import DenseFeaturizer
from rasa.nlu.constants import (
    DENSE_FEATURE_NAMES,
    DENSE_FEATURIZABLE_ATTRIBUTES,
    TEXT,
    TOKENS_NAMES,
)

if typing.TYPE_CHECKING:
    from rasa.nlu.model import Metadata


class BytePairFeaturizer(DenseFeaturizer):
    """This component adds BPEmb features."""

    @classmethod
    def required_components(cls) -> List[Type[Component]]:
        return [Tokenizer]

    @classmethod
    def required_packages(cls) -> List[Text]:
        return ["bpemb"]

    defaults = {
        # specifies the language of the subword segmentation model
        "lang": None,
        # specifies the dimension of the subword embeddings
        "dim": None,
        # specifies the vocabulary size of the segmentation model
        "vs": None,
        # if set to True and the given vocabulary size can't be loaded for the given
        # model, the closest size is chosen
        "vs_fallback": True,
        # specifies the folder in which downloaded BPEmb files will be cached
        "cache_dir": str(Path.home() / Path(".cache/bpemb")),
        # specifies the path to a custom SentencePiece model file
        "model_file": None,
        # specifies the path to a custom embedding file
        "emb_file": None,
    }

    language_list = ["vi"]

    def __init__(self, component_config: Optional[Dict[Text, Any]] = None) -> None:
        super().__init__(component_config)

        model_file, emb_file = (
            self.component_config[k] for k in ["model_file", "emb_file"]
        )
        if model_file:
            if not os.path.exists(model_file):
                raise FileNotFoundError(
                    f"BytePair model {model_file} not found. Please check config."
                )
        if emb_file:
            if not os.path.exists(emb_file):
                raise FileNotFoundError(
                    f"BytePair embedding file {emb_file} not found. Please check config."
                )

        if not self.component_config["lang"]:
            raise ValueError(
                "You must specify the `lang` parameter for BytePairEmbedding in `config.yml`."
            )

        if not self.component_config["vs"]:
            raise ValueError(
                "You must specify the `vs` parameter for BytePairEmbedding in `config.yml`."
            )

        if not self.component_config["dim"]:
            raise ValueError(
                "You must specify the `dim` parameter for BytePairEmbedding in `config.yml`."
            )

        self.model = BPEmb(
            lang=self.component_config["lang"],
            dim=self.component_config["dim"],
            vs=self.component_config["vs"],
            vs_fallback=self.component_config["vs_fallback"],
            cache_dir=self.component_config["cache_dir"],
            model_file=self.component_config["model_file"],
            emb_file=self.component_config["emb_file"],
        )

    def train(
        self,
        training_data: TrainingData,
        config: Optional[RasaNLUModelConfig] = None,
        **kwargs: Any,
    ) -> None:
        for example in training_data.intent_examples:
            for attribute in DENSE_FEATURIZABLE_ATTRIBUTES:
                self.set_bpemb_features(example, attribute)

    def create_word_vector(self, document: Text) -> np.ndarray:
        encoded_ids = self.model.encode_ids(document)
        if encoded_ids:
            return self.model.vectors[encoded_ids[0]]

        return np.zeros((self.component_config["dim"],), dtype=np.float32)

    def set_bpemb_features(self, message: Message, attribute: Text = TEXT) -> None:
        tokens = message.get(TOKENS_NAMES[attribute])

        if not tokens:
            return None

        text_vector = self.create_word_vector(document=message.text)
        word_vectors = [
            self.create_word_vector(document=t.text)
            for t in train_utils.tokens_without_cls(message, attribute)
        ]
        X = np.array(word_vectors + [text_vector])

        features = self._combine_with_existing_dense_features(
            message, additional_features=X, feature_name=DENSE_FEATURE_NAMES[attribute]
        )
        message.set(DENSE_FEATURE_NAMES[attribute], features)

    def process(self, message: Message, **kwargs: Any) -> None:
        self.set_bpemb_features(message)

    def persist(self, file_name: Text, model_dir: Text) -> Optional[Dict[Text, Any]]:
        pass

    @classmethod
    def load(
        cls,
        meta: Dict[Text, Any],
        model_dir: Optional[Text] = None,
        model_metadata: Optional["Metadata"] = None,
        cached_component: Optional["Component"] = None,
        **kwargs: Any,
    ) -> "Component":
        if cached_component:
            return cached_component

        return cls(meta)
