#!/bin/bash

#####################################################
# File :    FirstPass.sh                            #
# Brief :   ASR first pass and speaker diarization  #
# Author :  Jean-François Rey                       #
#	        (base on Emmanuel Ferreira              #
#	        and Hugo Mauchrétien works)             #
# Version : 1.1                                     #
# Date :    18/06/13                                #
#####################################################

# Check OTMEDIA_HOME env var
if [ -z ${OTMEDIA_HOME} ]
then 
    OTMEDIA_HOME=$(dirname $(dirname $(readlink -e $0)))
    export OTMEDIA_HOME=$OTMEDIA_HOME
fi

# where is FirstPass.sh
MAIN_SCRIPT_PATH=$(dirname $(readlink -e $0))

# scripts path
SCRIPT_PATH=$OTMEDIA_HOME/tools/scripts

# Include scripts
. $SCRIPT_PATH"/Tools.sh"
. $SCRIPT_PATH"/CheckFirstPass.sh"

# where is FirstPass.cfg
FIRSTPASS_CONFIG_FILE=$OTMEDIA_HOME"/cfg/FirstPass.cfg"
if [ -e $FIRSTPASS_CONFIG_FILE ]
then
	. $FIRSTPASS_CONFIG_FILE
else
	echo "ERROR : Can't find configuration file $FIRSTPASS_CONFIG_FILE" >&2
	exit 1
fi

#---------------#
# Parse Options #
#---------------#
while getopts ":hDv:cf:r" opt
do
	case $opt in
		h)
			echo -e "$0 [OPTIONS] <WAV_FILE> <OUTPUT_DIRECTORY>\n"
            echo -e "\t Options:"
            echo -e "\t\t-h :\tprint this message"
            echo -e "\t\t-D :\tDEBUG mode on"
            echo -e "\t\t-v l :\tVerbose mode, l=(1|2|3) level mode"
            echo -e "\t\t-c :\tCheck process, stop if error detected"
            echo -e "\t\t-f n :\tspecify a speeral forks number (default 1)"
            echo -e "\t\t-r :\tforce rerun the wav file"
			exit 1
			;;
		D)
			DEBUG=1
			;;
        v)
            VERBOSE=$OPTARG
            ;;
        c)
            CHECK=1
            ;;
        f)
            FORKS="--forks $OPTARG"
            ;;
        r)
            RERUN=1
            ;;
		:)
			echo "Option -$OPTARG requires an argument." >&2
			exit 1
			;;
		\?)
			echo "BAD USAGE : unknow opton -$OPTARG"
			exit 1
			;;
	esac
done

# mode debug enable
if [ $DEBUG -eq 1 ]
then
       set -x
       echo -e "## Mode DEBUG ON ##"
       REDIRECTION_OUTPUT=""
else
       REDIRECTION_OUTPUT=" > /dev/null 2>&1"
fi

# mode verbose enable
if [ $VERBOSE -gt 0 ]; then echo -e "## Verbose level : $VERBOSE ##" ; REDIRECTION_OUTPUT=" 2> /dev/null"; fi

