from parameterized import parameterized

from torchaudio._internal.module_utils import is_module_available
from torchaudio_unittest.common_utils import TorchaudioTestCase, skipIfNoModule

if is_module_available("unidecode") and is_module_available("inflect"):
    from pipeline_tacotron2.text.text_preprocessing import text_to_sequence
    from pipeline_tacotron2.text.numbers import (
        _remove_commas,
        _expand_pounds,
        _expand_dollars,
        _expand_decimal_point,
        _expand_ordinal,
        _expand_number,
    )


@skipIfNoModule("unidecode")
@skipIfNoModule("inflect")
class TestTextPreprocessor(TorchaudioTestCase):

    @parameterized.expand(
        [
            ["dr.  Strange?", [15, 26, 14, 31, 26, 29, 11, 30, 31, 29, 12, 25, 18, 16, 10]],
            ["ML, is        fun.", [24, 23, 6, 11, 20, 30, 11, 17, 32, 25, 7]],
            ["I love torchaudio!", [20, 11, 23, 26, 33, 16, 11, 31, 26, 29, 14, 19, 12, 32, 15, 20, 26, 2]],
            # 'one thousand dollars, twenty cents'
            ["$1,000.20", [26, 25, 16, 11, 31, 19, 26, 32, 30, 12, 25, 15, 11, 15, 26, 23, 23,
                           12, 29, 30, 6, 11, 31, 34, 16, 25, 31, 36, 11, 14, 16, 25, 31, 30]],
        ]
    )
    def test_text_to_sequence(self, sent, seq):

        assert (text_to_sequence(sent) == seq)

    @parameterized.expand(
        [
            ["He, she, and I have $1,000", "He, she, and I have $1000"],
        ]
    )
    def test_remove_commas(self, sent, truth):

        assert (_remove_commas(sent) == truth)

    @parameterized.expand(
        [
            ["He, she, and I have £1000", "He, she, and I have 1000 pounds"],
        ]
    )
    def test_expand_pounds(self, sent, truth):

        assert (_expand_pounds(sent) == truth)

    @parameterized.expand(
        [
            ["He, she, and I have $1000", "He, she, and I have 1000 dollars"],
            ["He, she, and I have $3000.01", "He, she, and I have 3000 dollars, 1 cent"],
            ["He has $500.20 and she has $1000.50.",
             "He has 500 dollars, 20 cents and she has 1000 dollars, 50 cents."],
        ]
    )
    def test_expand_dollars(self, sent, truth):

        assert (_expand_dollars(sent) == truth)

    @parameterized.expand(
        [
            ["1000.20", "1000 point 20"],
            ["1000.1", "1000 point 1"],
        ]
    )
    def test_expand_decimal_point(self, sent, truth):

        assert (_expand_decimal_point(sent) == truth)

    @parameterized.expand(
        [
            ["21st centry", "twenty-first centry"],
            ["20th centry", "twentieth centry"],
            ["2nd place.", "second place."],
        ]
    )
    def test_expand_ordinal(self, sent, truth):

        assert (_expand_ordinal(sent) == truth)
        _expand_ordinal,

    @parameterized.expand(
        [
            ["100020 dollars.", "one hundred thousand twenty dollars."],
            ["1234567890!", "one billion, two hundred thirty-four million, "
                            "five hundred sixty-seven thousand, eight hundred ninety!"],
        ]
    )
    def test_expand_number(self, sent, truth):

        assert (_expand_number(sent) == truth)
