#!/bin/bash

hid_dim=1024
hid_num=6
pretrain_dir=exp/dnn4_pretrain-dbn
nerv_kaldi=/speechlab/users/mfy43/nerv/speech/kaldi_io/

[ -f path.sh ] && . ./path.sh
. parse_options.sh || exit 1;

data=$1
data_cv=$2
lang=$3
alidir=$4
alidir_cv=$5
dir=$6

[[ -z $data_fmllr ]] && data_fmllr=data-fmllr-tri3
[[ -z $alidir ]] && alidir=exp/tri3_ali
[[ -z $dir ]] && dir=exp/dnn4_nerv_prepare
[[ -z $data ]] && data=$data_fmllr/train_tr90
[[ -z $data_cv ]] && data_cv=$data_fmllr/train_cv10
kaldi_to_nerv=$nerv_kaldi/tools/kaldi_to_nerv
mkdir $dir -p
mkdir $dir/log -p
###### PREPARE DATASETS ######
cp $data/feats.scp $dir/train_sorted.scp
cp $data_cv/feats.scp $dir/cv.scp
utils/shuffle_list.pl --srand ${seed:-777} <$dir/train_sorted.scp >$dir/train.scp

feats_tr="ark:copy-feats scp:$dir/train.scp ark:- |"

###### INITIALIZE OUTPUT LAYER ######
[ -z $num_tgt ] && \
    num_tgt=$(hmm-info --print-args=false $alidir/final.mdl | grep pdfs | awk '{ print $NF }')
nnet_proto=$dir/nnet_output.proto
echo "# genrating network prototype $nnet_proto"
utils/nnet/make_nnet_proto.py \
    $hid_dim $num_tgt 0 $hid_dim >$nnet_proto || exit 1
nnet_init=$dir/nnet_output.init
nnet-initialize --binary=false $nnet_proto $nnet_init

###### MODEL PARAMETER CONVERSION ######
$kaldi_to_nerv $nnet_init $dir/nnet_output.nerv $hid_num
$kaldi_to_nerv <(nnet-copy --binary=false $pretrain_dir/${hid_num}.dbn -) $dir/nnet_init.nerv
$kaldi_to_nerv <(nnet-copy --binary=false $pretrain_dir/final.feature_transform -) $dir/nnet_trans.nerv
###### PREPARE FOR DECODING #####
echo "Using PDF targets from dirs '$alidir' '$alidir_cv'"
# training targets in posterior format,
labels_tr="ark:ali-to-pdf $alidir/final.mdl \"ark:gunzip -c $alidir/ali.*.gz |\" ark:- | ali-to-post ark:- ark:- |"
labels_cv="ark:ali-to-pdf $alidir/final.mdl \"ark:gunzip -c $alidir_cv/ali.*.gz |\" ark:- | ali-to-post ark:- ark:- |"
# training targets for analyze-counts,
labels_tr_pdf="ark:ali-to-pdf $alidir/final.mdl \"ark:gunzip -c $alidir/ali.*.gz |\" ark:- |"
labels_tr_phn="ark:ali-to-phones --per-frame=true $alidir/final.mdl \"ark:gunzip -c $alidir/ali.*.gz |\" ark:- |"

# get pdf-counts, used later for decoding/aligning,
analyze-counts --verbose=1 --binary=false "$labels_tr_pdf" $dir/ali_train_pdf.counts 2>$dir/log/analyze_counts_pdf.log || exit 1
# copy the old transition model, will be needed by decoder,
copy-transition-model --binary=false $alidir/final.mdl $dir/final.mdl || exit 1
# copy the tree
cp $alidir/tree $dir/tree || exit 1

# make phone counts for analysis,
[ -e $lang/phones.txt ] && analyze-counts --verbose=1 --symbol-table=$lang/phones.txt "$labels_tr_phn" /dev/null 2>$dir/log/analyze_counts_phones.log || exit 1