#!/bin/bash

# This script automates the running of PISCES logic locally using Argus sim and
# git checkouts.
set -euo pipefail

# Initialize a timer.
SECONDS=0

CMD=$(basename $0)
function usage {
  echo "============================"
  echo "Error: $1"
  echo ""
  echo "----------------------------"
  echo "Simulate two SHAs and compare results with PISCES"
  echo ""
  echo "Comparisons will be rendered via local Argus server, assumed to be"
  echo "running and hosted on port 8080. If no local Argus server exists,"
  echo "please follow the instructions on this confluence page:"
  echo "https://confluence.zooxlabs.com/display/PLAN/Argus+Sim+Viewer."
  echo ""
  echo "Alternatively, ping the #argus-sim-help channel on Slack."
  echo ""
  echo "Usage:"
  echo "      $CMD control_sha candidate_sha"
  echo "      $CMD control_sha candidate_sha -c chum_uri"
  echo "      $CMD control_sha candidate_sha [-i] -s [15min_28cover|30min_38cover|40min_48cover] {default: 15min_28cover}"
  echo ""
  echo "NOTE: 15min_28cover - takes about 15 mins to run and cover 28% of planner behaviors (full pisces cover 56%)"
  echo "      30min_38cover - takes about 30 mins to run and cover 38% of planner behaviors (full pisces cover 56%)"
  echo "      40min_48cover - takes about 40 mins to run and cover 48% of planner behaviors (full pisces cover 56%)"
  exit 1
}

# ========================================================================================= #
# NOTE: Changing argument list would break the build plan "PISCES_PLANS",                   #
#       please also update 'ci/cinder/plans/pisces.py' if the public API is to be changed.  #
# ========================================================================================= #

