import multiprocessing as mp
import subprocess
import shutil
import os

from ..helper import make_path_safe, thirdparty_binary, filter_scp
from ..exceptions import CorpusError


def mfcc_func(directory, job_name, mfcc_config_path):  # pragma: no cover
    log_directory = os.path.join(directory, 'log')
    raw_mfcc_path = os.path.join(directory, 'raw_mfcc.{}.ark'.format(job_name))
    raw_scp_path = os.path.join(directory, 'feats.{}.scp'.format(job_name))
    log_path = os.path.join(log_directory, 'make_mfcc.{}.log'.format(job_name))
    segment_path = os.path.join(directory, 'segments.{}'.format(job_name))
    scp_path = os.path.join(directory, 'wav.{}.scp'.format(job_name))

    with open(log_path, 'w') as f:
        if os.path.exists(segment_path):
            seg_proc = subprocess.Popen([thirdparty_binary('extract-segments'),
                                         'scp,p:' + scp_path, segment_path, 'ark:-'],
                                        stdout=subprocess.PIPE, stderr=f)
            comp_proc = subprocess.Popen([thirdparty_binary('compute-mfcc-feats'), '--verbose=2',
                                          '--config=' + mfcc_config_path,
                                          'ark:-', 'ark:-'],
                                         stdout=subprocess.PIPE, stderr=f, stdin=seg_proc.stdout)
        else:

            comp_proc = subprocess.Popen([thirdparty_binary('compute-mfcc-feats'), '--verbose=2',
                                          '--config=' + mfcc_config_path,
                                          'scp,p:' + scp_path, 'ark:-'],
                                         stdout=subprocess.PIPE, stderr=f)
        copy_proc = subprocess.Popen([thirdparty_binary('copy-feats'),
                                      '--compress=true', 'ark:-',
                                      'ark,scp:{},{}'.format(raw_mfcc_path, raw_scp_path)],
                                     stdin=comp_proc.stdout, stderr=f)
        copy_proc.wait()


def init(env):
    os.environ = env

def mfcc(mfcc_directory, num_jobs, feature_config, frequency_configs):
    """
    Multiprocessing function that converts wav files into MFCCs

    See http://kaldi-asr.org/doc/feat.html and
    http://kaldi-asr.org/doc/compute-mfcc-feats_8cc.html for more details on how
    MFCCs are computed.

    Also see https://github.com/kaldi-asr/kaldi/blob/master/egs/wsj/s5/steps/make_mfcc.sh
    for the bash script this function was based on.

    Parameters
    ----------
    mfcc_directory : str
        Directory to save MFCC feature matrices
    log_directory : str
        Directory to store log files
    num_jobs : int
        The number of processes to use in calculation
    mfcc_configs : list of :class:`~aligner.config.MfccConfig`
        Configuration object for generating MFCCs

    Raises
    ------
    CorpusError
        If the files per speaker exceeds the number of files that are
        allowed to be open on the computer (for Unix-based systems)
    """
    child_env = os.environ.copy()

    os.makedirs(os.path.join(mfcc_directory, 'log'), exist_ok=True)
    paths = []
    for j, p in frequency_configs:
        paths.append(feature_config.write(mfcc_directory, j, p))
    jobs = [(mfcc_directory, x, paths[x])
            for x in range(num_jobs)]
    with mp.Pool(processes=num_jobs, initializer=init, initargs=(child_env,)) as pool:
        r = False
        try:
            results = [pool.apply_async(mfcc_func, args=i) for i in jobs]
            output = [p.get() for p in results]
        except OSError as e:
            print(dir(e))
            if e.errno == 24:
                r = True
            else:
                raise
    if r:
        raise (CorpusError(
            'There were too many files per speaker to process based on your OS settings.  Please try to split your data into more speakers.'))


