#!/usr/bin/env bash
set -e

########################################
# 기본 설정 (환경에 맞게 수정)
########################################
CMD="python train.py"
CFG_STAGE1="/workspace/stormer/configs/pretrain_one_step.yaml"
CFG_STAGE2="/workspace/stormer/configs/finetune_multi_step_stage2.yaml"
CFG_STAGE3="/workspace/stormer/configs/finetune_multi_step_stage3.yaml"

DATA_ROOT="/workspace/stormer/wb2_processed"

########################################
# (hidden_size, depth, num_heads) 세트 정의
########################################
SETTINGS=(
  "384-18-6"
)

########################################
# sweep loop
########################################
for setting in "${SETTINGS[@]}"; do
  # "hidden-depth-head" 파싱
  IFS="-" read -r h d nh <<< "${setting}"

  RUN_ID="h${h}_d${d}_n${nh}"
  echo "============================================================"
  echo "RUN: ${RUN_ID}  (hidden=${h}, depth=${d}, num_heads=${nh})"
  echo "============================================================"

  ######################################
  # STAGE 1 : PRETRAIN (one step)
  ######################################
  # ${CMD} \
    # --config "${CFG_STAGE1}" \
    # --trainer.default_root_dir "/workspace/stormer/results/pretrain/${RUN_ID}" \
    # --trainer.logger.init_args.name "pretrain_one_step_${RUN_ID}" \
    # --model.net.init_args.hidden_size "${h}" \
    # --model.net.init_args.depth "${d}" \
    # --model.net.init_args.num_heads "${nh}" \
    # --model.net.patch_size 4 \
    # --data.root_dir "${DATA_ROOT}" \
    # --data.steps 1 \
    # --data.batch_size 1

  # yaml에서 dirpath: "${trainer.default_root_dir}/pretrain_one_step/checkpoints"
  PRETRAIN_CKPT="/workspace/stormer/results/pretrain/${RUN_ID}/pretrain_one_step_${RUN_ID}/checkpoints/last.ckpt"

  ######################################
  # STAGE 2 : FINETUNE (4 steps)
  ######################################
  ${CMD} \
    --config "${CFG_STAGE2}" \
    --trainer.default_root_dir "/workspace/stormer/results/finetune/4steps_${RUN_ID}" \
    --trainer.logger.init_args.name "finetune_4_steps_${RUN_ID}" \
    --model.net.init_args.hidden_size "${h}" \
    --model.net.init_args.depth "${d}" \
    --model.net.init_args.num_heads "${nh}" \
    --model.net.patch_size 4 \
    --model.pretrained_path "${PRETRAIN_CKPT}" \
    --data.root_dir "${DATA_ROOT}" \
    --data.steps 4 \
    --data.batch_size 1

  FINETUNE4_CKPT="/workspace/stormer/results/finetune/4steps_${RUN_ID}/finetune_4_steps_${RUN_ID}/checkpoints/last.ckpt"

  ######################################
  # STAGE 3 : FINETUNE (8 steps)
  ######################################
  ${CMD} \
    --config "${CFG_STAGE3}" \
    --trainer.default_root_dir "/workspace/stormer/results/finetune/8steps_${RUN_ID}" \
    --trainer.logger.init_args.name "finetune_8_steps_${RUN_ID}" \
    --model.net.init_args.hidden_size "${h}" \
    --model.net.init_args.depth "${d}" \
    --model.net.init_args.num_heads "${nh}" \
    --model.net.patch_size 4 \
    --model.pretrained_path "${FINETUNE4_CKPT}" \
    --data.root_dir "${DATA_ROOT}" \
    --data.steps 8 \
    --data.batch_size 1

done