'''@file train_asr.py
this file will do the asr training'''

import os
from functools import partial
import tensorflow as tf
from six.moves import configparser
from nabu.distributed import create_server
from nabu.processing import batchdispenser, feature_reader, target_coder
from nabu.neuralnetworks.classifiers.asr import asr_factory
from nabu.neuralnetworks.trainers import trainer_factory
from nabu.neuralnetworks.decoders import decoder_factory

def train_asr(clusterfile,
              job_name,
              task_index,
              ssh_command,
              expdir):

    ''' does everything for asr training
    Args:
        clusterfile: the file where all the machines in the cluster are
            specified if None, local training will be done
        job_name: one of ps or worker in the case of distributed training
        task_index: the task index in this job
        ssh_command: the command to use for ssh, if 'None' no tunnel will be
            created
        expdir: the experiments directory
    '''

    #read the database config file
    parsed_database_cfg = configparser.ConfigParser()
    parsed_database_cfg.read(os.path.join(expdir, 'database.cfg'))
    database_cfg = dict(parsed_database_cfg.items('database'))

    #read the features config file
    parsed_feat_cfg = configparser.ConfigParser()
    parsed_feat_cfg.read(os.path.join(expdir, 'model', 'features.cfg'))
    feat_cfg = dict(parsed_feat_cfg.items('features'))

    #read the asr config file
    parsed_nnet_cfg = configparser.ConfigParser()
    parsed_nnet_cfg.read(os.path.join(expdir, 'model', 'asr.cfg'))
    nnet_cfg = dict(parsed_nnet_cfg.items('asr'))

    #read the trainer config file
    parsed_trainer_cfg = configparser.ConfigParser()
    parsed_trainer_cfg.read(os.path.join(expdir, 'trainer.cfg'))
    trainer_cfg = dict(parsed_trainer_cfg.items('trainer'))

    #read the decoder config file
    parsed_decoder_cfg = configparser.ConfigParser()
    parsed_decoder_cfg.read(os.path.join(expdir, 'model', 'decoder.cfg'))
    decoder_cfg = dict(parsed_decoder_cfg.items('decoder'))

    #make distinction between three implemented different kind of training forms
    if database_cfg['train_mode'] == 'supervised':
        nonsupervised = False
    elif database_cfg['train_mode'] == 'nonsupervised' or\
            database_cfg['train_mode'] == 'semisupervised':
        nonsupervised = True
    else:
        raise Exception('Wrong kind of training mode')

    #when (partly) nonsupervised, what features are used for the reconstruction
    #currently two possible options implemented
    if nonsupervised:
        if trainer_cfg['reconstruction_features'] == 'audio_samples':
            audio_used = True
        elif trainer_cfg['reconstruction_features'] == 'input_features':
            audio_used = False
        else:
            raise Exception(
                'Unknown specification for the reconstruction features')

    #read the quant config file if nonsupervised training and samples used
    if nonsupervised:
        if audio_used:
            parsed_quant_cfg = configparser.ConfigParser()
            parsed_quant_cfg.read(os.path.join(expdir,
                                               'model', 'quantization.cfg'))
            quant_cfg = dict(parsed_quant_cfg.items('features'))

    #based on the other settings, compute and overwrite samples_per_hlfeature
    #and unpredictable_samples in the classifier config dictionary
    if nonsupervised:
        if audio_used:
            rate_after_quant = int(quant_cfg['quant_rate'])
            win_lenght = float(feat_cfg['winlen'])
            win_shift = float(feat_cfg['winstep'])
            samples_one_window = int(win_lenght*rate_after_quant)
            samples_one_shift = int(win_shift*rate_after_quant)
            #### THIS IS ONLY RELEVANT WHEN USING A LISTENER WITH PYRAM STRUCT
            # and this line should be adapted otherwise
            time_compression = 2**int(nnet_cfg['listener_numlayers'])
            #store values in config dictionary
            nnet_cfg['samples_per_hlfeature'] = samples_one_shift\
                                        *time_compression
            nnet_cfg['unpredictable_samples'] = (samples_one_window+\
                                    (time_compression-1)\
                        *samples_one_shift)-nnet_cfg['samples_per_hlfeature']


    #create the cluster and server
    server = create_server.create_server(
        clusterfile=clusterfile,
        job_name=job_name,
        task_index=task_index,
        expdir=expdir,
        ssh_command=ssh_command)

    #the ps should just wait
    if job_name == 'ps':
        server.join()

    # path to where the training samples are stored
    featdir = os.path.join(database_cfg['train_dir'], feat_cfg['name'])

    #create the coder
    with open(os.path.join(database_cfg['train_dir'], 'alphabet')) as fid:
        alphabet = fid.read().split(' ')
    coder = target_coder.TargetCoder(alphabet)

    #create a feature reader for the training data
    with open(featdir + '/maxlength', 'r') as fid:
        max_length = int(fid.read())

    featreader = feature_reader.FeatureReader(
        scpfile=featdir + '/feats_shuffled.scp',
        cmvnfile=featdir + '/cmvn.scp',
        utt2spkfile=featdir + '/utt2spk',
        max_length=max_length)

    #read the feature dimension
    with open(featdir + '/dim', 'r') as fid:
        input_dim = int(fid.read())

    #the path to the text file
    textfile = os.path.join(database_cfg['train_dir'], 'targets')


    # If nonsupervised and audio used, we also need to read samples
    # these can be done with a second feature reader
    if nonsupervised:
        if audio_used:
            featdir2 = os.path.join(database_cfg['train_dir'],
                                    quant_cfg['name'])

            with open(featdir2 + '/maxlength', 'r') as fid:
                max_length_audio = int(fid.read())

            audioreader = feature_reader.FeatureReader(
                scpfile=featdir2 + '/feats_shuffled.scp',
                cmvnfile=None,
                utt2spkfile=None,
                max_length=max_length_audio)

    ## create a batch dispenser, depending on which situation we're in
    if not nonsupervised:
    # in the normal supervised training mode, regular dispenser is needed
        if 'las_ignoring_mode' in trainer_cfg:
            if trainer_cfg['las_ignoring_mode'] == 'True':
            # if we ignore unlabeled examples
                dispenser = batchdispenser.AsrTextBatchDispenser(
                    feature_reader=featreader,
                    target_coder=coder,
                    size=int(trainer_cfg['batch_size']),
                    target_path=textfile)
            elif trainer_cfg['las_ignoring_mode'] == 'False':
            # if we choose to process the unlabeled examples
                if 'fixed_ratio' in trainer_cfg:
                    if trainer_cfg['fixed_ratio'] == 'True':
                    # if we choose to process with batches with fixed
                    # labeled/unlabeled ratio
                        dispenser = \
                            batchdispenser.AsrTextBatchDispenserAltFixRatio(
                                feature_reader=featreader,
                                target_coder=coder,
                                size=int(trainer_cfg['batch_size']),
                                target_path=textfile,
                                percentage_unlabeled=1-float(
                                    database_cfg['part_labeled']))
                    elif trainer_cfg['fixed_ratio'] == 'False':
                        # if the fixed ratio is not used
                        dispenser = batchdispenser.AsrTextBatchDispenserAlt(
                            feature_reader=featreader,
                            target_coder=coder,
                            size=int(trainer_cfg['batch_size']),
                            target_path=textfile)
                    else:
                        raise Exception('wrong information in fixed_ratio var')
                else:
                # if fixed ratio is not specified, we choose to do without it
                    dispenser = batchdispenser.AsrTextBatchDispenserAlt(
                        feature_reader=featreader,
                        target_coder=coder,
                        size=int(trainer_cfg['batch_size']),
                        target_path=textfile)
            else:
                raise Exception('wrong information in LAS_ignoring_mode var')
        else:
        # if no specification is made about the ignoring, ignore the unlabeled
            dispenser = batchdispenser.AsrTextBatchDispenser(
                feature_reader=featreader,
                target_coder=coder,
                size=int(trainer_cfg['batch_size']),
                target_path=textfile)
    else:
        # when doing (partly) nonsupervised extra reconstruction features needed
        if audio_used:
        # when the audio is the reconstruction feature
            if 'fixed_ratio' in trainer_cfg:
                if trainer_cfg['fixed_ratio'] == 'True':
                    # if specified to work with fixed lab/unlab ratio batches
                    dispenser = \
                        batchdispenser.AsrTextAndAudioBatchDispenserFixRatio(
                            feature_reader=featreader,
                            audio_reader=audioreader,
                            target_coder=coder,
                            size=int(trainer_cfg['batch_size']),
                            target_path=textfile,
                            percentage_unlabeled=1-float(
                                database_cfg['part_labeled']))
                elif trainer_cfg['fixed_ratio'] == 'False':
                # if specified to not use the fixed ratio
                    dispenser = batchdispenser.AsrTextAndAudioBatchDispenser(
                        feature_reader=featreader,
                        audio_reader=audioreader,
                        target_coder=coder,
                        size=int(trainer_cfg['batch_size']),
                        target_path=textfile)
                else:
                    raise Exception('wrong information in fixed_ratio var')
            else:
            # without specification, suppose no fixed ratio batches
                dispenser = batchdispenser.AsrTextAndAudioBatchDispenser(
                    feature_reader=featreader,
                    audio_reader=audioreader,
                    target_coder=coder,
                    size=int(trainer_cfg['batch_size']),
                    target_path=textfile)
        else:
        # if no audio is used, the input features are used
            if 'fixed_ratio' in trainer_cfg:
                if trainer_cfg['fixed_ratio'] == 'True':
                # if specified to work with fixed labeled/unlabled ratio batches
                    dispenser = \
                        batchdispenser.AsrTextAndFeatBatchDispenserFixRatio(
                            feature_reader=featreader,
                            target_coder=coder,
                            size=int(trainer_cfg['batch_size']),
                            target_path=textfile,
                            percentage_unlabeled=1-float(
                                database_cfg['part_labeled']))
                elif trainer_cfg['fixed_ratio'] == 'False':
                # if specified to not use the fixed ratio
                    dispenser = batchdispenser.AsrTextAndFeatBatchDispenser(
                        feature_reader=featreader,
                        target_coder=coder,