#!/usr/bin/env python3

# 2019 Dongji Gao
# Apache 2.0.

from make_lexicon_fst import read_lexiconp
import argparse
import math
import sys

# see get_args() below for usage mesage
def get_args():
    parser = argparse.ArgumentParser(description="""This script creates the
        text form of a subword lexicon FST to be compiled by fstcompile using
        the appropriate symbol tables (phones.txt and words.txt). It will mostly
        be invoked indirectly via utils/prepare_lang_subword.sh. The output
        goes to the stdout. This script is the subword version of make_lexicon_fst.py.
        It only allows optional silence to appear after end-subword or singleton-subword,
        (i.e., subwords without separator). In this version we do not support
        pronunciation probability. (i.e., pron-prob = 1.0)""")

    parser.add_argument('--sil-phone', type=str, help="""Text form of
        optional-silence phone, e.g. 'SIL'. See also the --sil-prob option.""")
    parser.add_argument('--sil-prob', type=float, default=0.0, help="""Probability
        of silence between words (including the beginning and end of word sequence).
        Must be in range [0.0, 1.0). This refer to the optional silence inserted by
        the lexicon; see the --sil-phone option.""")
    parser.add_argument('--sil-disambig', type=str, help="""Disambiguation symbol
        to disambiguate silence, e.g. #5. Will only be supplied if you are creating 
        the version of L.fst with disambiguation symbols, intended for use with cyclic 
        G.fst. This symbol was introduced to fix a rather obscure source of nondeterminism 
        of CLG.fst, that has to do with reordering of disambiguation symbols and phone symbols.""")
    parser.add_argument('--position-dependent', action="store_true", help="""Whether 
        the input lexicon is position-dependent.""")
    parser.add_argument("--separator", type=str, default="@@", help="""Separator
        indicates the position of a subword in a word.
        Subword followed by separator can only appear at the beginning or middle of a word.
        Subword without separator can only appear at the end of a word or is a word itself.
        E.g. "international -> inter@@ nation@@ al";
             "nation        -> nation"
    The separator should match the separator used in the input lexicon.""")
    parser.add_argument('lexiconp', type=str, help="""Filename of lexicon with
        pronunciation probabilities (normally lexiconp.txt), with lines of the
        form 'subword prob p1 p2...', e.g. 'a, 1.0 ay'""")
    args = parser.parse_args()
    return args

def contain_disambig_symbol(phones):
    """Return true if the phone sequence contains disambiguation symbol.
    Return false otherwise. Disambiguation symbol is at the end of phones 
    in the form of #1, #2... There is at most one disambiguation 
    symbol for each phone sequence"""
    return True if phones[-1].startswith("#") else False

def print_arc(src, dest, phone, word, cost):
    print('{}\t{}\t{}\t{}\t{}'.format(src, dest, phone, word, cost))

def is_end(word, separator):
    """Return true if the subword can appear at the end of a word (i.e., the subword
    does not end with separator). Return false otherwise."""
    return not word.endswith(separator)

def get_suffix(phone):
    """Return the suffix of a phone. The suffix is in the form of '_B', '_I'..."""
    if len(phone) < 3:
        print("{}: invalid phone {} (please check if the phone is position-dependent)".format(
              sys.argv[0], phone), file=sys.stderr)
        sys.exit(1)
    return phone[-2:]