def apply_cmvn_func(directory, job_name, config):
    normed_scp_path = os.path.join(directory, config.raw_feature_id + '.{}.scp'.format(job_name))
    normed_ark_path = os.path.join(directory, config.raw_feature_id + '.{}.ark'.format(job_name))
    with open(os.path.join(directory, 'log', 'norm.{}.log'.format(job_name)), 'w') as logf:
        utt2spkpath = os.path.join(directory, 'utt2spk.{}'.format(job_name))
        cmvnpath = os.path.join(directory, 'cmvn.{}.scp'.format(job_name))
        featspath = os.path.join(directory, 'feats.{}.scp'.format(job_name))
        if not os.path.exists(normed_scp_path):
            cmvn_proc = subprocess.Popen([thirdparty_binary('apply-cmvn'),
                                          '--utt2spk=ark:' + utt2spkpath,
                                          'scp:' + cmvnpath,
                                          'scp:' + featspath,
                                          'ark,scp:{},{}'.format(normed_ark_path, normed_scp_path)],
                                         stderr=logf
                                         )
            cmvn_proc.communicate()


def apply_cmvn(directory, num_jobs, config):
    child_env = os.environ.copy()
    jobs = [(directory, x, config)
            for x in range(num_jobs)]
    with mp.Pool(processes=num_jobs, initializer=init, initargs=(child_env,)) as pool:
        results = [pool.apply_async(apply_cmvn_func, args=i) for i in jobs]
        output = [p.get() for p in results]


def add_deltas_func(directory, job_name, config):
    normed_scp_path = os.path.join(directory, config.raw_feature_id + '.{}.scp'.format(job_name))
    ark_path = os.path.join(directory, config.feature_id + '.{}.ark'.format(job_name))
    scp_path = os.path.join(directory, config.feature_id + '.{}.scp'.format(job_name))
    with open(os.path.join(directory, 'log', 'add_deltas.{}.log'.format(job_name)), 'w') as logf:
        if config.fmllr_path is not None and os.path.exists(config.fmllr_path):
            deltas_proc = subprocess.Popen([thirdparty_binary('add-deltas'),
                                            'scp:' + normed_scp_path, 'ark:-'],
                                           stderr=logf,
                                           stdout=subprocess.PIPE)
            trans_proc = subprocess.Popen([thirdparty_binary('transform-feats'),
                                           'ark:' + config.fmllr_path, 'ark:-',
                                           'ark,scp:{},{}'.format(ark_path, scp_path)],
                                          stdin=deltas_proc.stdout,
                                          stderr=logf)
            trans_proc.communicate()
        else:
            deltas_proc = subprocess.Popen([thirdparty_binary('add-deltas'),
                                            'scp:' + normed_scp_path, 'ark,scp:{},{}'.format(ark_path, scp_path)],
                                           stderr=logf)
            deltas_proc.communicate()


def add_deltas(directory, num_jobs, config):
    child_env = os.environ.copy()
    jobs = [(directory, x, config)
            for x in range(num_jobs)]
    with mp.Pool(processes=num_jobs, initializer=init, initargs=(child_env,)) as pool:
        results = [pool.apply_async(add_deltas_func, args=i) for i in jobs]
        output = [p.get() for p in results]


