#!/bin/env python

# Given threee files:
# token_timings (input) - containing:  start_t stop_t "token"
# master label file (input) - label file generated by HTK with class information
# output file
#
# merge the two input files to create a new timing file with the class label

# standard libraries
import os
import sys
import re

import pdb

# custom libraries
sys.path.append("c:/bin-htk/lib")
sys.path.append(os.path.expanduser("~/bin-htk/lib"))

import labels

def read_times(triton_labels):
    # read_times - Parse Triton .lab file with token names and times
    # Returns dictionary keyed on token

    label_file = labels.WaveSurferReader(triton_labels)
    (filename, tokens, starts, stops) = label_file.get_label_file()

    # build hash table
    dictionary = dict()
    for i in xrange(len(tokens)):
        dictionary[tokens[i]] = (starts[i], stops[i])

    return dictionary
    
    
def main():

    if len(sys.argv) != 4:        raise ValueError, "Need token file and master label file"

    token_timestamps = read_times(sys.argv[1])

    htk_h = open(sys.argv[2])

    output_labels = labels.WaveSurferWriter(sys.argv[3])

    re_file = re.compile('File: (?P<file>.*)')
    re_score = re.compile('(?P<class>\w+)\s*==\s*\[(?P<frames>\d+)\s*frames\]\s*(?P<loglkhd>-?\d+\.?\d*)\s+.*')
    
    # Search for the first line with File:
    line = htk_h.readline()
    file_match = re_file.match(line)
    while not file_match:
        line = htk_h.readline()
        file_match = re_file.match(line)

    done = False
    while not done:
        line = htk_h.readline()
        score_match = re_score.match(line)
        if not score_match:
            raise ValueError, "Expected line with class and score"

        # have current file and score

        # find class of token
        filename = file_match.group("file")
        # Get information for current region
        (class_name, frame_count, loglikelihood) = score_match.groups()

        # retrieve the representations of start/stop times
        (start, stop) = token_timestamps['"%s"'%(os.path.basename(filename))]

        output_labels.add_entry(['%s:F%d:L%d'%(
            class_name, int(frame_count), int(round(float(loglikelihood))))],
                                float(start), float(stop))

        # ready for next one
        line = htk_h.readline()
        if line == "":
            done = True         # end of file
        else:
            file_match = re_file.match(line)
            if not file_match:
                raise ValueError, "Expected next file"

            
        
main()