#!/bin/bash
set -euo pipefail

# FACT full training (from scratch, no LoRA) on a T-Rex-format LeRobot-v2.1
# dataset: offline caches (frozen Wan2.2 VAE latents sharded over all GPUs,
# umT5-XXL text embeddings on node 0), then torch.distributed DDP training with
# checkpoints every SAVE_STEPS, W&B scalars and a periodic monitoring video.
#
# Usage (inside tmux):
#   [source your env]                       # optional: FACT_ENV_FILE=/path/to/env.sh
#   bash scripts/train_full.sh [extra train_ddp.py args]
# Multi-node: run on every node with NNODES / NODE_RANK / MASTER_ADDR set
# (node rank 0 = MASTER_ADDR; caches are built cooperatively on shared storage).
# Env overrides: DATA_ROOT OUTPUT_DIR CACHE_DIR FACT_CONFIG WAN22_CKPT_DIR NUM_GPUS
#                PER_DEVICE_BATCH LR MAX_STEPS SAVE_STEPS VIDEO_STEPS WANDB_PROJECT
#                WANDB_RUN_NAME REQUIRE_FORCE(0/1)

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
DATA_ROOT="${DATA_ROOT:-${TREX_DATA_ROOT:-$REPO_ROOT/data/trex_full_force}}"
OUTPUT_DIR="${OUTPUT_DIR:-$REPO_ROOT/checkpoints/$(basename "$DATA_ROOT")_fact}"
CACHE_DIR="${CACHE_DIR:-${FACT_CACHE_DIR:-$REPO_ROOT/cache/$(basename "$DATA_ROOT")}}"
FACT_CONFIG="${FACT_CONFIG:-$REPO_ROOT/configs/fact_base_text.json}"
WAN22_CKPT_DIR="${WAN22_CKPT_DIR:-$REPO_ROOT/checkpoints/Wan2.2-TI2V-5B}"
NUM_GPUS="${NUM_GPUS:-16}"
NNODES="${NNODES:-1}"
NODE_RANK="${NODE_RANK:-0}"
MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
MASTER_PORT="${MASTER_PORT:-29701}"
PER_DEVICE_BATCH="${PER_DEVICE_BATCH:-8}"
LR="${LR:-3e-4}"
MAX_STEPS="${MAX_STEPS:-100000}"
SAVE_STEPS="${SAVE_STEPS:-10000}"
VIDEO_STEPS="${VIDEO_STEPS:-1000}"
VIDEO_SAVE_LOCAL="${VIDEO_SAVE_LOCAL:-0}"
WANDB_PROJECT="${WANDB_PROJECT:-fact}"
WANDB_RUN_NAME="${WANDB_RUN_NAME:-$(basename "$OUTPUT_DIR")}"
REQUIRE_FORCE="${REQUIRE_FORCE:-0}"      # 1 = dataset must carry the trex_track_force tactile schema/manifest

# ---- environment ------------------------------------------------------------
if [ -n "${FACT_ENV_FILE:-}" ] && [ -f "$FACT_ENV_FILE" ]; then
    set +eu
    # shellcheck disable=SC1090
    source "$FACT_ENV_FILE"
    set -eu
fi
export PYTHONPATH="$REPO_ROOT${PYTHONPATH:+:$PYTHONPATH}"
cd "$REPO_ROOT"
filt() { grep --line-buffered -v -E "libhggcrt1|ALINPU|^\s*$"; }

if [ ! -f "$DATA_ROOT/meta/info.json" ]; then
    echo "ERROR: LeRobot-v2 dataset not found: $DATA_ROOT" >&2; exit 1
fi
for required in "$WAN22_CKPT_DIR/Wan2.2_VAE.pth" "$WAN22_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth" \
                "$WAN22_CKPT_DIR/google/umt5-xxl/tokenizer.json" "$FACT_CONFIG"; do
    if [ ! -f "$required" ]; then echo "ERROR: missing file: $required" >&2; exit 1; fi
done

READY_ARGS=(--dataset-root "$DATA_ROOT")
if [ "$REQUIRE_FORCE" = "1" ]; then READY_ARGS+=(--require-force); fi
python "$REPO_ROOT/scripts/check_dataset_ready.py" "${READY_ARGS[@]}" 2>&1 | filt