# Check USAGE by arguments number
if [ $(($#-($OPTIND-1))) -ne 2 ]
then
	echo "BAD USAGE : FirstPass.sh [OPTIONS] <WAV_FILE> <OUTPUT_DIR>"
	echo "$0 -h for more info"
	exit 1
fi

shift $((OPTIND-1))
# check audio file - First argument
if [ -e $1 ] && [ -s $1 ]
then
	# absolute path to wav file
	WAV_FILE=$(readlink -e $1)
	# wav filename
	FILENAME=$(basename $WAV_FILE)
	# wav filename without extension
	BASENAME=${FILENAME%.*}

    print_info "=> $BASENAME P1 | $(date +'%d/%m/%y %H:%M:%S')" 1
    print_info "$WAV_FILE OK" 1
else
	print_error "can't find $1 OR file is empty"
	exit 1
fi

# check output directory - Second argument
if [ ! -e $2 ]
then
   	mkdir -p $2
    print_info "Make directory $2" 1
fi


#-------------#
# GLOBAL VARS #
#-------------#
OUTPUT_DIR=$(readlink -e $2)			# Output directory absolute path
OUTPUT_DIR_BASENAME="$OUTPUT_DIR/$BASENAME/" 	# New OUTPUT_DIR with BASENAME
PLP_FILE="$OUTPUT_DIR_BASENAME/$BASENAME.plp"   # Global PLP file
PLP_DIR="$OUTPUT_DIR_BASENAME/PLP/"		# Segmented PLP files directory
SEG_FILE="$OUTPUT_DIR_BASENAME/$BASENAME.seg"	# Global Seg file
LBL_FILE="$OUTPUT_DIR_BASENAME/$BASENAME.lbl"	# Global LBL file
RES_DIR=$OUTPUT_DIR_BASENAME"/res_p1"
LOGFILE="$OUTPUT_DIR/info_p1.log"
ERRORFILE="$OUTPUT_DIR/error_p1.log"

#------------------#
# Create WORKSPACE #
#------------------#
if [ ! -e $OUTPUT_DIR_BASENAME ]
then
	mkdir -p $OUTPUT_DIR_BASENAME
    print_info "Make directory $OUTPUT_DIR_BASENAME" 1
fi

# Lock directory
if [ -e $OUTPUT_DIR_BASENAME/FIRSTPASS.lock ] && [ $RERUN -eq 0 ]; then exit 1; fi
rm "$OUTPUT_DIR_BASENAME/FIRSTPASS.unlock" > /dev/null 2>&1
touch "$OUTPUT_DIR_BASENAME/FIRSTPASS.lock" > /dev/null 2>&1

rm -r $PLP_DIR > /dev/null 2>&1; 
mkdir -p $PLP_DIR
print_info "Make directory $PLP_DIR" 1
if [ $RERUN -eq 0 ];
then
    rm -r $RES_DIR > /dev/null 2>&1;
else
    rm $RES_DIR/*.lock > /dev/null 2>&1
fi
mkdir -p $RES_DIR $REDIRECTION_OUTPUT
print_info "Make directory $RES_DIR" 1

#--------------------#
# Save configuration #
#--------------------#
cp $FIRSTPASS_CONFIG_FILE $OUTPUT_DIR_BASENAME/FirstPass.cfg
echo "FIRSTPASS_SCRIPT_PATH=$MAIN_SCRIPT_PATH" >> $OUTPUT_DIR_BASENAME/FirstPass.cfg
echo "WAV_FILE=$WAV_FILE" >> $OUTPUT_DIR_BASENAME/FirstPass.cfg
echo "BASENAME=$BASENAME" >> $OUTPUT_DIR_BASENAME/FirstPass.cfg
echo "OUTPUT_DIR=$OUTPUT_DIR" >> $OUTPUT_DIR_BASENAME/FirstPass.cfg
echo "OUTPUT_DIR_BASENAME=$OUTPUT_DIR_BASENAME" >> $OUTPUT_DIR_BASENAME/FirstPass.cfg
echo "PLP_FILE=$PLP_FILE" >> $OUTPUT_DIR_BASENAME/FirstPass.cfg
echo "PLP_DIR=$PLP_DIR" >> $OUTPUT_DIR_BASENAME/FirstPass.cfg
echo "SEG_FILE=$SEG_FILE" >> $OUTPUT_DIR_BASENAME/FirstPass.cfg
echo "LBL_FILE=$LBL_FILE" >> $OUTPUT_DIR_BASENAME/FirstPass.cfg
echo "RES_DIR=$RES_DIR" >> $OUTPUT_DIR_BASENAME/FirstPass.cfg
print_info "save config in $OUTPUT_DIR_BASENAME/FirstPass.cfg" 1

#-------------------------#
# Check Audio File Format #
#-------------------------#
error=0
temp=$(avconv -i $WAV_FILE 2>&1 | grep "16000 Hz")
if [ -z "$temp" ]; then error=1; fi
temp=$(avconv -i $WAV_FILE 2>&1 | grep "1 channels")
if [ -z "$temp" ]; then error=1; fi
temp=$(avconv -i $WAV_FILE 2>&1 | grep "s16")
if [ -z "$temp" ]; then error=1; fi

if [ $error -eq 1 ]
then
    print_message $WARNING 2 "$WAV_FILE is not a wav file at 16000 Hz, 1 channel, 16bits\nhave to convert"
    print_message $INFO 3 "avconv -i $WAV_FILE -threads 4 -vn -f wav -ac 1 -ar 16000 -ab 256000 $OUTPUT_DIR_BASENAME/$BASENAME.wav"
	avconv -i $WAV_FILE -threads 4 -vn -f wav -ac 1 -ar 16000 -ab 256000 $OUTPUT_DIR_BASENAME/$BASENAME.wav $REDIRECTION_OUTPUT
	WAV_FILE=$OUTPUT_DIR_BASENAME/$BASENAME.wav
	FILENAME=$BASENAME.wav
    print_message $INFO 1 "new wav file : $WAV_FILE"
fi

#---------------#
# Get SRT file  #
#---------------#
if [ -s $(dirname $WAV_FILE)/$BASENAME.SRT ]
then
    cp $(dirname $WAV_FILE)/$BASENAME.SRT $OUTPUT_DIR_BASENAME/$BASENAME.SRT
    print_info "copy $BASENAME.SRT file into workingspace" 1
fi

#------------#
# WAV -> PLP #
#------------#
print_info "convert WAV -> PLP" 1
echo $FILENAME > $OUTPUT_DIR_BASENAME/list.tmp
print_info "$BIN_PATH/lia_plp_mt.32 --lst $OUTPUT_DIR_BASENAME/list.tmp --input_dir $(dirname $WAV_FILE) --output_dir $OUTPUT_DIR_BASENAME --input_type WAV --output_type HTK --nb_coef 12 --cms
" 2

$BIN_PATH/lia_plp_mt$ARCH --lst $OUTPUT_DIR_BASENAME/list.tmp --input_dir $(dirname $WAV_FILE) --output_dir $OUTPUT_DIR_BASENAME --input_type WAV --output_type HTK --nb_coef 12 --cms $REDIRECTION_OUTPUT

if [ $CHECK -eq 1 ]
then
    check_first_pass_plp "$PLP_FILE"
    if [ $? -eq 1 ]
    then
        echo "ERROR [$(date +'%d/%m/%y %H:%M:%S')] $PLP_FILE" >> $ERRORFILE
        exit 1
    fi
fi

rm $OUTPUT_DIR_BASENAME/list.tmp

#------------------------------#
# S/NS + SPEAKERS SEGMENTATION #
#------------------------------#
print_info "Launch speakers diarization" 1
# Calcul seg file
print_info "java -Xmx4096m -jar $BIN_PATH/LIUM_SpkDiarization-4.2.jar --fInputMask=${WAV_FILE} --sOutputMask=${SEG_FILE} $BASENAME" 2
#java -Xmx8000m -Xms2048 -jar $BIN_PATH/LIUM_SpkDiarization-4.2.jar --fInputMask=${WAV_FILE} --sOutputMask=${SEG_FILE} $BASENAME
java -Xmx4096m -jar $BIN_PATH/LIUM_SpkDiarization-4.2.jar --fInputMask=${WAV_FILE} --sOutputMask=${SEG_FILE} $BASENAME $REDIRECTION_OUTPUT #–doCEClustering 

if [ $CHECK -eq 1 ] && ( [ ! -e $SEG_FILE ] || [ -z $SEG_FILE ] )
then
    echo "ERROR [$(date +'%d/%m/%y %H:%M:%S')] $SEG_FILE" >> $ERRORFILE
    exit 1
fi


# Create LBL file
print_info "Extract LBL file from SEG file" 2

cat $SEG_FILE | grep -v ";;" | cut -f3,4,5,8 -d" " | tr " " "#" | sort -k1 -n | tr "#" " " > $LBL_FILE

if [ $CHECK -eq 1 ] && ( [ ! -e $LBL_FILE ] || [ -z $LBL_FILE ] )
then
    echo "ERROR [$(date +'%d/%m/%y %H:%M:%S')] $LBL_FILE" >> $ERRORFILE
    exit 1
fi


#----------------------------------------------------#
# Cut global PLP file depending to LBL segmentations #
#----------------------------------------------------#
print_info "Cut PLP file depending to LBL segmentations" 1
print_info "$BIN_PATH/gcep $PLP_FILE $LBL_FILE 500 $PLP_DIR -FSEG" 2

$SPEERAL_TOOLS/gcep $PLP_FILE $LBL_FILE 500 $PLP_DIR -FSEG $REDIRECTION_OUTPUT

if [ $CHECK -eq 1 ]
then
    check_first_pass_plps_lbl $PLP_DIR $LBL_FILE
    if [ $? -eq 1 ]
    then
        echo "ERROR [$(date +'%d/%m/%y %H:%M:%S')] $PLP wrong .plp files number" >> $ERRORFILE
        exit 1
    fi
fi

# change plp files names
cd $PLP_DIR;
rename -f s/_/#/g *plp
rename -f s/#/_/ *plp
cd $OLDPWD

#---------------------------------------------#
# PLP files list depending to acoustic models #
#---------------------------------------------#
print_info "Create PLP list depending of the model" 1
# Create a list of plp files
find $PLP_DIR -type f -exec basename "{}" .plp \; | sort > $OUTPUT_DIR_BASENAME/plp.lst

rm $OUTPUT_DIR_BASENAME/plp_*.lst > /dev/null 2>&1
for (( i=0; $i<${#MTAG[@]} ; i++ ))
do
	a=`grep -e "${MTAG[$i]}" $OUTPUT_DIR_BASENAME/plp.lst`
    if [ -n "$a" ]; then
        print_info "$OUTPUT_DIR_BASENAME/plp_${MODS[$i]}.lst" 3
       	grep -e "${MTAG[$i]}" $OUTPUT_DIR_BASENAME/plp.lst | sort > $OUTPUT_DIR_BASENAME/plp_${MODS[$i]}.lst
    fi
done

#-----------------------#
# First Pass (DECODING) #
#-----------------------#
#
# For all AM do decoding
# if Check error -> iter on undone decoding (max 1 times)
#
print_info "Launch decoding" 1
for (( i=0; $i<${#MTAG[@]} ; i++ ))
do
    redo=1; # nb of try if not all segs is decoded
	if [ -e $OUTPUT_DIR_BASENAME/plp_${MODS[$i]}.lst ]
	then
        todo=$OUTPUT_DIR_BASENAME/plp_${MODS[$i]}.lst
        while [ $redo -gt 0 ]; do
            rm $RES_DIR/*.lock > /dev/null 2>&1
            print_info "$SPEERAL_BIN $todo $RES_DIR ${SPEERAL_CFG[$i]} -r $PLP_DIR -m $SPEERAL_AM/${MODS[$i]}.hmm -c $SPEERAL_AM/${MODS[$i]}.cls $FORKS --lock $REDIRECTION_OUTPUT" 2
            # Run speeral
            $SPEERAL_BIN  ${todo} $RES_DIR ${SPEERAL_CFG[$i]} -r $PLP_DIR -m $SPEERAL_AM/${MODS[$i]}.hmm -c $SPEERAL_AM/${MODS[$i]}.cls $FORKS --lock $REDIRECTION_OUTPUT

            # Check if error
            if [ $CHECK -eq 1 ]
            then
                check_first_pass_output_speeral "${OUTPUT_DIR_BASENAME}/plp_${MODS[$i]}.lst" "$RES_DIR"
                # if error
                if [ $? -eq 1 ]
                then
                    # rerun 
                    redo=$(($redo - 1));
                    echo -e "WARN : Speeral output ERROR ${OUTPUT_DIR_BASENAME}/plp_${MODS[$i]}.lst" >> $ERRORFILE
                    # new plp list
                    # list .seg done and compare to list of seg to do
                    ls $RES_DIR/*.seg | grep -e "${MTAG[$i]}" | sed -e "s|$RES_DIR\/||g" | sed -e 's/\.seg//' | sort > ${OUTPUT_DIR_BASENAME}/.tmp
                    diff ${OUTPUT_DIR_BASENAME}/plp_${MODS[$i]}.lst ${OUTPUT_DIR_BASENAME}/.tmp | grep -e "^< " | sed -e "s/< //" > ${OUTPUT_DIR_BASENAME}/todo.lst
                    rm ${OUTPUT_DIR_BASENAME}/.tmp
                    # log seg to do
                    cat ${OUTPUT_DIR_BASENAME}/todo.lst >> $ERRORFILE
                    todo=${OUTPUT_DIR_BASENAME}/todo.lst
                    echo -e "WARN : Try $redo" >> $ERRORFILE
                fi
            else
                redo=-5;
            fi
        done
        if [ $redo -eq 0 ]
        then
            echo -e "ERROR : Speeral $todo" >> $ERRORFILE
            cat ${OUTPUT_DIR_BASENAME}/todo.lst >> $ERRORFILE
            #exit 1
        fi
        rm ${OUTPUT_DIR_BASENAME}/todo.lst > /dev/null 2>&1
	    #rm $OUTPUT_DIR_BASENAME/plp_${MODS[$i]}.lst
        rm $RES_DIR/*.lock > /dev/null 2>&1
    fi
done

print_info "<= End P1 $BASENAME | $(date +'%d/%m/%y %H:%M:%S')" 1

## Check missing seg and log it
ls $RES_DIR/*.seg | sed -e "s|$RES_DIR\/||g" | sed -e 's/\.seg//' | sort > ${OUTPUT_DIR_BASENAME}/.tmp
echo -e "$BASENAME P1 END\n[" >> $LOGFILE
diff ${OUTPUT_DIR_BASENAME}/plp.lst ${OUTPUT_DIR_BASENAME}/.tmp | grep -e "^< " | sed -e "s/< //" >> $LOGFILE
todo=$(cat ${OUTPUT_DIR_BASENAME}/plp.lst | wc -l)
notdone=$(($todo - $(cat ${OUTPUT_DIR_BASENAME}/.tmp | wc -l)))
pourcentage=$((($notdone*100)/$todo))
echo -e "] $pourcentage%  $BASENAME" >> $LOGFILE
rm ${OUTPUT_DIR_BASENAME}/.tmp


#---------------#
# Convert res   #
#---------------#

# .res => .ctm
$SCRIPT_PATH/res2out.pl --dir $RES_DIR --format CTM --ignore $RULES/asupp --out $OUTPUT_DIR_BASENAME/$BASENAME.1pass.ctm $REDIRECTION_OUTPUT
# .res => .trs
echo -e "name $AUTHOR\nfileName $BASENAME\nfileExt wav\nsegFile $OUTPUT_DIR_BASENAME/$BASENAME.seg" > $OUTPUT_DIR_BASENAME/$BASENAME.trs_cfg
$SCRIPT_PATH/res2out.pl --dir $RES_DIR --format TRS --ignore $RULES/asupp --out $OUTPUT_DIR_BASENAME/$BASENAME.1pass.trs --trs_config $OUTPUT_DIR_BASENAME/$BASENAME.trs_cfg $REDIRECTION_OUTPUT
rm $OUTPUT_DIR_BASENAME/$BASENAME.trs_cfg 2> /dev/null
# .res => .txt
$SCRIPT_PATH/res2out.pl --dir $RES_DIR --format TXT --ignore $RULES/asupp --out $OUTPUT_DIR_BASENAME/$BASENAME.1pass.txt $REDIRECTION_OUTPUT

# unlock directory
mv "$OUTPUT_DIR_BASENAME/FIRSTPASS.lock" "$OUTPUT_DIR_BASENAME/FIRSTPASS.unlock"