def apply_lda_func(directory, job_name, config):
    normed_scp_path = os.path.join(directory, config.raw_feature_id + '.{}.scp'.format(job_name))
    ark_path = os.path.join(directory, config.feature_id + '.{}.ark'.format(job_name))
    scp_path = os.path.join(directory, config.feature_id + '.{}.scp'.format(job_name))
    ivector_scp_path = os.path.join(directory, 'ivector.{}.scp'.format(job_name))
    with open(os.path.join(directory, 'log', 'lda.{}.log'.format(job_name)), 'a') as logf:
        if os.path.exists(config.lda_path):
            splice_feats_proc = subprocess.Popen([thirdparty_binary('splice-feats'),
                                                  '--left-context={}'.format(config.splice_left_context),
                                                  '--right-context={}'.format(config.splice_right_context),
                                                  'scp:' + normed_scp_path,
                                                  'ark:-'],
                                                 stdout=subprocess.PIPE,
                                                 stderr=logf)
            if config.ivectors and os.path.exists(ivector_scp_path):
                transform_feats_proc = subprocess.Popen([thirdparty_binary("transform-feats"),
                                                         config.lda_path,
                                                         'ark:-',
                                                         'ark:-'],
                                                        stdin=splice_feats_proc.stdout,
                                                        stdout=subprocess.PIPE,
                                                        stderr=logf)
                paste_proc = subprocess.Popen([thirdparty_binary('paste-feats'),
                                               'ark:-',
                                               'scp:' + ivector_scp_path,
                                               'ark,scp:{},{}'.format(ark_path, scp_path)],
                                              stdin=transform_feats_proc.stdout,
                                              stderr=logf)
                paste_proc.communicate()
            else:
                transform_feats_proc = subprocess.Popen([thirdparty_binary("transform-feats"),
                                                         config.lda_path,
                                                         'ark:-',
                                                         'ark,scp:{},{}'.format(ark_path, scp_path)],
                                                        stdin=splice_feats_proc.stdout,
                                                        stderr=logf)
                transform_feats_proc.communicate()
        else:
            logf.write('could not find "{}"\n'.format(config.lda_path))
            splice_feats_proc = subprocess.Popen([thirdparty_binary('splice-feats'),
                                                  '--left-context={}'.format(config.splice_left_context),
                                                  '--right-context={}'.format(config.splice_right_context),
                                                  'scp:' + normed_scp_path,
                                                  'ark,scp:{},{}'.format(ark_path, scp_path)],
                                                 stderr=logf)
            splice_feats_proc.communicate()


def apply_lda(directory, num_jobs, config):
    jobs = [(directory, x, config)
            for x in range(num_jobs)]
    with mp.Pool(processes=num_jobs, initializer=init, initargs=(os.environ.copy(),)) as pool:
        results = [pool.apply_async(apply_lda_func, args=i) for i in jobs]
        output = [p.get() for p in results]

[(11, 'os.path.join', 'os.path.join', (['directory', '"""log"""'], {}), False, 'import os\n'), (71, 'os.environ.copy', 'os.environ.copy', ([], {}), False, 'import os\n'), (114, 'os.environ.copy', 'os.environ.copy', ([], {}), False, 'import os\n'), (146, 'os.environ.copy', 'os.environ.copy', ([], {}), False, 'import os\n'), (19, 'os.path.exists', 'os.path.exists', (['segment_path'], {}), False, 'import os\n'), (73, 'os.path.join', 'os.path.join', (['mfcc_directory', '"""log"""'], {}), False, 'import os\n'), (79, 'multiprocessing.Pool', 'mp.Pool', ([], {'processes': 'num_jobs', 'initializer': 'init', 'initargs': '(child_env,)'}), True, 'import multiprocessing as mp\n'), (117, 'multiprocessing.Pool', 'mp.Pool', ([], {'processes': 'num_jobs', 'initializer': 'init', 'initargs': '(child_env,)'}), True, 'import multiprocessing as mp\n'), (149, 'multiprocessing.Pool', 'mp.Pool', ([], {'processes': 'num_jobs', 'initializer': 'init', 'initargs': '(child_env,)'}), True, 'import multiprocessing as mp\n'), (160, 'os.path.exists', 'os.path.exists', (['config.lda_path'], {}), False, 'import os\n'), (102, 'os.path.exists', 'os.path.exists', (['normed_scp_path'], {}), False, 'import os\n'), (127, 'os.path.exists', 'os.path.exists', (['config.fmllr_path'], {}), False, 'import os\n'), (168, 'os.path.exists', 'os.path.exists', (['ivector_scp_path'], {}), False, 'import os\n'), (205, 'os.environ.copy', 'os.environ.copy', ([], {}), False, 'import os\n')]