#!/bin/bash
# ----------------------------------------------------------
# Copyright (c) Microsoft Corporation.  All rights reserved.
# Licensed under the MIT license. See LICENSE.md file in the project root for full license information.
# ----------------------------------------------------------
# Helper script containing common code used by run-test scripts of E2E tests

BinaryPath=$TEST_CNTK_BINARY

export V2_LIB_TESTING=1

if [ "$TEST_DEVICE" == "cpu" ]; then
  CNTKDeviceId=-1
elif [ "$TEST_DEVICE" == "gpu" ]; then
  CNTKDeviceId=0
else
  echo "Error: Unknown TEST_DEVICE specified!"
  exit 3
fi

LogFileName=

ConfigDir=$TEST_DIR
RunDir=$TEST_RUN_DIR
DataDir=$TEST_DATA_DIR
OutputDir=$TEST_RUN_DIR

MPIMode=0
MPIArgs=

DeleteExistingModels=1
DeleteModelsAfterTest=1

# Print info needed by MetricsDriver.py to extract tests metrics
printHardwareInfo()
{
  cpuName=$(cat /proc/cpuinfo 2> /dev/null | grep -m 1 'model name' | cut -d : -f 2- | tr -s " " | cut -c 2-)
  totalMemory=$(cat /proc/meminfo 2> /dev/null | grep 'MemTotal' | cut -d : -f 2- | tr -s " " | cut -c 2-)
  nproc=$(nproc)

  # Note that MetricsDriver.py depends on this format
  echo "CPU info:"
  echo "    CPU Model Name: $cpuName"
  echo "    Hardware threads: $nproc"
  echo "    Total Memory: $totalMemory"
  echo "-------------------------------------------------------------------"
}

# Helper function to print and run a command
run()
{
  cmd=$1
  shift
  if [ "$DRY_RUN" == "1" ]; then
    workingDir=$PWD
    if [ "$OS" == "Windows_NT" ]; then
      workingDir=$(cygpath -aw $workingDir)
      if [[ $MPIMode == 0 ]]; then
        cmd=$(cygpath -aw $cmd)
        TEST_ROOT_DIR_ESCAPED=`echo -n $(cygpath -aw $TEST_ROOT_DIR) | sed 's/\\\\/\\\\\\\\/g'`
        workingDir=`echo "$workingDir" | sed "s/$TEST_ROOT_DIR_ESCAPED/\\$\\(SolutionDir\\)\\\\\\\\Tests\\\\\\\\EndToEndTests/g"`
      fi
    fi
    echo Working Directory: $workingDir
    echo Full command: "$cmd" "$@"
    if [ "$OS" == "Windows_NT" ]; then
      if [[ $MPIMode == 0 ]]; then
        echo VS debugging command args: "$@" | sed "s/$TEST_ROOT_DIR_ESCAPED/\\$\\(SolutionDir\\)\\\\Tests\\\\EndToEndTests/g"
      fi
    fi
    return 1
  else
    echo === Running "$cmd" "$@"
    "$cmd" "$@"
    return $?
  fi
}

# Function for launching the CNTK executable
# cntkrun <CNTK config file name> <additional CNTK args>
cntkrun()
{
  configFileName=$1
  additionalCNTKArgs=$2

  if [ "$OS" == "Windows_NT" ]; then
    # When running on cygwin translating /cygdrive/xxx paths to proper windows paths:
    ConfigDir=$(cygpath -aw $ConfigDir)
    RunDir=$(cygpath -aw $RunDir)
    DataDir=$(cygpath -aw $DataDir)
    OutputDir=$(cygpath -aw $OutputDir)
  fi

  CNTKArgs="configFile=$ConfigDir/$configFileName currentDirectory=$DataDir RunDir=$RunDir DataDir=$DataDir ConfigDir=$ConfigDir OutputDir=$OutputDir DeviceId=$CNTKDeviceId timestamping=true $additionalCNTKArgs"
  if [ "$LogFileName" != "" ]; then
    CNTKArgs="$CNTKArgs stderr=$RunDir/$LogFileName"
  fi

  modelsDir=$TEST_RUN_DIR/Models
  if [[ $DeleteExistingModels == 1 ]]; then
    [ -d $modelsDir ] && rm -rf "$modelsDir"
  fi
  mkdir -p $modelsDir || exit $?

  if [[ $MPIMode == 0 ]]; then
    run "$BinaryPath" $CNTKArgs
  else
    run "$MPI_BINARY" $MPIArgs $BinaryPath $CNTKArgs
  fi
  local ExitCode=$?
  if [[ $DeleteModelsAfterTest == 1 ]]; then
    [ -d $modelsDir ] && rm -rf "$modelsDir"
  fi

  return $ExitCode
}

# Given number of instances, return number of hardware threads we can use per
# instance
threadsPerInstance()
{
  local threads=$((`nproc` / $1))
  [[ $threads -eq 0 ]] && echo 1 || echo $threads
}

# Function for launching a parallel CNTK run with MPI
# cntkmpirun <MPI args> <CNTK config file name> <additional CNTK args>
cntkmpirun()
{
  # Since we use the MS MPI program on Windows, the CNTK binary path argument
  # passed to mpiexec must be in the windows format
  if [ "$OS" == "Windows_NT" ]; then
    BinaryPath=$(cygpath -aw $BinaryPath)
  fi

  MPIMode=1
  MPIArgs=$1

  cntkrun "$2" "$3"
  return $?
}

#Function to check/copy original baselines
# checkOriginalBaseline <original baseline relative dir> <target dir>
checkOriginalBaseline()
{
  TargetPath=$1
  OriginRelativePath=$2
  
  (cd $TargetPath/$OriginRelativePath && md5sum baseline*) | (cd $TargetPath && md5sum --status -c -)
  if [ $? != 0 ]; then
    echo Error: Baselines must match original test. Copy from $OriginalTestDir.
    exit 1
    #cp $TargetPath/$OriginRelativePath/baseline* $TargetPath
  fi
}

# place printHardwareInfo here, so that all tests print it
printHardwareInfo