#!/usr/bin/env bash
# SUM Parts - unattended end-to-end run: bare machine to trained model
#
# Built for a weekend with nobody at the keyboard. Everything that can stop the
# run is checked in PREFLIGHT, before any long step, so a failure surfaces in
# the first minute rather than after three hours of setup.
#
#   preflight -> bootstrap -> env -> patches -> verify -> data -> vram -> train -> eval
#
# The one thing this cannot do for you is accept the HuggingFace dataset gate:
# it needs a browser and a logged-in account. It IS per-account though, so if
# you already accepted it elsewhere, copying the token to this machine is
# enough. Preflight fails immediately if the token is missing or rejected.
#
# Usage:
#   bash scripts/run_all.sh                 # blocking, logs to stdout + file
#   bash scripts/run_all.sh --detach        # survives terminal/session close
#
# Watch it later:
#   tail -f ~/sum-parts/runs/run_all/run.log
#   cat ~/sum-parts/runs/run_all/STATUS
set -uo pipefail

SCRIPTS="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
OUT="$HOME/sum-parts/runs/run_all"
LOG="$OUT/run.log"
STATUS="$OUT/STATUS"

CFG="${CFG:-pointvector-xl}"
EPOCHS="${EPOCHS:-100}"
VAL_FREQ="${VAL_FREQ:-5}"
# candidates tried high to low; first one that fits in VRAM wins.
# 64000 is the paper setting and needs ~16.5 GB.
VOXEL_CANDIDATES="${VOXEL_CANDIDATES:-64000 48000 40000 32000 24000}"

# ---------------------------------------------------------------- detach

if [ "${1:-}" = "--detach" ]; then
    mkdir -p "$OUT"
    echo "detaching; log: $LOG"
    setsid nohup bash "${BASH_SOURCE[0]}" > "$OUT/nohup.out" 2>&1 < /dev/null &
    sleep 2
    pgrep -af "run_all.sh" | grep -v detach || true
    exit 0
fi

mkdir -p "$OUT"
exec > >(tee -a "$LOG") 2>&1

PHASE="starting"
STARTED=$(date '+%F %T')

say()  { echo "[$(date '+%F %T')] $*"; }
head_() { echo; echo "════ $* ════"; }

write_status() {
    {
        echo "state    : $1"
        echo "phase    : $PHASE"
        echo "cfg      : $CFG"
        echo "voxel_max: ${VOXEL_MAX:-(not chosen yet)}"
        echo "started  : $STARTED"
        echo "updated  : $(date '+%F %T')"
        [ -n "${EXTRA:-}" ] && echo "note     : $EXTRA"
        echo "log      : $LOG"
    } > "$STATUS"
}

die() {
    say "GIVING UP in phase '$PHASE': $*"
    EXTRA="$*" write_status "FAILED"
    exit 1
}

# Every phase is idempotent -- bootstrap skips an existing install, the patches
# detect themselves, download_data skips cached archives, training resumes from
# its checkpoint. So retrying a phase is always safe, and so is rerunning the
# whole script from the top (see keepalive.sh).
STEP_RETRIES="${STEP_RETRIES:-4}"
STEP_BACKOFF="${STEP_BACKOFF:-60}"

step() {
    PHASE="$1"; shift
    head_ "$PHASE"
    write_status "running"

    local attempt=1 wait=$STEP_BACKOFF
    while :; do
        if "$@"; then
            [ "$attempt" -gt 1 ] && say "phase '$PHASE' succeeded on attempt $attempt"
            return 0
        fi
        if [ "$attempt" -ge "$STEP_RETRIES" ]; then
            die "$* (failed $attempt times)"
        fi
        say "phase '$PHASE' failed (attempt $attempt/$STEP_RETRIES); retrying in ${wait}s"
        EXTRA="retrying $PHASE ($attempt/$STEP_RETRIES)" write_status "retrying"
        sleep "$wait"
        attempt=$((attempt + 1))
        wait=$((wait * 2))
    done
}

# Preflight is the exception: a missing HF token or absent GPU will not fix
# itself, so retrying just burns the weekend. Fail loudly and immediately.
step_once() {
    PHASE="$1"; shift
    head_ "$PHASE"
    write_status "running"
    "$@" || die "$*"
}