if (($# == 0)); then
  usage "No arguments provided."
fi

if (($# < 2)); then
  usage "At least 2 arguments are required."
fi

if [[ $(git diff --stat) != '' ]]; then
  echo 'Local changes in your repository. Please stash or commit them.'
  echo
  exit 1
fi

# Required argments
CONTROL=$1
CANDIDATE=$2
SCENARIO_SET="15min_28cover"     # Default to the first scenario set
CI=false                         # Default to not running CI

if (($# > 2)); then
  shift 2 # shift the argument position, to process for the optional arguments
  while getopts ":c:s:i" opt; do
    case ${opt} in
      c )
        CHUM_URI=$OPTARG;;
      s )
        SCENARIO_SET=$OPTARG;;
      i )
        CI=true;;
      \? )
        usage "Invalid option: -$OPTARG";;
      : )
        usage "Invalid option: -$OPTARG requires an argument.";;
    esac
  done
fi

# cd to the top-level git repo
pushd $(git rev-parse --show-toplevel)

CURRENT_BRANCH="$(git rev-parse --abbrev-ref HEAD)"
if [[ "${CURRENT_BRANCH}" == "HEAD" ]]; then
  CURRENT_BRANCH="$(git rev-parse HEAD)"
fi

# Defer cleanup to any EXIT and restore repo to whatever state it was before we
# ran this script
function restore_repo {
  # Restore to your original SHA
  git checkout "${CURRENT_BRANCH}"

  # restore us to the original directory
  popd
}
trap restore_repo EXIT

function get_chum_uris {
  # set the input variable of reference type, so can modify it
  local -n chum_uri_arr=$1
  TMP_PATH=$(mktemp -dt)/chums.pbtxt
  cp "/mnt/nautilus_rw/dzeng/pisces/local_test_set_${SCENARIO_SET}.pbtxt" $TMP_PATH

  for t in $(cat $TMP_PATH); do
    if [[ "$t" =~ chum://.* ]]; then
      chum=`echo $t | sed 's/.*\"\(.*\)\".*$/\1/g'`
      chum_uri_arr+=($chum)
    fi
  done
  rm $TMP_PATH
}

# Check out a SHA and run a log test
function simulate {
  set -euo pipefail
  SHA=$1
  chum_uri=$2
  git checkout $1 > /dev/null

  CHUM_OUTPUT=$(./sim/argus_sim/run.sh "${chum_uri}" \
    in_process_planner:=true \
    override_autonomous_mode:=PURE_PLAYBACK_MODE \
    rviz:=false \
    chum_output_path:=/tmp/chum \
    --config=cuda \
    --config=kraken | \
    tee /dev/stderr |
    grep -A1 "Simulated Chum URI:" | \
    tail -n1)
  echo $CHUM_OUTPUT
}

#########################################
## Stage-1: Re-simulate Selected Chums ##
#########################################

# Initialize an array to define the chum uris to run, if the user didn't specify
# a CHUM_URI, then run it with number of predefined scenarios; otherwise, run ONLY
# the chum scenario defined by the user.
chum_uris_to_run=()
if [ -z ${CHUM_URI+x} ]; then
  # Get chum uris from a set defined inside of the get function
  get_chum_uris chum_uris_to_run
else
  chum_uris_to_run+=($CHUM_URI)
fi
# Run the simulation sequentially and populate the arrays below with resulting chum uris.
CONTROL_RESULTS=()
echo "====== Running CONTROL sims ==========="
for uri in ${chum_uris_to_run[@]}; do
  echo "====== Running sim $uri"
  CONTROL_RESULTS+=($(simulate $CONTROL $uri))
  echo "====== Done running control sim $uri"
done
CANDIDATE_RESULTS=()
echo "====== Running CANDIDATE sims ==========="
for uri in ${chum_uris_to_run[@]}; do
  echo "====== Running sim $uri"
  CANDIDATE_RESULTS+=($(simulate $CANDIDATE $uri))
  echo "====== Done running candidate sim $uri"
done

# run PISCES extraction and analysis on current branch
restore_repo

# remove EXIT trap as we are in the current repo state
trap - EXIT

#####################################
## Stage-2: Extract Simulated Data ##
#####################################

DATA_DIR=$(mktemp -p /tmp -d -t pisces-XXXXXXXXXX)
mkdir -p $DATA_DIR/baseline
mkdir -p $DATA_DIR/comparison

BAZEL_RUN='bazel --nohome_rc run --config=kraken'

for idx in ${!chum_uris_to_run[@]}; do
  # Remove the prefix of `chum://` to make scenario name
  SCENARIO_NAME=${chum_uris_to_run[$idx]##*chum://}
  echo "Extracting data to $DATA_DIR/*/$SCENARIO_NAME*"
  $BAZEL_RUN //mined_metric/builder/metrics_impl/pisces:extract_data -- \
    --chum_uri ${CONTROL_RESULTS[$idx]} \
    --filepath $DATA_DIR/baseline/${SCENARIO_NAME}.csv

  $BAZEL_RUN //mined_metric/builder/metrics_impl/pisces:extract_data -- \
    --chum_uri ${CANDIDATE_RESULTS[$idx]} \
    --filepath $DATA_DIR/comparison/${SCENARIO_NAME}.csv
  echo "finished running data extraction to $DATA_DIR/*/$SCENARIO_NAME*"
done

#######################################
## Stage-3: Aggregate Simulated Data ##
#######################################

EXPERIMENT_ID=$(basename $DATA_DIR)
ARGS="""
  --directory /tmp \
  --id $EXPERIMENT_ID \
  --control baseline \
  --control-sha $CONTROL \
  --candidate comparison \
  --candidate-sha $CANDIDATE \
  --argus-host $(hostname -I | cut -d' ' -f1):8080
"""
if ! $CI; then
# if this is being launched on ci using github "ci-pisces-xxxx" label
# then it will not have the right user email. So don't do notification.
  ARGS=$ARGS" --email $(git config user.email)"
else
# if it is launched from CI, then make sure to raise exception when changes
# are detected, this will block merging of the behavioral changing PRs
  ARGS=$ARGS" --raise-on-changed"
fi
echo "Analyze and aggregate extracted data:"
$BAZEL_RUN //mined_metric/builder/metrics_impl/pisces/utils:analyze_data -- $ARGS

echo "========================================================================"
echo "Result posted to http://pisces.web.zooxlabs.com/$EXPERIMENT_ID"
if $CI; then
  echo ""
  echo """NOTE: argus links in the pisces report is NOT working when launched from ci,
      but results are valid. If further debugging needed for a specific scenario,
      regenerate debug report on your dev machine, using the following command:
      ./mined_metric/builder/metrics_impl/pisces/run_local.sh control_sha candidate_sha -c {chum with delta}"""
fi
echo "========================================================================="
echo "Total run time $SECONDS seconds."
