# -*- coding:utf-8 -*-
#!/usr/bin/python

''' Language modelling for automatic speech recognition based on n-gram
author:

      iiiiiiiiiiii            iiiiiiiiiiii         !!!!!!!             !!!!!!
      #        ###            #        ###           ###        I#        #:
      #      ###              #      I##;             ##;       ##       ##
            ###                     ###               !##      ####      #
           ###                     ###                 ###    ## ###    #'
         !##;                    `##%                   ##;  ##   ###  ##
        ###                     ###                     $## `#     ##  #
       ###        #            ###        #              ####      ####;
     `###        -#           ###        `#               ###       ###
     ##############          ##############               `#         #

date:2017-04-17
'''
import numpy as np
import os
import operator
import pickle

def save_obj(name, obj):
    with open(name + '.pkl', 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)

class NGram:
    def __init__(self, rootdir):
        self.rootdir = rootdir

    def get_corpus(self):
        corpus = []
        word_count = {}
        biword_count = {}
        bigram = {}
        bigram['SOS'] = {}
        trigram = {}
        for subdir, dirs, files in os.walk(self.rootdir):
            for f in files:
                fullFilename = os.path.join(subdir, f)
                filenameNoSuffix =  os.path.splitext(fullFilename)[0]
            if f.endswith('.label'):
                with open(fullFilename, 'r') as f:
                    line = f.readline()
                    corpus.append(line)
                    line = line.strip().split(' ')
                    len_sent = range(len(line))
                    for idx in len_sent:
                        word = line[idx]
                        word_count = inc_dict(word_count, word)
                        if not bigram.has_key(word):
                            bigram[word] = {}
                        if idx == 0:
                            bigram['SOS'] = inc_dict(bigram['SOS'], word)
                        elif idx != len(line)-1:
                            bigram[word] = inc_dict(bigram[word], line[idx+1])
                        else:
                            bigram[word] = inc_dict(bigram[word], 'EOS')
                        if idx == 0:
                            tri_key = 'SOS ' + word
                        else:
                            tri_key = line[idx-1]+' '+word
                        if not trigram.has_key(tri_key):
                            trigram[tri_key] = {}
                        if idx == len(line)-1:
                            trigram[tri_key] = inc_dict(trigram[tri_key], 'EOS')
                        else:
                            trigram[tri_key] = inc_dict(trigram[tri_key], line[idx+1])

        return corpus, word_count, bigram, trigram

def inc_dict(dic, key):
    if not dic.has_key(key):
        dic[key] = 0
        dic[key] += 1
    return dic


if __name__ == '__main__':
    ngram = NGram('/media/pony/Seagate Expansion Drive/学习/语音识别/ASR数据库/LibriSpeech/')
    corpus, word_count, bigram, trigram = ngram.get_corpus()
    savedir = '/home/pony/github/data/libri/ngram/'
    save_obj(savedir+'corpus', corpus)
    save_obj(savedir+'word_count', word_count)
    save_obj(savedir+'bigram', bigram)
    save_obj(savedir+'trigram', trigram)
    #sorted_word_count = sorted(word_count.items(), key=operator.itemgetter(1), reverse=True)
