#!/bin/bash
# Unattended pipeline. Every stage is idempotent and resumable: re-running the
# script after any interruption picks up where it stopped, because each stage
# checks for its own output before doing work.
#
#   ./run_overnight.sh            start (or resume) in the background
#   tail -f logs/overnight.log    watch
#   cat STATUS_AUTO.md            what finished, what is pending
#
# Stage 1 must stay on the LOGIN NODE: /mnt/data (the NAS holding pgb for 2001
# and 2005-2011) is not mounted on the compute nodes -- verified, job 1400800.

set -u
ROOT=/home/tanpv/Progs/claude/HW_CFS_raw
cd $ROOT
source $(conda info --base)/etc/profile.d/conda.sh
conda activate ehd
export OMP_NUM_THREADS=1
export PYTHONPATH=$ROOT/src

LOG=$ROOT/logs/overnight.log
ST=$ROOT/STATUS_AUTO.md
NPROC=10
TARGET=$(wc -l < outputs/meta/init_both.txt)

say() { echo "[$(date '+%m-%d %H:%M:%S')] $*" | tee -a $LOG; }

status() {
  {
    echo "# Trạng thái tự động — cập nhật $(date '+%Y-%m-%d %H:%M:%S')"
    echo
    echo "| Hạng mục | Trạng thái |"
    echo "|---|---|"
    echo "| pgb đã trích | $(find outputs/Pgb2x -name '*.npy' 2>/dev/null | wc -l) / $TARGET |"
    for f in "Cube/pgb_weekly.npz|đặc trưng cao không theo tuần" \
             "AI_Results/exp_features.json|thí nghiệm biến & lớp mô hình" \
             "AI_Results/bcsd.json|thang downscaling D0/D1/D2"; do
      p=${f%%|*}; n=${f##*|}
      if [ -s "outputs/$p" ]; then echo "| $n | XONG |"; else echo "| $n | chưa |"; fi
    done
    echo
    echo "## 20 dòng nhật ký cuối"
    echo '```'
    tail -20 $LOG 2>/dev/null
    echo '```'
  } > $ST
}

# ---------------------------------------------------------------- stage 1: pgb
say "STAGE 1 — trích pgb (login node, NAS)"
for attempt in 1 2 3 4 5 6; do
  n=$(find outputs/Pgb2x -name '*.npy' 2>/dev/null | wc -l)
  if [ "$n" -ge "$TARGET" ]; then say "  pgb đủ $n/$TARGET"; break; fi
  say "  lượt $attempt: hiện $n/$TARGET, chạy $NPROC tiến trình"
  rm -f logs/expgb_failed_*.txt
  for c in $(seq 0 $((NPROC-1))); do
    nice -n 10 python src/extract_pgb.py --chunk $c --nchunk $NPROC \
        >> logs/pgbnas_$c.out 2>> logs/pgbnas_$c.err &
  done
  wait
  status
  n2=$(find outputs/Pgb2x -name '*.npy' 2>/dev/null | wc -l)
  say "  lượt $attempt xong: $n -> $n2"
  # no progress twice in a row means the rest is genuinely unreadable, not flaky
  if [ "$n2" -le "$n" ]; then say "  không tiến triển, dừng thử lại"; break; fi
done
say "STAGE 1 kết thúc: $(find outputs/Pgb2x -name '*.npy' | wc -l)/$TARGET"
status

# Stages 2-4 go to the compute nodes, not this login shell. Measured 2026-09-03:
# build_pgb_weekly used 0.09 of 16 allocated cores and bcsd 5.8 of 16, while nine
# nodes sat idle. The loop over inits (stage 2) and the loop over the six weeks
# (stages 3-4) are independent, so they are spread across nodes; ~100 min -> ~2 min.
# Unlike stage 1, which is pinned to this node by the NAS mount and capped by its
# bandwidth, so no amount of parallelism helps there.

run_week_array() {   # $1 = script, $2 = exports, $3 = result path, $4 = stage name
  sbatch --wait --job-name=ovn --export=ALL,SCRIPT="$1",EXPORTS="$2" \
         submit_week_array.slurm >> $LOG 2>&1
  eval "$2"
  python "$1" --merge >> $LOG 2>&1
  [ -s "$3" ] && say "  $4 xong" || say "  $4 LỖI"
}

# ------------------------------------------------- stage 2: upper-air features
if [ ! -s outputs/Cube/pgb_weekly.npz ]; then
  say "STAGE 2 — dẫn xuất biến cao không theo tuần (1 node x 20 core)"
  sbatch --wait --job-name=ovnpgb submit_pgbweek.slurm >> $LOG 2>&1
  [ -s outputs/Cube/pgb_weekly.npz ] && say "  STAGE 2 xong" || say "  STAGE 2 LỖI"
else
  say "STAGE 2 — đã có, bỏ qua"
fi
status

# ------------------------------------------- stage 3: features & model families
if [ ! -s outputs/AI_Results/exp_features.json ]; then
  say "STAGE 3 — biến cao không và mô hình phi tuyến mua được gì (6 node)"
  run_week_array src/exp_features.py "" outputs/AI_Results/exp_features.json "STAGE 3"
else
  say "STAGE 3 — đã có, bỏ qua"
fi
status

# ------------------------------------------------------ stage 4: downscaling
if [ ! -s outputs/AI_Results/bcsd.json ]; then
  say "STAGE 4 — thang downscaling D0 / D1 (BCSD) / D2 (ML) (6 node)"
  run_week_array src/bcsd.py "" outputs/AI_Results/bcsd.json "STAGE 4"
else
  say "STAGE 4 — đã có, bỏ qua"
fi
status
say "TOÀN BỘ CHUỖI KẾT THÚC"