#!/usr/bin/env python3

# Copyright 2020 Nagoya University (Wen-Chin Huang)
#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)

import argparse
import codecs
import nltk
import os

from text.cleaners import custom_english_cleaners

E_lang_tag="en_US"

try:
    # For phoneme conversion, use https://github.com/Kyubyong/g2p.
    from g2p_en import G2p
    f_g2p = G2p()
    f_g2p("")
except ImportError:
    raise ImportError("g2p_en is not installed. please run `. ./path.sh && pip install g2p_en`.")
except LookupError:
    # NOTE: we need to download dict in initial running
    import ssl
    try:
        _create_unverified_https_context = ssl._create_unverified_context
    except AttributeError:
        pass
    else:
        ssl._create_default_https_context = _create_unverified_https_context
    nltk.download("punkt")


def g2p(text):
    """Convert grapheme to phoneme."""
    tokens = filter(lambda s: s != " ", f_g2p(text))
    return ' '.join(tokens)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('transcription_path', type=str, help='path for the transcription text file')
    parser.add_argument('utt2spk', type=str, help='utt2spk file for the speaker')
    parser.add_argument("trans_type", type=str, default="char",
                        choices=["char", "phn"],
                        help="Input transcription type")
    parser.add_argument('lang_tag', type=str, help='lang tag')
    parser.add_argument('spk', type=str, help='speaker name')
    parser.add_argument('--transcription_path_en', type=str, default=None,
                        help='path for the English transcription text file')
    args = parser.parse_args()

    # clean every line in transcription file first
    transcription_dict = {}
    with codecs.open(args.transcription_path, 'r', 'utf-8') as fid:
        for line in fid.readlines():
            segments = line.split(" ")
            lang_char = args.transcription_path.split('/')[-1][0]
            id = args.spk + '_' + lang_char + segments[0] # ex. TGF1_M10001
            content = ' '.join(segments[1:])
            text = content.rstrip().upper()
            if args.trans_type == "phn":
                # NOTE: we don't have phone for German yet.
                clean_content = text
            else:
                clean_content = text

            if args.lang_tag:
                transcription_dict[id] = "<" + args.lang_tag + "> " + clean_content
            else:
                transcription_dict[id] = clean_content
    
    if args.transcription_path_en:
        with codecs.open(args.transcription_path_en, 'r', 'utf-8') as fid:
            for line in fid.readlines():
                segments = line.split(" ")
                id = args.spk + '_' + "E" + segments[0] # ex. TGF1_E10001
                content = ' '.join(segments[1:])
                clean_content = custom_english_cleaners(content.rstrip())
                if args.trans_type == "phn":
                    text = clean_content.lower()
                    clean_content = g2p(text)

                transcription_dict[id] = "<" + E_lang_tag + "> " + clean_content
    
    # read the utt2spk file and actually write
    with codecs.open(args.utt2spk, 'r', 'utf-8') as fid:
        for line in fid.readlines():
            segments = line.split(" ")
            id = segments[0] # ex. E10001
            content = transcription_dict[id]

            print("%s %s" % (id, content))