# ---------------------------------------------------------------- preflight

preflight() {
    local fail=0

    say "checking tools"
    for t in git curl python3 tar; do
        command -v "$t" > /dev/null || { say "  MISSING: $t"; fail=1; }
    done

    say "checking GPU"
    if ! command -v nvidia-smi > /dev/null; then
        say "  MISSING: nvidia-smi -- the driver is not exposing a GPU here"
        fail=1
    else
        nvidia-smi --query-gpu=name,memory.total,driver_version \
                   --format=csv,noheader | sed 's/^/  /'
        VRAM_MB=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | head -1)
        say "  usable VRAM: ${VRAM_MB} MiB"
    fi

    say "checking disk (need 35 GB free in \$HOME)"
    local free_gb
    free_gb=$(df -BG --output=avail "$HOME" | tail -1 | tr -dc '0-9')
    say "  free: ${free_gb} GB"
    [ "${free_gb:-0}" -lt 35 ] && { say "  NOT ENOUGH"; fail=1; }

    say "checking HuggingFace credentials"
    # the dataset is gated; the token must exist AND the account must already
    # have accepted the licence in a browser
    local tok=""
    [ -n "${HF_TOKEN:-}" ] && tok="$HF_TOKEN"
    [ -z "$tok" ] && [ -f "$HOME/.cache/huggingface/token" ] \
        && tok=$(tr -d '\r\n' < "$HOME/.cache/huggingface/token")
    if [ -z "$tok" ]; then
        say "  MISSING: no HF token"
        say "  fix: copy the token from a machine that already accepted the gate:"
        say "       mkdir -p ~/.cache/huggingface"
        say "       echo hf_xxxxx > ~/.cache/huggingface/token"
        say "  the gate itself is per-account and needs a browser once:"
        say "       https://huggingface.co/datasets/gwxgrxhyz/SUM-Parts"
        fail=1
    else
        local code
        code=$(curl -s -o /dev/null -w '%{http_code}' -I \
            -H "Authorization: Bearer $tok" \
            "https://huggingface.co/datasets/gwxgrxhyz/SUM-Parts/resolve/main/demo.zip")
        say "  gate probe: HTTP $code"
        case "$code" in
            200|302) say "  gate OK" ;;
            401|403) say "  REJECTED -- token invalid, or this account has not accepted the gate"
                     say "  accept it in a browser: https://huggingface.co/datasets/gwxgrxhyz/SUM-Parts"
                     fail=1 ;;
            *)       say "  unexpected response; continuing but the download may fail" ;;
        esac
    fi

    [ "$fail" -eq 0 ] || return 1
    say "preflight OK"
}

# ---------------------------------------------------------------- vram pick

pick_voxel_max() {
    say "measuring which voxel_max fits in VRAM (high to low, first fit wins)"
    say "NOTE: on WSL2 an oversized value does not OOM -- the driver spills into"
    say "      host RAM and the run completes 25-100x slower. So this measures"
    say "      peak allocation instead of trusting that it 'worked'."

    source "$HOME/miniconda3/etc/profile.d/conda.sh"
    conda activate sumparts
    export WANDB_MODE=disabled WANDB_SILENT=true CUDA_HOME="$CONDA_PREFIX"
    export PYTORCH_CUDA_ALLOC_CONF="garbage_collection_threshold:0.7,max_split_size_mb:128"

    cd "$HOME/sum-parts/semantic_segmentation/PointNeXt_bundle/examples/segmentation" \
        || return 1

    local vm log line fits
    for vm in $VOXEL_CANDIDATES; do
        log="$OUT/vram_${vm}.log"
        say "  trying voxel_max=$vm"
        python -u "$SCRIPTS/bench_models.py" --iters 4 --voxel-max "$vm" \
            --cfgs "$CFG" > "$log" 2>&1
        line=$(grep -aE "^${CFG} +[0-9]" "$log" | tail -1)
        if [ -z "$line" ]; then
            say "    no result (probably OOM or an error); see $log"
            continue
        fi
        fits=$(echo "$line" | grep -o 'yes$' || true)
        say "    $(echo "$line" | awk '{print "peak", $5, "s/iter", $6}')"
        if [ -n "$fits" ]; then
            VOXEL_MAX="$vm"
            say "  chosen: voxel_max=$VOXEL_MAX"
            return 0
        fi
        say "    does not fit -- spilling to host RAM"
    done

    say "  nothing fit; falling back to the smallest candidate"
    VOXEL_MAX=$(echo "$VOXEL_CANDIDATES" | awk '{print $NF}')
    return 0
}

