#!/usr/bin/env bash
set -euo pipefail

# ============================================================================
# [Decay-only] 선택한 Stage1 branch 체크포인트로부터 *decay(annealing)만* 수행.
#              (run_stage2_3_from_ckpt.sh 의 (1) decay 단계만 떼어낸 것. stage2/3 는 하지 않음.)
#
#  - SETTINGS  : run_stage1_sweep.sh 에서 쓴 것과 동일한 (hidden-depth-heads) 세트.
#  - EPOCHS    : 목표 stage1 "총 budget"(decay 포함) epoch 들. SAVE_EVERY 의 배수여야
#               stage1 에서 저장된 branch 체크포인트와 정확히 맞습니다.
#
#  각 (setting, budget) 조합마다:
#     branch ckpt(= interval_epoch_BRANCH_EP.ckpt) 의 가중치에서 시작해
#     LR 을 DECAY_LR -> DECAY_ETA_MIN 으로 DECAY_EPOCHS 만큼 decay 합니다.
#
#  ※ decay 길이는 *비율 기반*(stage1 의 branch 저장과 동일 공식):
#       DECAY_EPOCHS(E) = ceil(max(DECAY_RATIO*E, MIN_DECAY_EPOCHS))
#       BRANCH_EP(E)    = E - DECAY_EPOCHS(E)
#    총 budget E = BRANCH_EP stable + DECAY_EPOCHS decay (길이 E 단일 WSD run 과 등가).
#    ⚠️ DECAY_RATIO / MIN_DECAY_EPOCHS 는 run_stage1_sweep.sh 와 *반드시 동일*해야
#       BRANCH_EP 가 stage1 이 저장한 branch epoch 과 일치합니다.
#
#  ※ 출력 경로/이름은 run_stage2_3_from_ckpt.sh 의 decay 단계와 동일하게 맞춰서,
#    원하면 나중에 그 스크립트로 stage2/3 만 이어서 돌릴 수 있습니다
#    (그 경우 decay 는 last.ckpt 자동 resume 으로 즉시 통과).
# ============================================================================

CMD="python train.py"
CFG_DECAY="/workspace/stormer/configs/decay_finetune.yaml"
DATA_ROOT="/workspace/stormer/wb2_h5df"

S1_RESULTS_ROOT="/workspace/stormer/results/pretrain"
DECAY_RESULTS_ROOT="/workspace/stormer/results/decay"

# ---- stage1 과 동일한 (hidden-depth-heads) 세트 ----
SETTINGS=(
  "128-24-2"
)

# ---- ⭐ 목표 stage1 "총 budget"(decay 포함) epoch 들 (SAVE_EVERY 배수) ----
EPOCHS=(20 40 80 100)

PATCH_SIZE=4
BATCH_SIZE=1

# ---- decay(annealing) 설정 ----
# ⚠️ DECAY_RATIO / MIN_DECAY_EPOCHS 는 run_stage1_sweep.sh 와 동일하게!
DECAY_RATIO=0.1          # decay 길이 비율 (budget 대비)
MIN_DECAY_EPOCHS=1       # decay 최소 길이(epoch)
DECAY_LR=5e-4            # decay 시작 LR = stage1 의 stable LR 과 동일하게
DECAY_TYPE=sqrt          # cosine | linear | sqrt
DECAY_ETA_MIN=1e-8       # decay 도착 LR

# DECAY_EPOCHS(E) = ceil(max(DECAY_RATIO*E, MIN_DECAY_EPOCHS))  (stage1 branch 공식과 동일)
decay_len () {
  awk -v e="$1" -v r="${DECAY_RATIO}" -v mn="${MIN_DECAY_EPOCHS}" \
    'BEGIN{ d=r*e; if(d<mn) d=mn; c=int(d); if(c < d-1e-9) c++; print c }'
}

