#!/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" > /dev/stderr
	echo "exit" > /dev/stderr
	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, and log it into files, can 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." > /dev/stderr
			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=" 2> /dev/null"
fi

# mode verbose enable
if [ $VERBOSE -gt 0 ]; then echo -e "## Verbose level : $VERBOSE ##" ;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 start | $(date +'%d/%m/%y %H:%M:%S')" 1
    print_info "[${BASENAME}] $WAV_FILE OK" 2 
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 "[${BASENAME}] Make directory $2" 2 
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_BASENAME/info_p1.log"
ERRORFILE="$OUTPUT_DIR_BASENAME/error_p1.log"

#------------------#
# Create WORKSPACE #
#------------------#
if [ ! -e $OUTPUT_DIR_BASENAME ]
then
	mkdir -p $OUTPUT_DIR_BASENAME
    print_info "[${BASENAME}] Make directory $OUTPUT_DIR_BASENAME" 2
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 "[${BASENAME}] Make directory $PLP_DIR" 2
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 > /dev/null 2>&1
print_info "[${BASENAME}] Make directory $RES_DIR" 2
rm $LOGFILE $ERRORFILE > /dev/null 2>&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 "[${BASENAME}] 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 "[${BASENAME}] $WAV_FILE is not a wav file at 16000 Hz, 1 channel, 16bits\nhave to convert"
    print_message $INFO 3 "[${BASENAME}] 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 "[${BASENAME}] 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 "[${BASENAME}] copy $BASENAME.SRT file into $OUTPUT_DIR_BASENAME" 3
fi

#------------#
# WAV -> PLP #
#------------#
print_info "[${BASENAME}] convert WAV -> PLP" 1
echo $FILENAME > $OUTPUT_DIR_BASENAME/list.tmp
print_info "[${BASENAME}] $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" 3 

$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
        print_log_file "$ERROFILE" "ERROR [$(date +'%d/%m/%y %H:%M:%S')] Creating PLP file : $PLP_FILE"
        echo "[${BASENAME}] ERROR : ${BASENAME} check $ERRORFILE file" > /dev/stderr
        exit 1
    fi
fi

rm $OUTPUT_DIR_BASENAME/list.tmp

#------------------------------#
# S/NS + SPEAKERS SEGMENTATION #
#------------------------------#
print_info "[${BASENAME}] Launch speakers diarization" 1
# Calcul seg file
print_info "[${BASENAME}] java -Xmx4096m -jar $BIN_PATH/LIUM_SpkDiarization-4.2.jar --fInputMask=${WAV_FILE} --sOutputMask=${SEG_FILE} $BASENAME" 3
#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
    print_log_file "$ERRORFILE" "ERROR [$(date +'%d/%m/%y %H:%M:%S')] Creating SEG file : $SEG_FILE"
    print_error "[${BASENAME}] ERROR : check $ERRORFILE file"
    exit 1
fi


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

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
    print_log_file "$ERRORFILE" "ERROR [$(date +'%d/%m/%y %H:%M:%S')] Creating LBL file : $LBL_FILE"
    print_error "[${BASENAME}] ERROR : check $ERRORFILE file"
    exit 1
fi


#----------------------------------------------------#
# Cut global PLP file depending to LBL segmentations #
#----------------------------------------------------#
print_info "[${BASENAME}] Cut PLP file depending to LBL segmentations" 1
print_info "[${BASENAME}] $BIN_PATH/gcep $PLP_FILE $LBL_FILE 500 $PLP_DIR -FSEG" 3

$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
        print_log_file $ERRORFILE "ERROR [$(date +'%d/%m/%y %H:%M:%S')] $PLP wrong number of .plp files"
        print_error "[${BASENAME}] ERROR : check $ERRORFILE file" 
        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 "[${BASENAME}] 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 "[${BASENAME}] Creating $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 "[${BASENAME}] 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 "[${BASENAME}] $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" 3
            # 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));
                    print_warn "[${BASENAME}] Speeral output error : check $LOGFILE" 2
                    print_log_file $LOGFILE "WARN : Speeral number of output ERROR ${OUTPUT_DIR_BASENAME}/plp_${MODS[$i]}.lst"
                    # 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
                    print_log_file $LOGFILE "Segs not done :"
                    cat ${OUTPUT_DIR_BASENAME}/todo.lst >> $LOGFILE
                    todo=${OUTPUT_DIR_BASENAME}/todo.lst
                    print_warn "[${BASENAME}] Try $redo" 3
                fi
            fi
        done
        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

## Check missing seg and log it
if [ $CHECK -eq 1 ]
then
    ls $RES_DIR/*.seg | sed -e "s|$RES_DIR\/||g" | sed -e 's/\.seg//' | sort > ${OUTPUT_DIR_BASENAME}/.tmp
    todo=$(cat ${OUTPUT_DIR_BASENAME}/plp.lst | wc -l)
    if [ $todo -eq 0 ]; then todo=1;fi
    notdone=$(($todo - $(cat ${OUTPUT_DIR_BASENAME}/.tmp | wc -l)))
    pourcentage=$((($notdone*100)/$todo))

    if [ $notdone -ne 0 ]
    then
        print_error "[${BASENAME}] ERROR : check $ERRORFILE"
        print_log_file "$ERRORFILE" "ERROR : Segs not done ["
        diff ${OUTPUT_DIR_BASENAME}/plp.lst ${OUTPUT_DIR_BASENAME}/.tmp | grep -e "^< " | sed -e "s/< //" >> $ERRORFILE
        print_log_file "$ERRORFILE" "] $pourcentage% $BASENAME"
    else
        print_log_file "$LOGFILE" "P1 OK $BASENAME | $(date +'%d/%m/%y %H:%M:%S')"
    fi
    rm ${OUTPUT_DIR_BASENAME}/.tmp
fi

#---------------#
# Convert res   #
#---------------#
print_info "[${BASENAME}] Convert .res into .ctm" 1 
# .res => .ctm
$SCRIPT_PATH/res2out.pl --dir $RES_DIR --format CTM --ignore $RULES/asupp --out $OUTPUT_DIR_BASENAME/$BASENAME.1pass.ctm $REDIRECTION_OUTPUT
print_info "[${BASENAME}] Convert .res into .trs" 1 
# .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
print_info "[${BASENAME}] Convert .res into .txt" 1 
# .res => .txt
$SCRIPT_PATH/res2out.pl --dir $RES_DIR --format TXT --ignore $RULES/asupp --out $OUTPUT_DIR_BASENAME/$BASENAME.1pass.txt $REDIRECTION_OUTPUT

print_info "[${BASENAME}] <= P1 End | $(date +'%d/%m/%y %H:%M:%S')" 1
# unlock directory
mv "$OUTPUT_DIR_BASENAME/FIRSTPASS.lock" "$OUTPUT_DIR_BASENAME/FIRSTPASS.unlock"