import math


def write_nonterminal_arcs(start_state, sil_state, non_sil_state,
                           next_state, sil_phone,
                           nonterminals, left_context_phones):
    """This function relates to the grammar-decoding setup, see
    kaldi-asr.org/doc/grammar.html.  It is called from write_fst, and writes to
    the stdout some extra arcs in the lexicon FST that relate to nonterminal
    symbols.

    See the section "Special symbols in L.fst,
    kaldi-asr.org/doc/grammar.html#grammar_special_l.
       start_state: the start-state of L.fst.
       sil_state:  the state of high out-degree in L.fst where words leave
                   when preceded by optional silence
       non_sil_state:   the state of high out-degree in L.fst where words leave
                   when not preceded by optional silence
       next_state: the number from which this function can start allocating its
                  own states.  the updated value of next_state will be returned.
       sil_phone:  the optional-silence phone (a string, e.g 'sil')
       nonterminals: the user-defined nonterminal symbols as a list of
          strings, e.g. ['#nonterm:contact_list', ... ].
       left_context_phones: a list of phones that may appear as left-context,
          e.g. ['a', 'ah', ... '#nonterm_bos'].
    """
    shared_state = next_state
    next_state += 1
    final_state = next_state
    next_state += 1

    print("{src}\t{dest}\t{phone}\t{word}\t{cost}".format(
        src=start_state, dest=shared_state,
        phone='#nonterm_begin', word='#nonterm_begin',
        cost=0.0))

    for nonterminal in nonterminals:
        # What we are doing here could be viewed as a little lazy, by going to
        # 'shared_state' instead of a state specific to nonsilence vs. silence
        # left-context vs. unknown (for #nonterm_begin).  If we made them
        # separate we could improve (by half) the correctness of how it
        # interacts with sil-probs in the hard-to-handle case where
        # word-position-dependent phones are not used and some words end
        # in the optional-silence phone.
        for src in [sil_state, non_sil_state]:
            print("{src}\t{dest}\t{phone}\t{word}\t{cost}".format(
                src=src, dest=shared_state,
                phone=nonterminal, word=nonterminal,
                cost=0.0))

    # this_cost equals log(len(left_context_phones)) but the expression below
    # better captures the meaning.  Applying this cost to arcs keeps the FST
    # stochatic (sum-to-one, like an HMM), so that if we do weight pushing
    # things won't get weird.  In the grammar-FST code when we splice things
    # together we will cancel out this cost, see the function CombineArcs().
    this_cost = -math.log(1.0 / len(left_context_phones))

    for left_context_phone in left_context_phones:
        # The following line is part of how we get this to interact correctly with
        # the silence probabilities: if the left-context phone was the silence
        # phone, it goes to sil_state, else nonsil_state.  This won't always
        # do the right thing if you have a system without word-position-dependent
        # phones (--position-dependent-phones false to prepare_lang.sh) and
        # you have words that end in the optional-silence phone.
        dest = (sil_state if left_context_phone == sil_phone else non_sil_state)

        print("{src}\t{dest}\t{phone}\t{word}\t{cost}".format(
            src=shared_state, dest=dest,
            phone=left_context_phone, word='<eps>', cost=this_cost))

    # arc from sil_state and non_sil_state to a final-state with #nonterm_end as
    # ilabel and olabel.  The costs on these arcs are zero because if you take
    # that arc, you are not really terminating the sequence, you are just
    # skipping to sil_state or non_sil_state in the FST one level up.  It
    # takes the correct path because of the code around 'dest = ...' a few
    # lines above this, after reaching 'shared_state' because it saw the
    # user-defined nonterminal.
    for src in [sil_state, non_sil_state]:
        print("{src}\t{dest}\t{phone}\t{word}\t{cost}".format(
            src=src, dest=final_state,
            phone='#nonterm_end', word='#nonterm_end', cost=0.0))
    print("{state}\t{final_cost}".format(
        state=final_state, final_cost=0.0))
    return next_state