# stage1 branch 체크포인트를 epoch(=학습한 epoch 수, 1-index) 으로부터 찾는다
resolve_s1_ckpt () {
  local ckpt_dir="$1"; local epoch="$2"
  local c="${ckpt_dir}/interval_epoch_$(printf '%03d' "${epoch}").ckpt"
  if [ -f "${c}" ]; then echo "${c}"; return 0; fi
  {
    echo "ERROR: stage1 branch epoch '${epoch}' 체크포인트를 ${ckpt_dir} 에서 못 찾음."
    echo "  (DECAY_RATIO/MIN_DECAY_EPOCHS 가 run_stage1_sweep.sh 와 같은지, EPOCHS 가 SAVE_EVERY 배수인지 확인)"
    echo "사용 가능한 branch 체크포인트:"
    ls -1 "${ckpt_dir}"/interval_epoch_*.ckpt 2>/dev/null || echo "  (없음 - stage1 을 먼저 돌렸는지 확인)"
  } >&2
  return 1
}

for setting in "${SETTINGS[@]}"; do
  IFS="-" read -r h d nh <<< "${setting}"
  RUN_ID="h${h}_d${d}_n${nh}"
  S1_NAME="pretrain_one_step_${RUN_ID}"
  S1_CKPT_DIR="${S1_RESULTS_ROOT}/${RUN_ID}/${S1_NAME}/checkpoints"

  # 아키텍처는 stage1 체크포인트와 반드시 동일해야 하므로 모든 호출에 동일하게 전달
  ARCH_ARGS=(
    --model.net.init_args.hidden_size "${h}"
    --model.net.init_args.depth "${d}"
    --model.net.init_args.num_heads "${nh}"
    --model.net.patch_size "${PATCH_SIZE}"
  )

  for EP in "${EPOCHS[@]}"; do
    # 총 budget E 에서 비율 decay 를 carve-out: stable 은 (E - DECAY_EPOCHS) 에서 branch.
    DECAY_EPOCHS="$(decay_len "${EP}")"
    BRANCH_EP=$((EP - DECAY_EPOCHS))
    if [ "${BRANCH_EP}" -le 0 ]; then
      echo "SKIP: EP=${EP} <= DECAY_EPOCHS=${DECAY_EPOCHS} 라 branch 지점(${BRANCH_EP})이 유효하지 않음." >&2
      continue
    fi

    TAG="${RUN_ID}_ep${BRANCH_EP}_d${DECAY_EPOCHS}"   # 라벨: stable branch epoch + decay 길이
    echo "############################################################"
    echo "# ${TAG} :  decay only"
    echo "#   목표 총 budget=${EP}ep  =  stable@${BRANCH_EP}ep + decay ${DECAY_EPOCHS}ep"
    echo "#   decay: ${DECAY_TYPE}, lr ${DECAY_LR} -> ${DECAY_ETA_MIN}"
    echo "############################################################"

    S1_CKPT="$(resolve_s1_ckpt "${S1_CKPT_DIR}" "${BRANCH_EP}")" || exit 1
    echo "  [ckpt] stage1 branch @ ${BRANCH_EP}ep : ${S1_CKPT}"

    ##########################################################
    # DECAY (annealing) :  stage1 ckpt -> lr decay -> eta_min
    ##########################################################
    DECAY_NAME="decay_${TAG}"
    DECAY_ROOT="${DECAY_RESULTS_ROOT}/${DECAY_NAME}"
    echo "  --- DECAY (${DECAY_EPOCHS} ep, ${DECAY_TYPE}, lr ${DECAY_LR} -> ${DECAY_ETA_MIN}) ---"
    ${CMD} \
      --config "${CFG_DECAY}" \
      --trainer.max_epochs "${DECAY_EPOCHS}" \
      --trainer.default_root_dir "${DECAY_ROOT}" \
      --trainer.logger.init_args.name "${DECAY_NAME}" \
      --model.max_epochs "${DECAY_EPOCHS}" \
      --model.lr "${DECAY_LR}" \
      --model.eta_min "${DECAY_ETA_MIN}" \
      --model.decay_type "${DECAY_TYPE}" \
      --model.pretrained_path "${S1_CKPT}" \
      "${ARCH_ARGS[@]}" \
      --data.root_dir "${DATA_ROOT}" \
      --data.steps 1 \
      --data.batch_size "${BATCH_SIZE}"
    DECAY_CKPT="${DECAY_ROOT}/${DECAY_NAME}/checkpoints/last.ckpt"
    echo "  [done] ${TAG}: decayed ckpt -> ${DECAY_CKPT}"
  done
done

echo "ALL DECAY-ONLY RUNS DONE"