# shellcheck shell=bash
# Shared environment for every Predictor shell script in this repository.
#
#   source "$(dirname "${BASH_SOURCE[0]}")/env.sh"
#
# Every value can be overridden from the caller's environment, e.g.
#   GPUS="0 1" DATA_ROOT=/some/disk bash scripts/run_stage1_offline_four_gpu_and_restore.sh

REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

# Python environment (torch 2.5.1 + flash-attn 2.8.3).
PYTHON_BIN="${PYTHON_BIN:-/local/zoubin/cz/envs/self_forcing/bin/python}"
# The env's torchrun wrapper has a stale shebang; launch the module directly.
TORCHRUN_CMD=("$PYTHON_BIN" -m torch.distributed.run)
VBENCH_PYTHON="${VBENCH_PYTHON:-/local/zoubin/cz/envs/vbench_eval/bin/python}"

# Sibling checkouts and model assets.  Inside the repo, checkpoints/chunkwise,
# wan_models and prompts/* are symlinks into ../Causal-Forcing and ../Self-Forcing.
# Self-Forcing-a ships scripts/evaluate_single_block_fppf.py (the FPPF reference
# evaluator); fall back to the older Self-Forcing checkout otherwise.
if [[ -z "${SELF_FORCING_ROOT:-}" ]]; then
    for candidate in "$(cd "$REPO/.." && pwd)/Self-Forcing-a" "$(cd "$REPO/.." && pwd)/Self-Forcing"; do
        if [[ -f "$candidate/scripts/evaluate_single_block_fppf.py" ]]; then
            SELF_FORCING_ROOT=$candidate; break
        fi
    done
    SELF_FORCING_ROOT="${SELF_FORCING_ROOT:-$(cd "$REPO/.." && pwd)/Self-Forcing}"
fi
# Predictor input fusion for Stage-1 launchers: self_forcing | disca | atc.
INPUT_VARIANT="${INPUT_VARIANT:-self_forcing}"
# Extra CLI flags forwarded to the Stage-1 trainers (e.g. "--atc_previous_scope last_frame").
STAGE1_EXTRA_ARGS="${STAGE1_EXTRA_ARGS:-}"
CAUSAL_CHECKPOINT="${CAUSAL_CHECKPOINT:-$REPO/checkpoints/chunkwise/causal_forcing.pt}"
CAUSAL_CONFIG="${CAUSAL_CONFIG:-$REPO/configs/causal_forcing_dmd_chunkwise.yaml}"
WAN_14B_DIR="${WAN_14B_DIR:-$REPO/wan_models/Wan2.1-T2V-14B}"
PROMPT_PATH="${PROMPT_PATH:-$REPO/prompts/vidprom_filtered_extended.txt}"
VALIDATION_PROMPT_PATH="${VALIDATION_PROMPT_PATH:-$REPO/prompts/MovieGenVideoBench_extended.txt}"
MOVIEBENCH_ORIGINAL="${MOVIEBENCH_ORIGINAL:-$REPO/prompts/MovieGenVideoBench.txt}"
MOVIEBENCH_EXTENDED="${MOVIEBENCH_EXTENDED:-$REPO/prompts/MovieGenVideoBench_extended.txt}"

# Where offline datasets and training/evaluation runs are written.
DATA_ROOT="${DATA_ROOT:-$REPO/data}"
OUTPUT_ROOT="${OUTPUT_ROOT:-$REPO/output}"

# Physical GPU indices, space separated.  GPUs 0-3 on this host are shared with
# other jobs, so the default uses the idle upper half.
GPUS="${GPUS:-4 5 6 7}"
read -r -a GPU_LIST <<< "$GPUS"
NUM_GPUS=${#GPU_LIST[@]}
GPUS_CSV=$(IFS=,; echo "${GPU_LIST[*]}")

export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}"
export MKL_NUM_THREADS="${MKL_NUM_THREADS:-4}"

# The login shell points HF_HOME/TORCH_HOME at an NVMe path that does not exist
# on this host; use a writable local cache instead (lpips/VBench weights).
CACHE_ROOT="${CACHE_ROOT:-/local/zoubin/cz/cache}"
export HF_HOME="${HF_HOME_OVERRIDE:-$CACHE_ROOT/huggingface}"
export HF_HUB_CACHE="$HF_HOME/hub"
export HUGGINGFACE_HUB_CACHE="$HF_HOME/hub"
export TORCH_HOME="$CACHE_ROOT/torch"
export SELF_FORCING_ROOT
mkdir -p "$DATA_ROOT" "$OUTPUT_ROOT" "$HF_HOME/hub" "$TORCH_HOME"

if [[ ! -x "$PYTHON_BIN" ]]; then
    echo "[env] python not found: $PYTHON_BIN (set PYTHON_BIN)" >&2
    return 1 2>/dev/null || exit 1
fi

# --- helpers ---------------------------------------------------------------

# log <master_log> <message...>
log() {
    local file=$1; shift
    printf '%s %s\n' "$(date --iso-8601=seconds)" "$*" | tee -a "$file"
}

# strided_ids <shard_index> <num_shards> <count>  -> "i i+n i+2n ..."
strided_ids() {
    local shard=$1 shards=$2 count=$3 ids=()
    for ((id=shard; id<count; id+=shards)); do ids+=("$id"); done
    echo "${ids[@]}"
}

# wait_all <master_log> <label> pid...  -> returns 1 if any child failed
wait_all() {
    local file=$1 label=$2; shift 2
    local status=0 index=0 pid
    for pid in "$@"; do
        if wait "$pid"; then
            log "$file" "[done] $label shard=$index"
        else
            log "$file" "[failed] $label shard=$index exit=$?"
            status=1
        fi
        index=$((index + 1))
    done
    return "$status"
}