#!/usr/bin/env python3

# Copyright (c) 2010  Carnegie Mellon University
#
# You may copy and modify this freely under the same terms as
# Sphinx-III
"""
Convert Sphinx lattices to finite state automata and grammars.
"""

__author__ = "David Huggins-Daines <dhdaines@gmail.com>"
__version__ = "$Revision$"

import openfst
from . import lattice
import math
import sys
import os


def baseword(sym):
    """
    Returns base word (no pronunciation variant) for sym.
    """
    paren = sym.rfind('(')
    if paren != -1:
        return sym[0:paren]
    else:
        return sym


def realword(sym):
    """
    Test if a word is a real word (not silence or filler)
    """
    if sym.lower() in ('<s>', '<sil>', '</s>'):
        return False
    if sym.startswith("++"):
        return False
    return True


def build_lattice_fsg(dag,
                      syms=None,
                      ascale=0,
                      pscale=0,
                      addsyms=False,
                      determinize=True,
                      baseword=baseword):
    """
    Build an FSM from a Sphinx word lattice.
    """
    fst = openfst.StdVectorFst()
    if syms is None:
        fsgsyms = openfst.SymbolTable("words")
        fsgsyms.AddSymbol("&epsilon;")
        fsgsyms.AddSymbol("&sigma;")
        fsgsyms.AddSymbol("&rho;")
        fsgsyms.AddSymbol("&phi;")
        addsyms = True
    else:
        fsgsyms = syms
    statemap = {}
    for n in dag.nodes:
        # Skip fillers as they have been "bypassed" by PocketSphinx
        if n.sym.startswith("++") or n.sym == "<sil>":
            continue
        # These should not exist, but they do (!!)
        if n.sym == "<s>" and n.entry != 0:
            continue
        if n not in statemap:
            statemap[n] = fst.AddState()
        if addsyms:
            fsgsyms.AddSymbol(baseword(n.sym))
    statemap[dag.start] = fst.AddState()
    fst.SetStart(statemap[dag.start])
    for n in dag.nodes:
        if n not in statemap:
            continue
        sym = fsgsyms.Find(baseword(n.sym))
        for x in n.exits:
            if x.dest not in statemap:
                continue
            weight = 0
            # Turn OOVs and non-events into epsilons
            if sym == -1 or n.sym == "<s>":
                sym = 0
            if ascale:
                weight = -x.ascr * ascale
            elif pscale:
                weight = -x.post * pscale
            fst.AddArc(statemap[x.src], sym, sym, weight, statemap[x.dest])
    # Add a </s> transition if none exists
    if '</s>' not in [x.src.sym for x in dag.end.entries]:
        end = fst.AddState()
        sym = fsgsyms.AddSymbol("</s>")
        fst.AddArc(statemap[dag.end], sym, sym, 0, end)
        fst.SetFinal(end, 0)
    else:
        fst.SetFinal(statemap[dag.end], 0)
    # Epsilon-remove it (like bypassing fillers...) (FIXME: umm...)
    openfst.RmEpsilon(fst)
    # Don't determinize if it's weighted
    if ascale or pscale:
        determinize = False
    if determinize:
        outfst = openfst.StdVectorFst()
        openfst.Determinize(fst, outfst)
        fst = outfst
    fst.SetInputSymbols(fsgsyms)
    fst.SetOutputSymbols(fsgsyms)
    return fst


def build_fsg_fst(fst, outfile, name="fromfst"):
    """
    Build a Sphinx FSG from an FST.
    """
    outfile.write("FSG_BEGIN %s\n" % name)
    outfile.write("NUM_STATES %d\n" % fst.NumStates())
    outfile.write("START_STATE %d\n" % fst.Start())
    # Assume there's only one final state
    end = None
    for state in fst:
        if fst.IsFinal(state):
            end = state
            break
    if end is None:
        raise RuntimeError("No ending state in lattice, writing %s failed" %
                           outfile)
    outfile.write("FINAL_STATE %d\n" % end)
    syms = fst.OutputSymbols()
    for state in fst:
        for arc in fst.iterarcs(state):
            if arc.ilabel == 0 or syms.Find(arc.ilabel) in ('<s>', '</s>'):
                outfile.write(
                    "TRANSITION %d %d %g\n" %
                    (state, arc.nextstate, math.exp(-arc.weight.Value())))
            else:
                outfile.write(
                    "TRANSITION %d %d %g %s\n" %
                    (state, arc.nextstate, math.exp(-arc.weight.Value()),
                     syms.Find(arc.ilabel)))
    outfile.write("FSG_END\n")
    outfile.flush()


