#!/bin/bash

NAME=$(grep '^name:' config.yml | awk '{print$2}' | tr -d "'")
MODEL=$(grep '^model:' config.yml | awk '{print$2}' | tr -d "'")
OUT_DIR=$(grep '^out_dir:' config.yml | awk '{print$2}' | tr -d "'")

cp -r prompts prompts_temp
sed -i "s#{BASE_DIR}#$OUT_DIR#g" prompts_temp/*.txt

cp root_to_numpy.smk root_to_numpy_temp.smk
cp scores.smk scores_temp.smk
cp categorization.smk categorization_temp.smk
sed -i "s#{BASE_DIR}#$OUT_DIR#g" *_temp.smk

mkdir -p $OUT_DIR/generated_code
cp utils.py $OUT_DIR/generated_code/utils.py

rm -f $OUT_DIR/logs/success.npy
rm -f $OUT_DIR/logs/calls.npy
rm -f $OUT_DIR/logs/input_tokens.npy
rm -f $OUT_DIR/logs/output_tokens.npy

echo "Starting step 1: ROOT -> NumPy"
TIME1=$SECONDS
snakemake -s root_to_numpy_temp.smk -j 1 --forcerun summarize_root --rerun-incomplete --configfile config.yml --latency-wait 120 --verbose
TIME1=$((SECONDS-TIME1))

echo "Starting step 2: SB scores"
TIME2=$SECONDS
snakemake -s scores_temp.smk -j 1 --forcerun scores --rerun-incomplete --configfile config.yml --latency-wait 120 --verbose
TIME2=$((SECONDS-TIME2))

echo "Starting step 3: categorization"
TIME3=$SECONDS
snakemake -s categorization_temp.smk -j 1 --forcerun categorization --rerun-incomplete --configfile config.yml --latency-wait 120 --verbose
TIME3=$((SECONDS-TIME3))

echo "Checking results"
python check_soln.py --out_dir $OUT_DIR

echo "Writing stats"

read -r -a success_arr < <(python get_arr.py --name success --out_dir $OUT_DIR)
SUCCESS1=${success_arr[0]}
SUCCESS2=${success_arr[1]}
SUCCESS3=${success_arr[2]}

read -r -a call_arr < <(python get_arr.py --name calls --out_dir $OUT_DIR)
CALLS1=${call_arr[0]}
CALLS2=${call_arr[1]}
CALLS3=${call_arr[2]}

read -r -a input_token_arr < <(python get_arr.py --name input_tokens --out_dir $OUT_DIR)
INPUT_TOKENS1=${input_token_arr[0]}
INPUT_TOKENS2=${input_token_arr[1]}
INPUT_TOKENS3=${input_token_arr[2]}

read -r -a output_token_arr < <(python get_arr.py --name output_tokens --out_dir $OUT_DIR)
OUTPUT_TOKENS1=${output_token_arr[0]}
OUTPUT_TOKENS2=${output_token_arr[1]}
OUTPUT_TOKENS3=${output_token_arr[2]}

python update_stats.py --name $NAME --success1 $SUCCESS1 --time1 $TIME1 --calls1 $CALLS1 --input_tokens1 $INPUT_TOKENS1 --success2 $SUCCESS2 --time2 $TIME2 --calls2 $CALLS2 --input_tokens2 $INPUT_TOKENS2 --success3 $SUCCESS3 --time3 $TIME3 --calls3 $CALLS3 --input_tokens3 $INPUT_TOKENS3 --output_tokens1 $OUTPUT_TOKENS1 --output_tokens2 $OUTPUT_TOKENS2 --output_tokens3 $OUTPUT_TOKENS3

rm -r prompts_temp
rm root_to_numpy_temp.smk
rm scores_temp.smk
rm categorization_temp.smk

echo "Finished!"