#!/usr/bin/env bash

# Move up one level from the bash_scripts folder so we're in the main folder 
# where "python -m train" typically works.
cd "$(dirname "$0")"/..

# Use BASE_PATH from environment or set default if not defined
BASE_PATH=${BASE_PATH:-"./"}
# Remove trailing slash if present
BASE_PATH=${BASE_PATH%/}

# Set how many jobs to run in parallel.
MAX_CONCURRENT=4

for YEAR in {1990..2013}; do
  # Train: [YEAR/01/01 ... (YEAR+7)/12/31]
  TRAIN_BEG=$(printf "%04d0101" "${YEAR}")
  TRAIN_END=$(printf "%04d1231" $((YEAR + 7)))
  
  # Val:   [(YEAR+8)/01/01 ... (YEAR+8)/12/31]
  VAL_BEG=$(printf "%04d0101" $((YEAR + 8)))
  VAL_END=$(printf "%04d1231" $((YEAR + 8)))
  VAL_YEAR=$(printf "%04d" $((YEAR + 8)))
  # Test:  [(YEAR+9)/01/01 ... (YEAR+9)/12/31]
  TEST_BEG=$(printf "%04d0101" $((YEAR + 9)))
  TEST_END=$(printf "%04d1231" $((YEAR + 9 )))
  
  
  echo "==========================================="
  echo " Sliding Window Training:"
  echo "   Train: ${TRAIN_BEG} to ${TRAIN_END}"
  echo "   Val:   ${VAL_BEG}   to ${VAL_END}"
  echo "   Test:  ${TEST_BEG}  to ${TEST_END}"
  echo "==========================================="
  
  # Concurrency check: if we already have $MAX_CONCURRENT jobs running, wait
  while [ "$(jobs -p | wc -l)" -ge "$MAX_CONCURRENT" ]; do
    sleep 2
  done
  
  # Fire off the training in the background
  python -m train experiment=equities/equities_hyena_joint.yaml \
    dataset.dataset_config.train_beg_date="${TRAIN_BEG}" \
    dataset.dataset_config.train_end_date="${TRAIN_END}" \
    dataset.dataset_config.val_beg_date="${VAL_BEG}" \
    dataset.dataset_config.val_end_date="${VAL_END}" \
    dataset.dataset_config.test_beg_date="${TEST_BEG}" \
    dataset.dataset_config.test_end_date="${TEST_END}" \
    ++dataset.save_data=False \
    ++dataset.load_data=True \
    ++dataset.data_path="${BASE_PATH}/data/equities/equity_dataset_"${VAL_YEAR}".npz" \
    wandb.name="${VAL_BEG}_hyena" &
  
  # Sleep a few seconds to ensure no two Hydra runs start in the same second
  sleep 5

done

# Wait for all background jobs to finish
wait

echo "All sliding window training jobs have completed."

# chmod +x rolling_train_equities_hyena_joint.sh
# ./rolling_train_equities_hyena_joint.sh