#!/bin/bash
set -euo pipefail

# One-demo T-Rex Track-Force overfit over the FULL mini episode.
# Uses every valid causal anchor (~537 on trex_mini_force), not a pinned clip.
# Starts from Wan2.2 foundation weights (no prior T-Rex LoRA checkpoint).

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
DREAMZERO_ROOT="${DREAMZERO_ROOT:-$REPO_ROOT}"
TREX_DATA_ROOT="${TREX_DATA_ROOT:-$DREAMZERO_ROOT/data/trex_mini_force}"
OUTPUT_DIR="${OUTPUT_DIR:-$DREAMZERO_ROOT/checkpoints/trex_mini_force_wan22_lora_track_overfit_full}"
MAX_STEPS="${MAX_STEPS:-20000}"
SAVE_STEPS="${SAVE_STEPS:-1000}"
# Viz preview length: each chunk is 0.8s of 20 Hz track. Training still uses
# all anchors; this only controls how much of the demo W&B renders.
VIZ_NUM_CHUNKS="${VIZ_NUM_CHUNKS:-16}"
VIZ_START_CHUNK="${VIZ_START_CHUNK:-0}"

if [[ ! -f "$TREX_DATA_ROOT/meta/info.json" ]]; then
    echo "ERROR: mini dataset not found: $TREX_DATA_ROOT" >&2
    echo "Build it with: python scripts/data/build_trex_mini_force.py" >&2
    exit 1
fi
if ! [[ "$MAX_STEPS" =~ ^[1-9][0-9]*$ ]]; then
    echo "ERROR: MAX_STEPS must be a positive integer: $MAX_STEPS" >&2
    exit 1
fi
if ! [[ "$SAVE_STEPS" =~ ^[1-9][0-9]*$ ]]; then
    echo "ERROR: SAVE_STEPS must be a positive integer: $SAVE_STEPS" >&2
    exit 1
fi
if [[ -f "$OUTPUT_DIR/config.json" ]] \
    || compgen -G "$OUTPUT_DIR/checkpoint-*" >/dev/null; then
    echo "ERROR: scratch output already contains trained state: $OUTPUT_DIR" >&2
    echo "Choose a new OUTPUT_DIR to start again from step 0." >&2
    exit 1
fi

exec env \
    DREAMZERO_ROOT="$DREAMZERO_ROOT" \
    TREX_DATA_ROOT="$TREX_DATA_ROOT" \
    OUTPUT_DIR="$OUTPUT_DIR" \
    bash "$SCRIPT_DIR/trex_track_force_training_wan22.sh" \
    wandb_project=trex_mini_force_track_overfit \
    max_steps="$MAX_STEPS" \
    save_steps="$SAVE_STEPS" \
    save_total_limit=5 \
    dataset_shard_sampling_rate=1.0 \
    max_training_anchors=null \
    pin_anchor_rank=null \
    action_head_cfg.config.track_loss_weight=20.0 \
    action_head_cfg.config.track_position_loss_weight=5.0 \
    action_head_cfg.config.track_smoothness_loss_weight=2.0 \
    action_head_cfg.config.track_motion_loss_boost=4.0 \
    action_head_cfg.config.zero_invisible_track_xy=true \
    wandb_video_reconstruction_num_chunks="$VIZ_NUM_CHUNKS" \
    wandb_video_start_chunk_index="$VIZ_START_CHUNK" \
    wandb_video_track_trail_steps=2 \
    wandb_video_reconstruction_inference_steps=6 \
    wandb_video_compare_tracks_on_gt=true \
    "$@"