#!/usr/bin/env bash
# SUM Parts - emit training progress as discrete events, exit when the run ends
#
# Designed to be selective: a 100-epoch run writes ~100 epoch lines, which would
# be 100 notifications. So this emits every Nth epoch, plus anything that would
# change what to do next -- watchdog retries, tracebacks, OOM, GPU trouble, and
# the final outcome.
#
# Silence must not be mistakable for success, so the terminal check keys off the
# watchdog process disappearing, not off a success marker appearing.
set -uo pipefail

RUNROOT="$HOME/sum-parts/runs"
EVERY="${EVERY:-10}"          # emit one progress line per this many epochs
POLL="${POLL:-120}"
TOTAL="${TOTAL:-100}"         # target epoch count, for the "n/TOTAL" display

WORKDIR="${1:-}"
if [ -z "$WORKDIR" ]; then
    WORKDIR=$(cat "$RUNROOT/.latest" 2>/dev/null || true)
fi
if [ -z "$WORKDIR" ] || [ ! -d "$WORKDIR" ]; then
    echo "MONITOR-ERROR: no run directory (looked at $RUNROOT/.latest)"
    exit 1
fi

TRAINLOG="$WORKDIR/train.log"
WDLOG="$WORKDIR/watchdog.log"

echo "MONITOR-START $(basename "$WORKDIR") every=${EVERY}ep poll=${POLL}s"

last_epoch_reported=0
wd_lines=0
err_lines=0

while true; do
    # --- progress ------------------------------------------------------
    if [ -f "$TRAINLOG" ]; then
        line=$(grep -aE 'Epoch [0-9]+ LR' "$TRAINLOG" 2>/dev/null | tail -1)
        if [ -n "$line" ]; then
            # Report the epoch number the trainer itself prints. Counting log
            # lines is wrong after a resume: the new log starts mid-run, so a
            # run continuing at epoch 31 would be announced as epoch 11.
            ep=$(echo "$line" | grep -oE 'Epoch [0-9]+' | grep -oE '[0-9]+')
            if [ -n "$ep" ] && [ $((ep / EVERY)) -gt $((last_epoch_reported / EVERY)) ]; then
                best=$(grep -aE 'Find a better ckpt' "$TRAINLOG" | tail -1 \
                       | grep -oE 'val_miou [0-9.]+' | tail -1)
                echo "EPOCH $ep/$TOTAL | ${line#*] } | best: ${best:-none}"
                last_epoch_reported=$ep
            fi
        fi
    fi

    # --- watchdog events (retries, resumes, give-ups) ------------------
    if [ -f "$WDLOG" ]; then
        c=$(wc -l < "$WDLOG")
        if [ "$c" -gt "$wd_lines" ]; then
            tail -n +$((wd_lines + 1)) "$WDLOG" \
                | grep -E 'attempt|exited rc=|giving up|exhausted|finished cleanly|FATAL' || true
            wd_lines=$c
        fi
    fi

    # --- speed regression -----------------------------------------------
    # The first run silently fell from 0.29 s/it to 3.1 s/it when the caching
    # allocator pool overflowed VRAM into host RAM. No error, no OOM -- just a
    # 24-hour ETA. Watch the rate, not only the exit code.
    if [ -f "$TRAINLOG" ]; then
        r=$(grep -aoE '[0-9.]+(s/it|it/s)' "$TRAINLOG" | tail -1)
        if [ -n "$r" ]; then
            sec=$(awk -v x="$r" 'BEGIN {
                if (x ~ /it\/s$/) { sub(/it\/s$/,"",x); printf "%.3f", (x>0 ? 1.0/x : 0) }
                else              { sub(/s\/it$/,"",x); printf "%.3f", x }
            }')
            slow=$(awk -v v="$sec" 'BEGIN { print (v > 1.5) ? 1 : 0 }')
            if [ "$slow" = "1" ] && [ "${slow_reported:-0}" = "0" ]; then
                mem=$(nvidia-smi --query-gpu=memory.used,power.draw --format=csv,noheader 2>/dev/null)
                echo "SLOWDOWN: ${sec} s/it (healthy is ~0.29) | GPU $mem"
                slow_reported=1
            elif [ "$slow" = "0" ]; then
                slow_reported=0
            fi
        fi
    fi

    # --- failure signatures in the trainer itself ----------------------
    if [ -f "$TRAINLOG" ]; then
        e=$(grep -cE 'Traceback|CUDA error|out of memory|OutOfMemory|Killed|AssertionError' \
            "$TRAINLOG" 2>/dev/null || echo 0)
        if [ "$e" -gt "$err_lines" ]; then
            echo "ERROR-SIGNAL x$((e - err_lines)):"
            grep -E 'Traceback|CUDA error|out of memory|OutOfMemory|Killed|AssertionError' \
                "$TRAINLOG" | tail -3
            err_lines=$e
        fi
    fi

    # --- terminal check: watchdog gone means the run is over, either way
    if ! pgrep -f train_watchdog.sh > /dev/null 2>&1; then
        sleep 5
        state=$(grep -E '^state' "$WORKDIR/status.txt" 2>/dev/null | cut -d: -f2- | xargs)
        final=$(grep -aE 'Epoch [0-9]+ LR' "$TRAINLOG" 2>/dev/null | tail -1 \
                | grep -oE 'Epoch [0-9]+' | grep -oE '[0-9]+')
        echo "MONITOR-END state='${state:-unknown}' last_epoch=${final:-0}/$TOTAL"
        if [ -f "$TRAINLOG" ]; then
            echo "--- tail ---"
            tail -8 "$TRAINLOG"
        fi
        exit 0
    fi

    sleep "$POLL"
done