# -*- coding: utf-8 -*-
# Copyright (C) 2016-2023 PyThaiNLP Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


class Translate:
    """
    Machine Translation

    :param str src_lang: source language
    :param str target_lang: target language
    :param bool use_gpu: load model to gpu (Default is False)

    **Options for source & target language**
        * *th* - *en* - Thai to English
        * *en* - *th* - English to Thai
        * *th* - *zh* - Thai to Chinese
        * *zh* - *th* - Chinese to Thai
        * *th* - *fr* - Thai to French

    :Example:

    Translate text from Thai to English::

        from pythainlp.translate import Translate
        th2en = Translate('th', 'en')

        th2en.translate("ฉันรักแมว")
        # output: I love cat.
    """

    def __init__(
        self, src_lang: str, target_lang: str, engine: str="default", use_gpu: bool = False
    ) -> None:
        """
        :param str src_lang: source language
        :param str target_lang: target language
        :param str engine: Machine Translation engine
        :param bool use_gpu: load model to gpu (Default is False)

        **Options for engine*
            * *default* - The engine default by each a language.
            * *small100* - A multilingual machine translation model (covering 100 languages)

        **Options for source & target language**
            * *th* - *en* - Thai to English
            * *en* - *th* - English to Thai
            * *th* - *zh* - Thai to Chinese
            * *zh* - *th* - Chinese to Thai
            * *th* - *fr* - Thai to French
            * *th* - *xx* - Thai to xx (xx is language code). It uses small100 model.
            * *xx* - *th* - xx to Thai (xx is language code). It uses small100 model.

        :Example:

        Translate text from Thai to English::

            from pythainlp.translate import Translate
            th2en = Translate('th', 'en')

            th2en.translate("ฉันรักแมว")
            # output: I love cat.
        """
        self.model = None
        self.engine = engine
        self.src_lang = src_lang
        self.use_gpu = use_gpu
        self.target_lang = target_lang
        self.load_model()

    def load_model(self):
        src_lang = self.src_lang
        target_lang = self.target_lang
        use_gpu = self.use_gpu
        if self.engine == "small100":
            from .small100 import Small100Translator

            self.model = Small100Translator(use_gpu)
        elif src_lang == "th" and target_lang == "en":
            from pythainlp.translate.en_th import ThEnTranslator

            self.model = ThEnTranslator(use_gpu)
        elif src_lang == "en" and target_lang == "th":
            from pythainlp.translate.en_th import EnThTranslator

            self.model = EnThTranslator(use_gpu)
        elif src_lang == "th" and target_lang == "zh":
            from pythainlp.translate.zh_th import ThZhTranslator

            self.model = ThZhTranslator(use_gpu)
        elif src_lang == "zh" and target_lang == "th":
            from pythainlp.translate.zh_th import ZhThTranslator

            self.model = ZhThTranslator(use_gpu)
        elif src_lang == "th" and target_lang == "fr":
            from pythainlp.translate.th_fr import ThFrTranslator

            self.model = ThFrTranslator(use_gpu)
        else:
            raise ValueError("Not support language!")

    def translate(self, text) -> str:
        """
        Translate text

        :param str text: input text in source language
        :return: translated text in target language
        :rtype: str
        """
        if self.engine == "small100":
            return self.model.translate(text, tgt_lang=self.target_lang)
        return self.model.translate(text)
