#!/usr/bin/env bash
# rolling_train_equities_mha.sh
# Sliding-window training for the Multi-Head-Attention baseline
# – 5 random seeds per year, up to 5 jobs running at once.

cd "$(dirname "$0")"/..   # repo root

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

# ── Tunables ────────────────────────────────────────────────────────────────────
N_SEEDS=5          # launch seeds 1…5 for every rolling window
MAX_CONCURRENT=4   # cap on simultaneous Hydra jobs

# ── Outer loop over rolling years ───────────────────────────────────────────────
for YEAR in {1990..2013}; do
  TRAIN_BEG=$(printf "%04d0101" "${YEAR}")
  TRAIN_END=$(printf "%04d1231" $((YEAR + 7)))
  VAL_BEG=$(printf  "%04d0101" $((YEAR + 8)))
  VAL_END=$(printf  "%04d1231" $((YEAR + 8)))
  TEST_BEG=$(printf "%04d0101" $((YEAR + 9)))
  TEST_END=$(printf "%04d1231" $((YEAR + 9)))
  VAL_YEAR=$(printf "%04d"     $((YEAR + 8)))

  echo "────────────────────────────────────────────────────────────"
  echo "YEAR ${YEAR}: Train ${TRAIN_BEG}–${TRAIN_END},   Val ${VAL_BEG}–${VAL_END},   Test ${TEST_BEG}–${TEST_END}"
  echo "────────────────────────────────────────────────────────────"

  # ── Inner loop over seeds ─────────────────────────────────────────────────────
  for SEED in $(seq 1 "$N_SEEDS"); do
    # Throttle until running jobs < MAX_CONCURRENT
    while [ "$(jobs -p | wc -l)" -ge "$MAX_CONCURRENT" ]; do
      sleep 2
    done

    python -m train experiment=equities/equities_mha.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" \
      ++train.seed="${SEED}" \
      wandb.name="${VAL_BEG}_mha_seed${SEED}" &

    # Small pause to avoid Hydra timestamp collisions
    sleep 5
  done
done

# ── Final rendez-vous ──────────────────────────────────────────────────────────
wait
echo "All sliding-window training jobs have completed."

# chmod +x rolling_train_equities_mha.sh
# ./rolling_train_equities_mha.sh