# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# 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.

import pytest
from transformers import AutoTokenizer

from nemo.collections.nlp.data.spellchecking_asr_customization.bert_example import BertExampleBuilder
from nemo.collections.nlp.data.spellchecking_asr_customization.utils import (
    apply_replacements_to_text,
    substitute_replacements_in_text,
)


@pytest.mark.unit
def test_substitute_replacements_in_text():
    text = "we began the further diversification of our revenue base with the protterra supply agreement and the navastar joint development agreement"
    replacements = [(66, 75, 'pro-terra', 0.99986), (101, 109, 'navistar', 0.996)]
    gold_text = "we began the further diversification of our revenue base with the pro-terra supply agreement and the navistar joint development agreement"
    corrected_text = substitute_replacements_in_text(text, replacements, replace_hyphen_to_space=False)
    assert corrected_text == gold_text

    gold_text_no_hyphen = "we began the further diversification of our revenue base with the pro terra supply agreement and the navistar joint development agreement"
    corrected_text = substitute_replacements_in_text(text, replacements, replace_hyphen_to_space=True)
    assert corrected_text == gold_text_no_hyphen


@pytest.mark.unit
def test_apply_replacements_to_text():

    # min_prob = 0.5
    # dp_data = None,
    # min_dp_score_per_symbol: float = -99.9

    # test more than one fragment to replace, test multiple same replacements
    text = "we began the further diversification of our revenue base with the protterra supply agreement and the navastar joint development agreement"
    replacements = [
        (66, 75, 'proterra', 0.99986),
        (66, 75, 'proterra', 0.9956),
        (101, 109, 'navistar', 0.93),
        (101, 109, 'navistar', 0.91),
        (101, 109, 'navistar', 0.92),
    ]
    gold_text = "we began the further diversification of our revenue base with the proterra supply agreement and the navistar joint development agreement"
    corrected_text = apply_replacements_to_text(
        text, replacements, min_prob=0.5, replace_hyphen_to_space=False, dp_data=None
    )
    assert corrected_text == gold_text

    # test that min_prob works
    gold_text = "we began the further diversification of our revenue base with the proterra supply agreement and the navastar joint development agreement"
    corrected_text = apply_replacements_to_text(
        text, replacements, min_prob=0.95, replace_hyphen_to_space=False, dp_data=None
    )
    assert corrected_text == gold_text


@pytest.fixture()
def bert_example_builder():
    tokenizer = AutoTokenizer.from_pretrained("huawei-noah/TinyBERT_General_6L_768D")
    label_map = {"0": 0, "1": 1, "2": 2, "3": 3, "4": 4, "5": 5, "6": 6, "7": 7, "8": 8, "9": 9, "10": 10}
    semiotic_classes = {"PLAIN": 0, "CUSTOM": 1}
    max_seq_len = 256
    builder = BertExampleBuilder(label_map, semiotic_classes, tokenizer, max_seq_len)
    return builder


@pytest.mark.skip("Doesn't work download when testing on github, for unknown reason")
@pytest.mark.with_downloads
@pytest.mark.unit
def test_creation(bert_example_builder):
    assert bert_example_builder._tokenizer is not None


@pytest.mark.skip("Doesn't work download when testing on github, for unknown reason")
@pytest.mark.with_downloads
@pytest.mark.unit
def test_builder_get_spans(bert_example_builder):
    span_info_parts = ["CUSTOM 37 41", "CUSTOM 47 52", "CUSTOM 42 46", "CUSTOM 0 7"]
    gold_sorted_spans = [(1, 1, 8), (1, 38, 42), (1, 43, 47), (1, 48, 53)]
    spans = bert_example_builder._get_spans(span_info_parts)
    spans.sort()
    assert spans == gold_sorted_spans


