#!/usr/bin/env bash
# Resolution study. Trains one (model, image_size) cell at a time.
#
# Usage:
#   ./run_resolution.sh segformer_b0 192 256              # one model, two resolutions
#   ./run_resolution.sh unet 192 256
#   ./run_resolution.sh segformer_b5 192 256              # B5 at 256 auto-drops bs to 8
#   THROTTLE=0.4 ./run_resolution.sh segformer_b0 192 256 # optional duty-cycle cap
#
# First arg is the model. Remaining args are the image sizes to sweep.
# Batch size is 16 unless the (model, image_size) combo is in the override
# list below (currently: segformer_b5 at 256 -> bs=8 for VRAM headroom).

set -euo pipefail
DIR="$(cd "$(dirname "$0")" && pwd)"
cd "$DIR"

OVERALL_T0=$(date +%s)
fmt_hms() {
    local s=$1
    printf '%d:%02d:%02d' $((s/3600)) $(((s%3600)/60)) $((s%60))
}

if [ $# -lt 2 ]; then
    echo "usage: $0 <model> <image_size> [image_size ...]"
    echo "example: $0 segformer_b0 192 256"
    exit 1
fi

MODEL="$1"
shift
RESOLUTIONS=("$@")

mkdir -p logs results

THROTTLE_ARG=()
if [ -n "${THROTTLE:-}" ] && [ "$THROTTLE" != "0" ]; then
    THROTTLE_ARG=(--throttle "$THROTTLE")
    echo "[grid] GPU duty-cycle cap: ${THROTTLE}"
fi

pick_bs() {
    local model="$1"
    local size="$2"
    # SegFormer-B5 at 256 needs smaller batch for VRAM.
    if [ "$model" = "segformer_b5" ] && [ "$size" -ge 256 ]; then
        echo 8
    else
        echo 16
    fi
}

for SIZE in "${RESOLUTIONS[@]}"; do
    BS=$(pick_bs "$MODEL" "$SIZE")
    cfg_id="${MODEL}_res${SIZE}"
    t0=$(date +%s)
    echo ""
    echo "════════════════════════════════════════════════════════════"
    echo "  $cfg_id   batch_size=$BS"
    echo "════════════════════════════════════════════════════════════"
    "${PYTHON:-python}" train.py \
        --model "$MODEL" \
        --image-size "$SIZE" \
        --batch-size "$BS" \
        "${THROTTLE_ARG[@]}" \
        2>&1 | tee "logs/${cfg_id}.stdout.log"
    dt=$(( $(date +%s) - t0 ))
    echo "  → $cfg_id finished in $(fmt_hms $dt)"
done

OVERALL_DT=$(( $(date +%s) - OVERALL_T0 ))
echo ""
echo "[grid] complete in $(fmt_hms $OVERALL_DT). Results in results/resolution_results.csv"