#!/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>
  "
              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." >&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=" 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
  have 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"
      echo "[${BASENAME}] ERROR : check $ERRORFILE file" > /dev/stderr

      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"
      echo "[${BASENAME}] ERROR : check $ERRORFILE file" > /dev/stderr

      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"
          echo "[${BASENAME}] ERROR : check $ERRORFILE file" > /dev/stderr

          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/< //" >> $ERROFILE
          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
  fileName $BASENAME
  fileExt wav
  segFile $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"