#!/bin/bash

#
# Copyright (C) 2017 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

NNAPI_VERSION="
V1_0
V1_1
"

# Process one test spec, and optionally provide the log file argument
# for the slicing tool. The first argument is the test spec file; the
# second optional argument specifies the log file this test should dump
# results into. Only used by the test slicing tool to collect reference
# outputs from the CPU. Also, it outputs the right #includes in the
# test harness so the test would be invoked by TestGenerated.cpp
#
# This function shouldn't be directly called from other scripts. Use
# generate_wrapper below for generating models and examples and updating the
# test framework in one shot.

export NNAPI_BASE=$ANDROID_BUILD_TOP/frameworks/ml/nn
: ${TEST_DIR:=frameworks/ml/nn/runtime/test}

function generate_one_testcase {
  # Generate one testcase
  local LOGFILE=$2
  if [ -n "$2" ]; then
    local LOGFILE=", \"$2\""
  fi
  local BASENAME=`basename -s .mod.py $1`
  local EXAMPLE="-e $ANDROID_BUILD_TOP/$TEST_DIR/generated/examples/$BASENAME.example.cpp"

  $NNAPI_BASE/tools/test_generator/test_generator.py ./`basename $1`\
    -m $ANDROID_BUILD_TOP/$TEST_DIR/generated/models/$BASENAME.model.cpp $EXAMPLE
  ret=$?
  # Paste these lines into TestGenerated.cpp
  echo
  echo namespace $BASENAME {
  echo std::vector\<MixedTypedExample\> examples \= {
  echo // Generated $BASENAME test
  echo \#include \"generated/examples/$BASENAME.example.cpp\"
  echo }\;
  echo // Generated model constructor
  echo \#include \"generated/models/$BASENAME.model.cpp\"
  echo }  // namespace $BASENAME
  echo TEST_F\(GeneratedTests\, $BASENAME\) {
  echo '    execute'\($BASENAME\:\:CreateModel\,
  echo '            '$BASENAME\:\:is_ignored\,
  echo '            '$BASENAME\:\:examples${LOGFILE}\)\;
  echo }
  return $ret
}

# Driver for generate_one_testcase. Append the output of generate_one_testcase
# (which are C++ snippets that invokes the test harness) to the
# all_generated_tests.cpp
# Optionally, the "LOG" file ($2), only used by the slicing tool, would be
# passed to generate_one_testcase.
#
# This function should be called to process one test spec from other scripts.
function generate_wrapper {
  local LOGFILE=""
  if [ $1 = "log" ]; then
    local LOGFILE=$2
    shift
    shift
  fi
  cd $ANDROID_BUILD_TOP/$TEST_DIR/specs
  OUTFILE=$ANDROID_BUILD_TOP/$TEST_DIR/generated/all_generated_tests.cpp
  echo "// DO NOT EDIT;" > $OUTFILE
  echo "// Generated by ml/nn/runtime/test/specs/generate_test.sh" >> $OUTFILE
  FOUND=0

  for ver in $NNAPI_VERSION;
  do
    VER_DIR=$ANDROID_BUILD_TOP/$TEST_DIR/specs/$ver
    [ ! -d $VER_DIR ] && continue
    pushd $VER_DIR > /dev/null
    for f in $@;
    do
      if [ -f $(basename $f) ]; then
        generate_one_testcase $f "$LOGFILE" >> $OUTFILE
        if [ $? -ne 0 ]; then
          echo "Failed processing $f"
          return $?
        fi
        FOUND=1
      fi
    done
    popd > /dev/null
  done
  if [[ $FOUND -eq 0 ]]; then
    echo did not find any files for $@
    exit 1
  fi
  return $?
}

# Only run the following when not sourced by another script
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
  set -eu
  if [ $# -eq 0 ]; then
    FILES=${ANDROID_BUILD_TOP}/${TEST_DIR}/specs/V*/*.mod.py
  else
    FILES="$@"
  fi
  generate_wrapper $FILES
  if [ $? -ne 0 ]; then
    exit $?
  fi
  echo "Generated file in ml/nn/runtime/test/generated/"`basename $OUTFILE`
fi # [[ "${BASH_SOURCE[0]}" == "${0}" ]]