#!/usr/bin/env bash
# Precompute the T5 text embedding for a task prompt into the PERSISTENT cache
# (fastwam_train/text_embeds_cache/). The dataloader hard-fails on a cache
# miss, so this must exist before training. GPU strongly recommended (loads
# the Wan2.2 T5 encoder); works on CPU but slowly.
#
#   bash scripts/precompute_text_embeds.sh ["prompt string"] [task_cfg_name]
#   e.g.: bash scripts/precompute_text_embeds.sh \
#             "Lift the lid, put it aside, and pick the black plum."
#
# Uses the fastwam precompute script's `override_instruction` path: it encodes
# exactly this one prompt and does NOT scan the dataset's tasks.jsonl, so it can
# run even before the LeRobot conversion exists -- but the string MUST match the
# --task string baked into the dataset's tasks.jsonl, or the training-time hash
# lookup will miss. Writes <sha256(DEFAULT_PROMPT.format(task=...))>.t5_len128.wan22ti2v5b.pt.
set -euo pipefail

TASK_STR="${1:-place the cube in the bowl}"
# Any composable task config works here (override_instruction bypasses its
# actual data/prompt) -- just needs to exist. place_cube_flashwam_ft was
# removed 2026-08-08; pcbnew100_flashwam_scratch is the current stand-in.
TASK_CFG="${2:-pcbnew100_flashwam_scratch}"

REPO=/shared_work/amin-a/realworld-wam
FT="${REPO}/fastwam_train"
CONFIG_DIR="${FT}/configs"
CACHE="${FT}/text_embeds_cache"
FASTWAM_ROOT=/shared_work/physical_intelligence/policies/Fast-WAM/fastwam
PY=/shared_work/environments/miniconda3/envs/fastwam-robotwin-rui/bin/python

mkdir -p "${CACHE}"
cd "${FASTWAM_ROOT}"   # read-only: T5/tokenizer weights via relative checkpoints/ paths

# hydra/job_logging=stdout: else Hydra tries to write its log into the read-only
# checkout cwd and crashes at startup (George's env quirk #1).
# data.train.text_embedding_cache_dir override -> write to the PERSISTENT cache
# (the config's default points at /dev/shm, which is a training-time staging dir).
# Hydra's CLI parser treats a bare comma in a value as a list separator (e.g.
# a prompt like "Lift the lid, put it aside, ..." would otherwise be parsed as
# a 2-element list, not a string) -- wrap in escaped single-quotes to force
# string interpretation, per Hydra's own error message.
"${PY}" scripts/precompute_text_embeds.py \
    --config-dir "${CONFIG_DIR}" \
    "task=${TASK_CFG}" \
    "+override_instruction='${TASK_STR}'" \
    "data.train.text_embedding_cache_dir=${CACHE}" \
    hydra/job_logging=stdout

echo "cache contents:"
ls -la "${CACHE}"