#!/usr/bin/env bash

# It may be necessary to start the ray server manually.

set -e

for cfg_name in {9,21}-player; do
  cd repo/
    # Run training
    rm -rf log_dir
    python ../helper.py $cfg_name

    # Move generic log dir to properly named one (the code does not support
    # naming the log dir natively).
    log_dir=log_dir-$cfg_name
    rm -rf $log_dir
    mv log_dir $log_dir

    # Find the last checkpoint
    ray_dir=$log_dir/ray_results/APPO/*/
    checkpoint=$(ls $ray_dir | grep '^checkpoint' | sort -t_ -nk2 | tail -n1)
    if [[ -z "$checkpoint" ]]; then
      echo Could not find checkpoint in $ray_dir.
      exit 1
    fi
    ckpt_iter=$(echo $checkpoint | cut -d_ -f2)
    
    # Run the corpus generation code.  If a second argument is given to
    # helper, assume that we are doing corpus genernation.
    python ../helper.py $cfg_name $ray_dir/$checkpoint/checkpoint-$ckpt_iter
  cd ..

  # Make data dir
  data_dir=../data/$cfg_name
  mkdir -p $data_dir
  out_metadata=$data_dir/metadata.json

  # Copy corpora into data dir
  cp repo/log_dir/eval/*.jsonl $data_dir

  # Extract metrics from output
  system_metrics=$(
    tail -n100 repo/$ray_dir/result.json |\
      jq "
        select(.iterations_since_restore == $ckpt_iter).custom_metrics |
          pick(.accord_mean,.win_vil_mean)
      "
  )

  if [[ -z "$system_metrics" ]]; then
    echo Failed to extract system metrics
    exit 1
  fi

  # Update corpus metadata
  # TODO Put this in common file
  if ! [[ -s $out_metadata ]]; then
    echo '{}' >$out_metadata
  fi
  jq ".metrics.system=$system_metrics" \
    <$out_metadata >$out_metadata.tmp
  mv $out_metadata{.tmp,}
done
