#!/bin/bash
export CUBLAS_WORKSPACE_CONFIG=":16:8" # https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility
export PYTHONHASHSEED=0
export base_dir="/system/user/publicwork/hauzenbe/eva_cache"
export WANDB_PROJECT="EVA"
export WANDB_ENTITY="ml_eva"

n_proc=$1
model_card='gpt2.md'
task_name='e2e_nlg'
# array params
seeds=(0 10 101) # 0 10 101
lora_dims=(16 8 4 2) # 16 8 4 2
learning_rates=(6e-2 6e-3 6e-4) # 6e-2 6e-3 6e-4
exp_names=('pca_init_adaptive') # 'rand_init' 'pca_init' 'pca_init_adaptive'

# create chunks from array params
chunks=()
chunk=()
for a in ${seeds[@]}
do
  for b in ${lora_dims[@]}
  do
    for c in ${learning_rates[@]}
    do
      for d in ${exp_names[@]}
      do
        chunk+=("$a|$b|$c|$d")
        # add chunk to chunks once its equal to n_parallel
        if [[ ${#chunk[@]} -eq $n_proc ]]; then
          chunk_string=$( IFS="%"; echo "${chunk[*]}")
          chunks+=($chunk_string)
          chunk=()
        fi
      done
    done
  done
done

# Add the final chunk (if not empty)
if [[ ${#chunk[@]} -gt 0 ]]; then 
  chunk_string=$( IFS="%"; echo "${chunk[*]}")
  chunks+=($chunk_string)
fi

echo "${chunks[@]}"

# start scripts in parallel
for chunk in ${chunks[@]}
do
  parallel_combs=($(echo $chunk | tr "%" " "))
  for (( i=0; i<${#parallel_combs[@]}; i++ ))
  do
    params_array=($(echo ${parallel_combs[i]} | tr "|" " "))
    seed=${params_array[0]}
    lora_dim=${params_array[1]}
    lr=${params_array[2]}
    experiment_name=${params_array[3]}

    dir="r_${lora_dim}_lr_${lr}"
    output_dir=$base_dir/$model_card/$task_name/$experiment_name/$dir/$seed

    python src/gpt2_decode.py \
    --vocab ./vocab \
    --sample_file $output_dir/beam_prediction.jsonl \
    --input_file ./data/e2e/test_formatted.jsonl \
    --output_ref_file $output_dir/e2e_ref.txt \
    --output_pred_file $output_dir/e2e_pred.txt &
  done
  wait
done