'''@file readfiles.py
contains functions to read some files created by Kaldi
'''

import gzip
from collections import OrderedDict
import numpy as np

def read_alignments(filename):
    '''
    read the alignment file generated by kaldi

    Args:
        filename: path to alignment file

    Returns:
        a dictionary containing:
            - key: the utterance ID
            - value: the alignments
    '''

    with gzip.open(filename, 'rb') as fid:
        alignments = {}
        for line in fid:
            data = line.replace(' \n', '').split(' ')
            #segment:alignment
            alignments[data[0]] = np.asarray(map(int, data[1:len(data)]))
    return alignments

def read_segments(filename):
    '''
    read a segment file that is used in kaldi

    Args:
        filename: path to segment file

    Returns:
        a dictionary containing:
            - key: name of the recording
            - value: a triple containing:
                - the utterance ID
                - the begining of the utterance
                - the end of the utterance
    '''

    with open(filename) as fid:
        segments = OrderedDict()
        for line in fid:
            #seg utt begin end
            data = line.replace('\n', '').split(' ')
            if data[1] not in segments:
                #utt: [(seg , begin, end)]
                segments[data[1]] = [(data[0], float(data[2]), float(data[3]))]
            else:
                segments[data[1]].append((data[0], float(data[2]),
                                          float(data[3])))
    return segments

def read_wavfiles(filename):
    '''
     read the wav.scp file used in kaldi

     Args:
        filename: path to the wav scp file

    Returns:
        a dictionary containing:
            - key: the utterance ID
            - value: a pair containing
                - the filenames
                - bool wheter or not the filename is extended (with a
                    read command)
    '''

    with open(filename) as fid:
        wavfiles = OrderedDict()
        for line in fid:
            data = line.replace('\n', '').split(' ')
            #wav.scp contains filenames
            if len(data) == 2:
                #utterance:(filename, not extended)
                wavfiles[data[0]] = (data[1], False)
            #wav.scp contains extended filenames
            else:
                #utterance: (extended filename, extended)
                wavfiles[data[0]] = (line[len(data[0])+1:len(line)-1], True)
    return wavfiles

def read_utt2spk(filename):
    '''
    Read the utt2spk file used in kaldi

    Args:
        filename: path to the utt2spk file

    Returns:
        a dictionary containing the speaker names with the utterance IDs as keys
    '''

    with open(filename) as fid:
        utt2spk = {}
        for line in fid:
            data = line.replace('\n', '').split(' ')
            utt2spk[data[0]] = data[1]
    return utt2spk