@pytest.mark.skip("Doesn't work download when testing on github, for unknown reason")
@pytest.mark.with_downloads
@pytest.mark.unit
def test_builder_get_fragment_indices(bert_example_builder):
    hyp = "a b o u t _ o u r _ s h i p e r s _ b u t _ y o u _ k n o w"
    targets = [1]
    # a b o u t _ o u r _ s h i p e r s _ b u t _ y o u _ k n o w
    # 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0
    span_info_parts = ["CUSTOM 8 17"]
    gold_sorted_fragment_indices = [(7, 18, 1), (11, 18, 1)]
    fragment_indices = bert_example_builder._get_fragment_indices(hyp, targets, span_info_parts)
    fragment_indices.sort()
    assert fragment_indices == gold_sorted_fragment_indices

    # a b o u t _ o u r _ s h i p e r s _ b u t _ y o u _ k n o w
    # 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0
    span_info_parts = ["CUSTOM 10 16"]
    gold_sorted_fragment_indices = [(11, 18, 1)]
    fragment_indices = bert_example_builder._get_fragment_indices(hyp, targets, span_info_parts)
    fragment_indices.sort()
    assert fragment_indices == gold_sorted_fragment_indices


@pytest.mark.skip("Doesn't work download when testing on github, for unknown reason")
@pytest.mark.with_downloads
@pytest.mark.unit
def test_builder_get_input_features(bert_example_builder):
    hyp = "a s t r o n o m e r s _ d i d i e _ s o m o n _ a n d _ t r i s t i a n _ g l l o"
    ref = "d i d i e r _ s a u m o n;a s t r o n o m i e;t r i s t a n _ g u i l l o t;t r i s t e s s e;m o n a d e;c h r i s t i a n;a s t r o n o m e r;s o l o m o n;d i d i d i d i d i;m e r c y"
    targets = [1, 3]
    span_info_parts = ["CUSTOM 12 23", "CUSTOM 28 41"]

    gold_tags = [
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        0,
        0,
        0,
        0,
        0,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
    ]
    gold_input_ids = [
        101,
        1037,
        1055,
        1056,
        1054,
        1051,
        1050,
        1051,
        1049,
        1041,
        1054,
        1055,
        1035,
        1040,
        1045,
        1040,
        1045,
        1041,
        1035,
        1055,
        1051,
        1049,
        1051,
        1050,
        1035,
        1037,
        1050,
        1040,
        1035,
        1056,
        1054,
        1045,
        1055,
        1056,
        1045,
        1037,
        1050,
        1035,
        1043,
        1048,
        1048,
        1051,
        102,
        1040,
        1045,
        1040,
        1045,
        1041,
        1054,
        1035,
        1055,
        1037,
        1057,
        1049,
        1051,
        1050,
        102,
        1037,
        1055,
        1056,
        1054,
        1051,
        1050,
        1051,
        1049,
        1045,
        1041,
        102,
        1056,
        1054,
        1045,
        1055,
        1056,
        1037,
        1050,
        1035,
        1043,
        1057,
        1045,
        1048,
        1048,
        1051,
        1056,
        102,
        1056,
        1054,
        1045,
        1055,
        1056,
        1041,
        1055,
        1055,
        1041,
        102,
        1049,
        1051,
        1050,
        1037,
        1040,
        1041,
        102,
        1039,
        1044,
        1054,
        1045,
        1055,
        1056,
        1045,
        1037,
        1050,
        102,
        1037,
        1055,
        1056,
        1054,
        1051,
        1050,
        1051,
        1049,
        1041,
        1054,
        102,
        1055,
        1051,
        1048,
        1051,
        1049,
        1051,
        1050,
        102,
        1040,
        1045,
        1040,
        1045,
        1040,
        1045,
        1040,
        1045,
        1040,
        1045,
        102,
        1049,
        1041,
        1054,
        1039,
        1061,
        102,
    ]
    gold_input_mask = [
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
    ]
    gold_segment_ids = [
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        2,
        2,
        2,
        2,
        2,
        2,
        2,
        2,
        2,
        2,
        2,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
        3,
        4,
        4,
        4,
        4,
        4,
        4,
        4,
        4,
        4,
        4,
        5,
        5,
        5,
        5,
        5,
        5,
        5,
        6,
        6,
        6,
        6,
        6,
        6,
        6,
        6,
        6,
        6,
        7,
        7,
        7,
        7,
        7,
        7,
        7,
        7,
        7,
        7,
        7,
        8,
        8,
        8,
        8,
        8,
        8,
        8,
        8,
        9,
        9,
        9,
        9,
        9,
        9,
        9,
        9,
        9,
        9,
        9,
        10,
        10,
        10,
        10,
        10,
        10,
    ]
    gold_labels_mask = [
        0,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
    ]
    gold_input_ids_for_subwords = [
        101,
        26357,
        2106,
        2666,
        2061,
        8202,
        1998,
        13012,
        16643,
        2319,
        1043,
        7174,
        102,
        2106,
        3771,
        7842,
        2819,
        2239,
        102,
        28625,
        3630,
        9856,
        102,
        9822,
        26458,
        7174,
        2102,
        102,
        13012,
        13473,
        11393,
        102,
        13813,
        3207,
        102,
        3017,
        102,
        15211,
        102,
        9168,
        102,
        2106,
        28173,
        4305,
        4305,
        102,
        8673,
        102,
    ]
    gold_input_mask_for_subwords = [
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
    ]
    gold_segment_ids_for_subwords = [
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        1,
        1,
        1,
        1,
        1,
        1,
        2,
        2,
        2,
        2,
        3,
        3,
        3,
        3,
        3,
        4,
        4,
        4,
        4,
        5,
        5,
        5,
        6,
        6,
        7,
        7,
        8,
        8,
        9,
        9,
        9,
        9,
        9,
        10,
        10,
    ]
    gold_character_pos_to_subword_pos = [
        0,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        2,
        2,
        2,
        3,
        3,
        3,
        4,
        4,
        5,
        5,
        5,
        5,
        6,
        6,
        6,
        6,
        7,
        7,
        7,
        8,
        8,
        8,
        9,
        9,
        9,
        10,
        11,
        11,
        11,
        12,
        13,
        13,
        13,
        14,
        14,
        14,
        14,
        15,
        15,
        16,
        16,
        17,
        17,
        18,
        19,
        19,
        19,
        19,
        19,
        20,
        20,
        21,
        21,
        21,
        22,
        23,
        23,
        23,
        23,
        23,
        23,
        23,
        23,
        24,
        24,
        24,
        25,
        25,
        25,
        26,
        27,
        28,
        28,
        28,
        29,
        29,
        29,
        30,
        30,
        30,
        31,
        32,
        32,
        32,
        32,
        33,
        33,
        34,
        35,
        35,
        35,
        35,
        35,
        35,
        35,
        35,
        35,
        36,
        37,
        37,
        37,
        37,
        37,
        37,
        37,
        37,
        37,
        37,
        38,
        39,
        39,
        39,
        39,
        39,
        39,
        39,
        40,
        41,
        41,
        41,
        42,
        42,
        42,
        43,
        43,
        44,
        44,
        45,
        46,
        46,
        46,
        46,
        46,
        47,
    ]

    tags = [0 for _ in hyp.split()]
    for p, t in zip(span_info_parts, targets):
        c, start, end = p.split(" ")
        start = int(start)
        end = int(end)
        tags[start:end] = [t for i in range(end - start)]

    # get input features for characters
    (input_ids, input_mask, segment_ids, labels_mask, labels, _, _,) = bert_example_builder._get_input_features(
        hyp=hyp, ref=ref, tags=tags
    )

    # get input features for words
    hyp_with_words = hyp.replace(" ", "").replace("_", " ")
    ref_with_words = ref.replace(" ", "").replace("_", " ")
    (
        input_ids_for_subwords,
        input_mask_for_subwords,
        segment_ids_for_subwords,
        _,
        _,
        _,
        _,
    ) = bert_example_builder._get_input_features(hyp=hyp_with_words, ref=ref_with_words, tags=None)

    character_pos_to_subword_pos = bert_example_builder._map_characters_to_subwords(input_ids, input_ids_for_subwords)

    assert tags == gold_tags
    assert input_ids == gold_input_ids
    assert input_mask == gold_input_mask
    assert segment_ids == gold_segment_ids
    assert labels_mask == gold_labels_mask
    assert input_ids_for_subwords == gold_input_ids_for_subwords
    assert input_mask_for_subwords == gold_input_mask_for_subwords
    assert segment_ids_for_subwords == gold_segment_ids_for_subwords
    assert character_pos_to_subword_pos == gold_character_pos_to_subword_pos
