import math
import logging
import wave

from gentle import transcription

from multiprocessing.pool import ThreadPool as Pool

class MultiThreadedTranscriber:
    def __init__(self, kaldi_queue, chunk_len=20, overlap_t=2, nthreads=4):
        self.chunk_len = chunk_len
        self.overlap_t = overlap_t
        self.nthreads = nthreads
            
        self.kaldi_queue = kaldi_queue

    def transcribe(self, wavfile, progress_cb=None):
        wav_obj = wave.open(wavfile, 'r')
        duration = wav_obj.getnframes() / float(wav_obj.getframerate())
        n_chunks = int(math.ceil(duration / float(self.chunk_len - self.overlap_t)))

        chunks = []

        def transcribe_chunk(idx):
            wav_obj = wave.open(wavfile, 'r')
            start_t = idx * (self.chunk_len - self.overlap_t)
            # Seek
            wav_obj.setpos(int(start_t * wav_obj.getframerate()))
            # Read frames
            buf = wav_obj.readframes(int(self.chunk_len * wav_obj.getframerate()))

            k = self.kaldi_queue.get()
            k.push_chunk(buf)
            ret = k.get_final()
            k.reset()
            self.kaldi_queue.put(k)

            chunks.append({"start": start_t, "words": ret})
            logging.info('%d/%d' % (len(chunks), n_chunks))
            if progress_cb is not None:
                progress_cb({"message": ' '.join([X['word'] for X in ret]),
                             "percent": len(chunks) / float(n_chunks)})


        pool = Pool(min(n_chunks, self.nthreads))
        pool.map(transcribe_chunk, range(n_chunks))
        pool.close()
        
        chunks.sort(key=lambda x: x['start'])

        # Combine chunks
        words = [transcription.Word(**wd).shift(time=c['start']) for c in chunks for wd in c['words']]

        # Remove overlap:  Sort by time, then filter out any Word entries in
        # the list that are adjacent to another entry corresponding to the same
        # word in the audio.
        words.sort(key=lambda word: word.start)
        words.append(transcription.Word(word="__dummy__"))
        words = [words[i] for i in range(len(words)-1) if not words[i].corresponds(words[i+1])]

        return words


if __name__=='__main__':
    # full transcription
    from Queue import Queue
    from util import ffmpeg
    from gentle import standard_kaldi

    import sys

    import logging
    logging.getLogger().setLevel('INFO')
    
    k_queue = Queue()
    for i in range(3):
        k_queue.put(standard_kaldi.Kaldi())

    trans = MultiThreadedTranscriber(k_queue)

    with gentle.resampled(sys.argv[1]) as filename:
        out = trans.transcribe(filename)

    open(sys.argv[2], 'w').write(out.to_json())


@@ -1691,93 +1691,1056 @@
  = %5B
-transcription.Word(**wd).shift(time=c%5B'start'%5D) for c in chunks for wd in c%5B'words'%5D%5D
+%5D%0A        for c in chunks:%0A            chunk_start = c%5B'start'%5D%0A            chunk_end = chunk_start + self.chunk_len%0A%0A            chunk_words = %5Btranscription.Word(**wd).shift(time=chunk_start) for wd in c%5B'words'%5D%5D%0A%0A            # At chunk boundary cut points the audio often contains part of a%0A            # word, which can get erroneously identified as one or more different%0A            # in-vocabulary words.  So discard one or more words near the cut points%0A            # (they'll be covered by the ovlerap anyway).%0A            #%0A            trim = min(0.25 * self.overlap_t, 0.5)%0A            if c is not chunks%5B0%5D:%0A                while len(chunk_words) %3E 1:%0A                    chunk_words.pop(0)%0A                    if chunk_words%5B0%5D.end %3E chunk_start + trim:%0A                        break%0A            if c is not chunks%5B-1%5D:%0A                while len(chunk_words) %3E 1:%0A                    chunk_words.pop()%0A                    if chunk_words%5B-1%5D.start %3C chunk_end - trim:%0A                        break%0A%0A            words.extend(chunk_words)
 %0A%0A  

Discard Words adjacent to cut points between chunks, since they're likely to be erroneous identifications of part of a word that was split at the cut.