def lat2fsg(lat, fsgfile, lmfst, prune=15):
    if isinstance(lat, str):
        if lat.endswith(".slf"):
            dag = lattice.Dag(htkfile=lat)
        else:
            dag = lattice.Dag(lat)
    else:
        dag = lat
    fst = build_lattice_fsg(dag, lmfst.InputSymbols())
    # Compose it (intersect, really) with the language model to get
    # correct N-gram scores (otherwise it is just a unigram LM).  This
    # is the same thing as "lattice expansion".
    phi = lmfst.InputSymbols().Find("&phi;")
    if phi != -1:
        opts = openfst.StdPhiComposeOptions()
        opts.matcher1 = openfst.StdPhiMatcher(fst, openfst.MATCH_NONE)
        opts.matcher2 = openfst.StdPhiMatcher(lmfst, openfst.MATCH_INPUT, phi)
        cfst = openfst.StdComposeFst(fst, lmfst, opts)
    else:
        cfst = openfst.StdComposeFst(fst, lmfst)
    outfst = openfst.StdVectorFst(cfst)
    openfst.Prune(outfst, prune)
    # Write it back out as an FSG for PocketSphinx.
    build_fsg_fst(outfst, fsgfile)
    return outfst


def apply_errfst(fst, errfst):
    sigma = errfst.InputSymbols().Find("&sigma;")
    opts = openfst.StdSigmaComposeOptions()
    opts.matcher1 = openfst.StdSigmaMatcher(fst, openfst.MATCH_NONE)
    opts.matcher2 = openfst.StdSigmaMatcher(errfst, openfst.MATCH_INPUT, sigma,
                                            True)
    cfst = openfst.StdComposeFst(fst, errfst, opts)
    cfst = openfst.StdVectorFst(cfst)
    openfst.ProjectOutput(cfst)
    return cfst


def lat2fsg_posterior(lat, fsgfile, prune=5, errfst=None):
    if isinstance(lat, str):
        if lat.endswith(".slf"):
            dag = lattice.Dag(htk_file=lat)
        else:
            dag = lattice.Dag(lat)
    else:
        dag = lat
    dag.posterior_prune(-prune)
    fst = build_lattice_fsg(dag, pscale=1)
    if errfst:
        fst = build_lattice_fsg(dag,
                                errfst.InputSymbols(),
                                pscale=1,
                                addsyms=True)
        errfst.SetOutputSymbols(errfst.InputSymbols())
        fst = apply_errfst(fst, errfst)
    else:
        fst = build_lattice_fsg(dag, pscale=1)
    build_fsg_fst(fst, fsgfile)
    return fst


def lat2flat(latfile, fsgfile, lmfst):
    """
    Subset a language model using the vocabulary of a lattice.
    """
    dag = lattice.Dag(latfile)
    fst = openfst.StdVectorFst()
    fst.SetStart(fst.AddState())
    fst.SetFinal(0, 0)
    syms = lmfst.InputSymbols()
    seen = set()
    for n in dag.nodes:
        # Skip fillers as they have been "bypassed" by PocketSphinx
        if n.sym.startswith("++") or n.sym == "<sil>":
            continue
        if n.sym in seen:
            continue
        seen.add(n.sym)
        sym = syms.Find(baseword(n.sym))
        if sym == -1:
            continue
        fst.AddArc(0, sym, sym, 0, 0)
    fst.SetOutputSymbols(lmfst.InputSymbols())
    phi = lmfst.InputSymbols().Find("&phi;")
    if phi != -1:
        opts = openfst.StdPhiComposeOptions()
        opts.matcher1 = openfst.StdPhiMatcher(fst, openfst.MATCH_NONE)
        opts.matcher2 = openfst.StdPhiMatcher(lmfst, openfst.MATCH_INPUT, phi)
        cfst = openfst.StdComposeFst(fst, lmfst, opts)
    else:
        cfst = openfst.StdComposeFst(fst, lmfst)
    outfst = openfst.StdVectorFst()
    openfst.Determinize(cfst, outfst)
    # Write it back out as an FSG for PocketSphinx.
    build_fsg_fst(outfst, fsgfile)
    return outfst


if __name__ == '__main__':
    from optparse import OptionParser
    parser = OptionParser(usage="%prog CTL LATDIR")
    parser.add_option("--prune", type="int", default=5)
    parser.add_option("--errprune", type="float", default=1e-5)
    parser.add_option("--errpen", type="float", default=1.0)
    parser.add_option("--outext")
    parser.add_option("--errfst")
    opts, args = parser.parse_args(sys.argv[1:])
    ctlfile, latdir = args
    errfst = None
    if opts.errfst:
        errfst = openfst.StdVectorFst.Read(opts.errfst)
        openfst.Prune(errfst, opts.errprune)
        sigma = errfst.InputSymbols().Find("&sigma;")
        for node in errfst:
            for ai in errfst.mutable_iterarcs(node):
                arc = ai.Value()
                if arc.ilabel == sigma:
                    continue
                ai.SetValue(
                    openfst.StdArc(arc.ilabel, arc.olabel,
                                   arc.weight.Value() + opts.errpen,
                                   arc.nextstate))
        openfst.ArcSortInput(errfst)
    if opts.outext is None:
        opts.outext = ".fsg%d" % opts.prune
    for spam in open(ctlfile):
        latfile = os.path.join(latdir, spam.strip() + ".slf")
        fsgfile = os.path.join(latdir, spam.strip() + opts.outext)
        print(spam, end=' ')
        ofst = lat2fsg_posterior(latfile, fsgfile, opts.prune, errfst)