# -*- coding:utf-8 -*-
__author__ = 'Deeper'
import tensorflow as tf
import numpy as np
import os
import codecs
from six.moves import cPickle, reduce, map
from collections import Counter
import random
import shutil


class SpeechLoader():
    def __init__(self, wav_path=None, label_file=None, batch_size=1, n_mfcc=20, encoding='utf-8'):
        self.batch_size = batch_size
        self.encoding = encoding
        self.n_mfcc = n_mfcc

        # path setting
        data_dir = os.path.join(os.getcwd(), 'cache', 'mfcc' + str(n_mfcc))
        if os.path.isdir(data_dir):
            None
        # shutil.rmtree(data_dir)
        else:
            os.mkdir(data_dir)
        # data cache
        wavs_file = os.path.join(data_dir, "wavs.file")
        vocab_file = os.path.join(data_dir, "vocab.file")
        mfcc_tensor = os.path.join(data_dir, "mfcc.tensor")
        label_tensor = os.path.join(data_dir, "label.tensor")

        # data process
        if not (os.path.exists(vocab_file) and os.path.exists(mfcc_tensor) and os.path.exists(label_tensor)):
            print("reading wav files")
            self.preprocess(wav_path, label_file, wavs_file, vocab_file, mfcc_tensor, label_tensor)
        else:
            print("loading preprocessed files")
            self.load_preprocessed(vocab_file, mfcc_tensor, label_tensor)

        # minibatch
        self.create_batches()
        # pointer reset
        self.reset_batch_pointer()

    def preprocess(self, wav_path, label_file, wavs_file, vocab_file, mfcc_tensor, label_tensor):
        def handle_file(dirpath, filename):
            if filename.endswith('.wav') or filename.endswith('.WAV'):
                filename_path = os.path.join(dirpath, filename)
                if os.stat(filename_path).st_size < 240000:
                    return
                return filename_path

        # read label file
        labels_dict = {}
        with codecs.open(label_file, "r", encoding=self.encoding) as f:
            for label in f:
                label = label.strip('\n')
                labels_id = label.split(' ', 1)[0]
                labels_text = label.split(' ', 1)[1]
                labels_dict[labels_id] = labels_text
        # print("",len(labels_dict)) # 10000

        # wav files
        wav_files = []
        if wav_path:
            for (dirpath, dirnames, filenames) in os.walk(wav_path):
                for filename in filenames:
                    if handle_file(dirpath, filename):
                        wav_files.append(handle_file(dirpath, filename))
        print "初始样本数：", len(wav_files)  # 样本数

        # data filter and feature extraction
        wav_files_filter = []
        labels_filter = []
        self.mfcc_tensor = []
        self.wav_max_len = 0
        cnt = 0
        for wav_file in wav_files:
            wav_id = os.path.basename(wav_file).split('.')[0]
            if wav_id in labels_dict:
                print '样本' + str(cnt), wav_file
                labels_filter.append(labels_dict[wav_id])
                wav_files_filter.append(wav_file)
                # mfcc feature
                wav_file, sr = librosa.load(wav_file, mono=True)
                mfcc = np.transpose(librosa.feature.mfcc(wav_file, sr, n_mfcc=self.n_mfcc), [1, 0])
                self.mfcc_tensor.append(mfcc.tolist())
                cnt += 1
        self.wav_max_len = max(len(mfcc) for mfcc in self.mfcc_tensor)
        print "样本总数：", cnt
        print "最长的语音：", self.wav_max_len
        # print(len(wav_files_filter), len(labels_filter),len(wav2mfcc)) # assert check dimensions

        with open(wavs_file, 'wb') as f:
            cPickle.dump(wav_files_filter, f)

        with open(mfcc_tensor, 'wb') as f:
            cPickle.dump(self.mfcc_tensor, f)

        # vocab file
        vocabs = []
        for label in labels_filter:
            vocabs += [word for word in label]
        count = Counter(vocabs)
        count_pairs = sorted(count.items(), key=lambda x: -x[1])
        words, _ = zip(*count_pairs)
        self.wordmap = dict(zip(words, range(len(words))))

        self.vocab_size = len(words)
        print "词汇表大小：", len(words)

        with open(vocab_file, 'wb') as f:
            cPickle.dump(self.wordmap, f)

        # label vector
        label_encoder = lambda word: self.wordmap.get(word, len(words))
        self.label_tensor = [list(map(label_encoder, label)) for label in labels_filter]
        self.label_max_len = max(len(label) for label in self.label_tensor)
        print "最长的句子：", self.label_max_len

        with open(label_tensor, 'wb') as f:
            cPickle.dump(self.label_tensor, f)

    def load_preprocessed(self, vocab_file, mfcc_tensor, label_tensor):
        with open(vocab_file, 'rb') as f:
            self.wordmap = cPickle.load(f)
        self.vocab_size = len(self.wordmap)
        print "词汇表大小：", self.vocab_size

        with open(mfcc_tensor, 'rb') as f:
            self.mfcc_tensor = cPickle.load(f)
        self.wav_max_len = max(len(mfcc) for mfcc in self.mfcc_tensor)
        print "最长的语音：", self.wav_max_len

        with open(label_tensor, 'rb') as f:
            self.label_tensor = cPickle.load(f)
        self.label_max_len = max(len(label) for label in self.label_tensor)
        print "最长的句子：", self.label_max_len

    def create_batches(self):
        self.n_batches = len(self.mfcc_tensor) // self.batch_size
        if self.n_batches == 0:
            assert False, "Not enough data. Make seq_length and batch_size small."

        self.mfcc_tensor = self.mfcc_tensor[:self.n_batches * self.batch_size]
        self.label_tensor = self.label_tensor[:self.n_batches * self.batch_size]

        # random shuffle the data
        if len(self.mfcc_tensor) != len(self.label_tensor):
            assert False, "Data length does not match the label length!"

        data_tensor = []
        for i in range(len(self.mfcc_tensor)):
            data_tensor.append([self.mfcc_tensor[i], self.label_tensor[i]])

        random.shuffle(data_tensor)
        self.mfcc_tensor = []
        self.label_tensor = []
        for i in range(len(data_tensor)):
            self.mfcc_tensor.append(data_tensor[i][0])
            self.label_tensor.append(data_tensor[i][1])

        # create batches
        self.x_batches = []
        self.y_batches = []

        for i in range(self.n_batches):
            from_index = i * self.batch_size
            to_index = from_index + self.batch_size
            mfcc_batches = self.mfcc_tensor[from_index:to_index]
            label_batches = self.label_tensor[from_index:to_index]
            # 补零对齐
            for mfcc in mfcc_batches:
                while len(mfcc) < self.wav_max_len:
                    mfcc.append([0] * self.n_mfcc)
            for label in label_batches:
                while len(label) < self.label_max_len:
                    label.append(0)

            self.x_batches.append(mfcc_batches)
            self.y_batches.append(label_batches)

    def next_batch(self):
        x, y = self.x_batches[self.pointer], self.y_batches[self.pointer]
        self.pointer += 1
        return x, y

    def reset_batch_pointer(self):
        self.pointer = 0