# ---- offline caches ---------------------------------------------------------
N_EPISODES=$(python -c "import json;print(json.load(open('$DATA_ROOT/meta/info.json'))['total_episodes'])")
count_cached() { ls "$CACHE_DIR/latents" 2>/dev/null | grep -c '^episode_.*\.npy$' || true; }
TOTAL_SHARDS=$((NNODES * NUM_GPUS))
if [ "$(count_cached)" -lt "$N_EPISODES" ]; then
    echo ">>> building VAE latent cache: $(count_cached)/$N_EPISODES episodes cached; node $NODE_RANK runs shards $((NODE_RANK * NUM_GPUS))..$((NODE_RANK * NUM_GPUS + NUM_GPUS - 1)) of $TOTAL_SHARDS"
    mkdir -p "$CACHE_DIR/latents"
    for ((i = 0; i < NUM_GPUS; i++)); do
        SHARD=$((NODE_RANK * NUM_GPUS + i))
        CUDA_VISIBLE_DEVICES=$i python scripts/build_cache_dataset.py --what latents \
            --dataset-root "$DATA_ROOT" --out "$CACHE_DIR" \
            --vae "$WAN22_CKPT_DIR/Wan2.2_VAE.pth" --shard "$SHARD" --num-shards "$TOTAL_SHARDS" \
            2>&1 | filt | sed "s/^/[shard $SHARD] /" &
    done
    wait
fi
if [ "$NODE_RANK" = "0" ] && [ ! -f "$CACHE_DIR/text_cache.npz" ]; then
    echo ">>> building umT5 text cache"
    CUDA_VISIBLE_DEVICES=0 python scripts/build_cache_dataset.py --what text \
        --dataset-root "$DATA_ROOT" --out "$CACHE_DIR" \
        --text-encoder "$WAN22_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth" \
        --tokenizer "$WAN22_CKPT_DIR/google/umt5-xxl" --max-text-tokens 64 2>&1 | filt
fi
until [ -f "$CACHE_DIR/text_cache.npz" ] && [ "$(count_cached)" -ge "$N_EPISODES" ]; do
    echo ">>> node $NODE_RANK waiting for caches: $(count_cached)/$N_EPISODES latents, text=$([ -f "$CACHE_DIR/text_cache.npz" ] && echo yes || echo no)"
    sleep 60
done

# ---- W&B: online if credentials exist, otherwise offline --------------------
if [ -z "${WANDB_API_KEY:-}" ] && ! grep -qs "api.wandb.ai" "$HOME/.netrc"; then
    echo ">>> no W&B credentials found; using WANDB_MODE=offline"
    export WANDB_MODE="${WANDB_MODE:-offline}"
fi

mkdir -p "$OUTPUT_DIR"
ulimit -n 65536 2>/dev/null || true
echo ">>> FACT training: $NNODES node(s) x $NUM_GPUS GPUs x batch $PER_DEVICE_BATCH, $MAX_STEPS steps, save every $SAVE_STEPS -> $OUTPUT_DIR"
if [ "$NNODES" = "1" ]; then
    LAUNCH=(python -m torch.distributed.run --nproc_per_node="$NUM_GPUS" --standalone)
else
    LAUNCH=(python -m torch.distributed.run --nproc_per_node="$NUM_GPUS" --nnodes="$NNODES" --node_rank="$NODE_RANK" \
            --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT")
fi
"${LAUNCH[@]}" \
    "$REPO_ROOT/scripts/train_ddp.py" \
    --dataset-root "$DATA_ROOT" \
    --cache-dir "$CACHE_DIR" \
    --output-dir "$OUTPUT_DIR" \
    --config "$FACT_CONFIG" \
    --blocks 4 \
    --max-steps "$MAX_STEPS" \
    --per-device-batch "$PER_DEVICE_BATCH" \
    --lr "$LR" \
    --min-lr-ratio 0.1 \
    --warmup-steps 2000 \
    --weight-decay 0.01 \
    --clip 1.0 \
    --seed 42 \
    --workers 4 \
    --bf16 1 \
    --log-steps 20 \
    --save-steps "$SAVE_STEPS" \
    --save-total-limit 10 \
    --resume auto \
    --report-to wandb \
    --wandb-project "$WANDB_PROJECT" \
    --wandb-run-name "$WANDB_RUN_NAME" \
    --video-steps "$VIDEO_STEPS" \
    --video-episode 0 \
    --video-num-chunks 4 \
    --video-save-local "$VIDEO_SAVE_LOCAL" \
    "$@" 2>&1 | filt | tee -a "$OUTPUT_DIR/train_node${NODE_RANK}.log"