def write_fst_no_silence(lexicon, position_dependent, separator):
    """Writes the text format of L.fst to the standard output.  This version is for
    when --sil-prob=0.0, meaning there is no optional silence allowed.
    loop_state here is the start and final state of the fst. It goes to word_start_state
    via epsilon transition.
    In position-independent case, there is no difference between beginning word and 
    middle word. So all subwords with separator would leave from and enter word_start_state.
    All subword without separator would leave from word_start_state and enter loop_state.
    This guarantees that optional silence can only follow a word-end subword.

    In position-dependent case, there are 4 types of position-dependent subword:
    1) Beginning subword. The first phone suffix should be "_B" and other suffixes should be "_I"s:
        nation@@ 1.0 n_B ey_I sh_I ih_I n_I
        n@@      1.0 n_B
    2) Middle subword. All phone suffixes should be "_I"s:
        nation@@ 1.0 n_I ey_I sh_I ih_I n_I
    3) End subword. The last phone suffix should be "_E" and other suffixes be should "_I"s:
        nation   1.0 n_I ey_I sh_I ih_I n_E
        n        1.0 n_E
    4) Singleton subword (i.e., the subword is word it self).
       The first phone suffix should be "_B" and the last suffix should be "_E".
       All other suffix should be "_I"s. If there is only one phone, its suffix should be "_S":
        nation   1.0 n_B ey_I sh_I ih_I n_E
        n        1.0 n_S

    So we need an extra word_internal_state. The beginning word 
    would leave from word_start_state and enter word_internal_state and middle word
    would leave from and enter word_internal_state. The rest part is same.

      'lexicon' is a list of 3-tuples (subword, pron-prob, prons) as returned by
      'position_dependent', which is true is the lexicon is position-dependent.
      'separator' is a symbol which indicates the position of a subword in word.
    """
    # regular setting
    loop_state = 0
    word_start_state = 1
    next_state = 2

    print_arc(loop_state, word_start_state, "<eps>", "<eps>", 0.0)

    # optional setting for word_internal_state
    if position_dependent:
        word_internal_state = next_state
        next_state += 1

    for (word, pron_prob, phones) in lexicon:
        pron_cost = 0.0                # do not support pron_prob
        phones_len = len(phones)

        # set start and end state for different cases
        if position_dependent:
            first_phone_suffix = get_suffix(phones[0])
            last_phone = phones[-2] if contain_disambig_symbol(phones) else phones[-1]
            last_phone_suffix = get_suffix(last_phone)

            # singleton word
            if first_phone_suffix == "_S":
                current_state = word_start_state
                end_state = loop_state
            # set the current_state
            elif first_phone_suffix == "_B":
                current_state = word_start_state
            elif first_phone_suffix == "_I" or first_phone_suffix == "_E":
                current_state = word_internal_state
            # then set the end_state
            if last_phone_suffix == "_B" or last_phone_suffix == "_I":
                end_state = word_internal_state
            elif last_phone_suffix == "_E":
                end_state = loop_state
        else:
            current_state = word_start_state
            end_state = loop_state if is_end(word, separator) else word_start_state

        # print arcs (except the last one) for the subword
        for i in range(phones_len - 1):
            word = word if i == 0 else "<eps>"
            cost = pron_cost if i == 0 else 0.0
            print_arc(current_state, next_state, phones[i], word, cost)
            current_state = next_state
            next_state += 1

        # print the last arc
        i = phones_len - 1
        phone = phones[i] if i >=0 else "<eps>"
        word = word if i <= 0 else "<eps>"
        cost = pron_cost if i <= 0 else 0.0
        print_arc(current_state, end_state, phone, word, cost)

    # set the final state
    print("{state}\t{final_cost}".format(state=loop_state, final_cost=0.0))

