import argparse
import logging
import os
import unittest

from interactive_asr.utils import setup_asr, transcribe_file


class ASRTest(unittest.TestCase):
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)

    arguments_dict = {
        "path": "/scratch/jamarshon/downloads/model.pt",
        "input_file": "/scratch/jamarshon/audio/examples/interactive_asr/data/sample.wav",
        "data": "/scratch/jamarshon/downloads",
        "user_dir": "/scratch/jamarshon/fairseq-py/examples/speech_recognition",
        "no_progress_bar": False,
        "log_interval": 1000,
        "log_format": None,
        "tensorboard_logdir": "",
        "tbmf_wrapper": False,
        "seed": 1,
        "cpu": True,
        "fp16": False,
        "memory_efficient_fp16": False,
        "fp16_init_scale": 128,
        "fp16_scale_window": None,
        "fp16_scale_tolerance": 0.0,
        "min_loss_scale": 0.0001,
        "threshold_loss_scale": None,
        "criterion": "cross_entropy",
        "tokenizer": None,
        "bpe": None,
        "optimizer": "nag",
        "lr_scheduler": "fixed",
        "task": "speech_recognition",
        "num_workers": 0,
        "skip_invalid_size_inputs_valid_test": False,
        "max_tokens": 10000000,
        "max_sentences": None,
        "required_batch_size_multiple": 8,
        "dataset_impl": None,
        "gen_subset": "test",
        "num_shards": 1,
        "shard_id": 0,
        "remove_bpe": None,
        "quiet": False,
        "model_overrides": "{}",
        "results_path": None,
        "beam": 40,
        "nbest": 1,
        "max_len_a": 0,
        "max_len_b": 200,
        "min_len": 1,
        "match_source_len": False,
        "no_early_stop": False,
        "unnormalized": False,
        "no_beamable_mm": False,
        "lenpen": 1,
        "unkpen": 0,
        "replace_unk": None,
        "sacrebleu": False,
        "score_reference": False,
        "prefix_size": 0,
        "no_repeat_ngram_size": 0,
        "sampling": False,
        "sampling_topk": -1,
        "sampling_topp": -1.0,
        "temperature": 1.0,
        "diverse_beam_groups": -1,
        "diverse_beam_strength": 0.5,
        "print_alignment": False,
        "ctc": False,
        "rnnt": False,
        "kspmodel": None,
        "wfstlm": None,
        "rnnt_decoding_type": "greedy",
        "lm_weight": 0.2,
        "rnnt_len_penalty": -0.5,
        "momentum": 0.99,
        "weight_decay": 0.0,
        "force_anneal": None,
        "lr_shrink": 0.1,
        "warmup_updates": 0,
    }

    arguments_dict["path"] = os.environ.get("ASR_MODEL_PATH", None)
    arguments_dict["input_file"] = os.environ.get("ASR_INPUT_FILE", None)
    arguments_dict["data"] = os.environ.get("ASR_DATA_PATH", None)
    arguments_dict["user_dir"] = os.environ.get("ASR_USER_DIR", None)
    args = argparse.Namespace(**arguments_dict)

    def test_transcribe_file(self):
        task, generator, models, sp, tgt_dict = setup_asr(self.args, self.logger)
        _, transcription = transcribe_file(self.args, task, generator, models, sp, tgt_dict)

        expected_transcription = [["THE QUICK BROWN FOX JUMPS OVER THE LAZY DOG"]]
        self.assertEqual(transcription, expected_transcription, msg=str(transcription))


if __name__ == "__main__":
    unittest.main()