#!/usr/bin/env bash
# Run one DAgger round on a (drive, world) combo.
#
# Usage: tools/dagger_round.sh <drive> <world> [seeds_per_n] [round_idx]
#
# Collects DAgger demos using the current BC policy as the actor and the
# universal teacher as the labeller, in the HERDING_WEBOTS preset env
# (140° FOV, tight tracker — matches deployment). Concatenates with the
# original BC demos, re-trains BC, and saves to runs/bc_dagger_<combo>/.

set -euo pipefail
ROOT="$( cd "$( dirname "${BASH_SOURCE[0]}" )/.." && pwd )"
cd "$ROOT"

DRIVE="${1:-differential}"
WORLD="${2:-field}"
SEEDS="${3:-15}"
ROUND="${4:-1}"

TAG="${DRIVE}_${WORLD}"
ORIG_DEMOS="training/bc/demos_${TAG}.npz"
DAGGER_DEMOS="training/bc/dagger${ROUND}_${TAG}.npz"
COMBINED_DEMOS="training/bc/combined${ROUND}_${TAG}.npz"
BC_DIR="training/runs/bc_${TAG}"
OUT_DIR="training/runs/bc_dagger${ROUND}_${TAG}"

case "$WORLD" in
    field_round)
        EPOCHS=150
        ;;
    *)
        EPOCHS=60
        ;;
esac

echo "=== DAgger round ${ROUND}: ${DRIVE}/${WORLD} ==="
echo "    Actor policy: ${BC_DIR}/policy.zip"
echo "    Output:       ${OUT_DIR}/policy.zip"

# 1. Collect DAgger demos: BC drives, teacher labels (privileged + HERDING_WEBOTS).
python -m training.bc.collect \
    --teacher universal --out "$DAGGER_DEMOS" \
    --seeds-per-n "$SEEDS" --subsample 3 \
    --frame-stack 4 --drive-mode "$DRIVE" --world "$WORLD" \
    --max-steps 30000 \
    --privileged --use-webots-preset \
    --fp-rate 0.0 --action-smooth 0.55 --wheel-slip-std 0.05 \
    --dagger-policy "$BC_DIR"

# 2. Concatenate original demos + dagger demos.
python - <<PY
import numpy as np
orig = np.load("${ORIG_DEMOS}")
dag  = np.load("${DAGGER_DEMOS}")
obs = np.concatenate([orig["obs"], dag["obs"]], axis=0)
act = np.concatenate([orig["actions"], dag["actions"]], axis=0)
np.savez("${COMBINED_DEMOS}", obs=obs, actions=act,
         meta=np.concatenate([orig["meta"], dag["meta"]], axis=0))
print(f"[combine] orig={orig['obs'].shape[0]} + dagger={dag['obs'].shape[0]} = {obs.shape[0]}")
PY

# 3. Re-train BC on combined demos.
python -m training.bc.pretrain \
    --demos "$COMBINED_DEMOS" --out "$OUT_DIR" \
    --epochs "$EPOCHS" --net-arch 512,512

echo "=== DAgger round ${ROUND} done: ${OUT_DIR}/policy.zip ==="