def write_fst_with_silence(lexicon, sil_phone, sil_prob, sil_disambig, position_dependent, separator):
    """Writes the text format of L.fst to the standard output.  This version is for
    when --sil-prob=0.0, meaning there is no optional silence allowed.
    loop_state here is the start and final state of the fst. It goes to word_start_state
    via epsilon transition.

    In position-independent case, there is no difference between beginning word and 
    middle word. So all subwords with separator would leave from and enter word_start_state.
    All subword without separator would leave from word_start_state and enter sil_state.
    This guarantees that optional silence can only follow a word-end subword and such subwords
    must appear at the end of the whole subword sequence.

    In position-dependent case, there are 4 types of position-dependent subword:
    1) Beginning subword. The first phone suffix should be "_B" and other suffixes should be "_I"s:
        nation@@ 1.0 n_B ey_I sh_I ih_I n_I
        n@@      1.0 n_B
    2) Middle subword. All phone suffixes should be "_I"s:
        nation@@ 1.0 n_I ey_I sh_I ih_I n_I
    3) End subword. The last phone suffix should be "_E" and other suffixes be should "_I"s:
        nation   1.0 n_I ey_I sh_I ih_I n_E
        n        1.0 n_E
    4) Singleton subword (i.e., the subword is word it self).
       The first phone suffix should be "_B" and the last suffix should be "_E".
       All other suffix should be "_I"s. If there is only one phone, its suffix should be "_S":
        nation   1.0 n_B ey_I sh_I ih_I n_E
        n        1.0 n_S

    So we need an extra word_internal_state. The beginning word 
    would leave from word_start_state and enter word_internal_state and middle word
    would leave from and enter word_internal_state. The rest part is same.

      'lexicon' is a list of 3-tuples (subword, pron-prob, prons)
         as returned by read_lexiconp().
      'sil_prob', which is expected to be strictly between 0.0 and 1.0, is the
         probability of silence
      'sil_phone' is the silence phone, e.g. "SIL".
      'sil_disambig' is either None, or the silence disambiguation symbol, e.g. "#5".
      'position_dependent', which is True is the lexicion is position-dependent.
      'separator' is the symbol we use to indicate the position of a subword in word.
    """

    sil_cost = -math.log(sil_prob)
    no_sil_cost = -math.log(1 - sil_prob)

    # regular setting
    start_state = 0
    loop_state = 1         # also the final state
    sil_state = 2          # words terminate here when followed by silence; this state
                           # has a licence transition to loop_state
    word_start_state = 3   # subword leave from here
    next_state = 4         # the next un-allocated state, will be incremented as we go

    print_arc(start_state, loop_state, "<eps>", "<eps>", no_sil_cost)
    print_arc(start_state, sil_state, "<eps>", "<eps>", sil_cost)
    print_arc(loop_state, word_start_state, "<eps>", "<eps>", 0.0)

    # optional setting for disambig_state
    if sil_disambig is None:
        print_arc(sil_state, loop_state, sil_phone, "<eps>", 0.0)
    else:
        disambig_state = next_state
        next_state += 1
        print_arc(sil_state, disambig_state, sil_phone, "<eps>", 0.0)
        print_arc(disambig_state, loop_state, sil_disambig, "<eps>", 0.0)

    # optional setting for word_internal_state
    if position_dependent:
        word_internal_state = next_state
        next_state += 1

    for (word, pron_prob, phones) in lexicon:
        pron_cost = 0.0           # do not support pron_prob
        phones_len = len(phones)
        
        # set start and end state for different cases
        if position_dependent:
            first_phone_suffix = get_suffix(phones[0])
            last_phone = phones[-2] if contain_disambig_symbol(phones) else phones[-1]
            last_phone_suffix = get_suffix(last_phone)

            # singleton subword
            if first_phone_suffix == "_S":
                current_state = word_start_state
                end_state_list = [loop_state, sil_state]
                end_cost_list = [no_sil_cost, sil_cost]
            # first set the current_state
            elif first_phone_suffix == "_B":
                current_state = word_start_state
            elif first_phone_suffix == "_I" or first_phone_suffix == "_E":
                current_state = word_internal_state
            # then set the end_state (end_state_list)
            if last_phone_suffix == "_B" or last_phone_suffix == "_I":
                end_state_list = [word_internal_state]
                end_cost_list = [0.0]
            elif last_phone_suffix == "_E":
                end_state_list = [loop_state, sil_state]
                end_cost_list = [no_sil_cost, sil_cost]
        else:
            current_state = word_start_state
            if is_end(word, separator):
                end_state_list = [loop_state, sil_state]
                end_cost_list = [no_sil_cost, sil_cost]
            else:
                end_state_list = [word_start_state]
                end_cost_list = [0.0]

        # print arcs (except the last one) for the subword
        for i in range(phones_len - 1):
            word = word if i == 0 else "<eps>"
            cost = pron_cost if i == 0 else 0.0
            print_arc(current_state, next_state, phones[i], word, cost)
            current_state = next_state
            next_state += 1

        # print the last arc
        i = phones_len - 1
        phone = phones[i] if i >= 0 else "<eps>"
        word = word if i <= 0 else "<eps>"
        cost = pron_cost if i <= 0 else 0.0
        for (end_state, end_cost) in zip(end_state_list, end_cost_list):
            print_arc(current_state, end_state, phone, word, cost + end_cost)

    # set the final state
    print("{state}\t{final_cost}".format(state=loop_state, final_cost=0.0))

def main():
    args = get_args()
    if args.sil_prob < 0.0 or args.sil_prob >= 1.0:
        print("{}: invalid value specified --sil-prob={}".format(
              sys.argv[0], args.sil_prob), file=sys.stderr)
        sys.exit(1)
    lexicon = read_lexiconp(args.lexiconp)
    if args.sil_prob == 0.0:
        write_fst_no_silence(lexicon, args.position_dependent, args.separator)
    else:
        write_fst_with_silence(lexicon, args.sil_phone, args.sil_prob, 
            args.sil_disambig, args.position_dependent, args.separator)

if __name__ == "__main__":
    main()