# ---------------------------------------------------------------- phases

do_verify() {
    source "$HOME/miniconda3/etc/profile.d/conda.sh"
    conda activate sumparts
    WANDB_MODE=disabled python "$SCRIPTS/verify_env.py"
}

do_train() {
    # Already finished? Don't retrain on a rerun.
    local done_marker="$OUT/TRAIN_DONE"
    if [ -f "$done_marker" ]; then
        say "training already completed (marker: $done_marker)"
        return 0
    fi

    # Pick up an earlier run's progress. train_watchdog only auto-discovers
    # checkpoints written after IT started, so a fresh invocation would restart
    # from epoch 1 without this.
    local ckpt
    ckpt=$(find "$HOME/sum-parts/semantic_segmentation/PointNeXt_bundle/examples/segmentation/log/sumv2_triangle" \
           -name '*_ckpt_latest.pth' -printf '%T@ %p\n' 2>/dev/null \
           | sort -rn | head -1 | cut -d' ' -f2-)

    if [ -n "$ckpt" ]; then
        say "resuming from $(basename "$ckpt")"
    else
        say "no checkpoint found; starting fresh"
    fi

    say "training $CFG for $EPOCHS epochs at voxel_max=$VOXEL_MAX"
    CFG_VOXEL_MAX="$VOXEL_MAX" \
    VAL_VOXEL_MAX="$VOXEL_MAX" \
    EPOCHS="$EPOCHS" VAL_FREQ="$VAL_FREQ" MAX_RETRIES=8 \
    RESUME_CKPT="$ckpt" \
        bash "$SCRIPTS/train_watchdog.sh" "$CFG" || return 1

    touch "$done_marker"
    return 0
}

do_eval() {
    bash "$SCRIPTS/final_eval.sh"  || say "final_eval reported non-zero (test split is blind; that is expected)"
    bash "$SCRIPTS/eval_coarse.sh" || say "eval_coarse reported non-zero"
    return 0
}

# ---------------------------------------------------------------- run

say "run_all starting -- cfg=$CFG epochs=$EPOCHS"
say "log: $LOG"
write_status "running"

step_once "preflight"   preflight

# RUN_ALL_DRYRUN lets you prove the preflight checks pass without starting the
# long phases -- worth doing before walking away for the weekend.
if [ -n "${RUN_ALL_DRYRUN:-}" ]; then
    say "DRYRUN set -- preflight passed, stopping before the real work"
    write_status "dryrun-ok"
    exit 0
fi

step "bootstrap"        bash "$SCRIPTS/bootstrap.sh"
step "conda env"        bash "$SCRIPTS/setup_env.sh"
step "cuda extensions"  bash "$SCRIPTS/setup_pointnext.sh"
step "patch numpy"      bash "$SCRIPTS/patch_numpy_aliases.sh"
step "patch test split" bash "$SCRIPTS/patch_unlabeled_test.sh"
step "patch val mode"   bash "$SCRIPTS/patch_val_mode.sh"
step "verify env"       do_verify
step "download data"    bash "$SCRIPTS/download_data.sh" all
step "prepare splits"   bash "$SCRIPTS/prepare_full_split.sh"
step "link data"        bash "$SCRIPTS/link_data.sh"
step "choose voxel_max" pick_voxel_max
step "train"            do_train
step "evaluate"         do_eval

PHASE="done"
write_status "DONE"

head_ "SUMMARY"
say "cfg       : $CFG"
say "voxel_max : $VOXEL_MAX"
grep -aE 'Best ckpt' "$OUT/../"*/train.log 2>/dev/null | tail -2
[ -f "$HOME/sum-parts/runs/coarse_eval/coarse.txt" ] && {
    echo
    echo "--- coarse (building / vegetation / vehicle / ground) ---"
    cat "$HOME/sum-parts/runs/coarse_eval/coarse.txt"
}
say "RUN ALL DONE"