#!/usr/bin/env bash
set -o pipefail

# CONFIG SETTING
NUM_GPUS=${NUM_GPUS:-1}
MASTER_PORT=${MASTER_PORT:-29110}
OFFLOAD=${OFFLOAD:-0}   # 1 or true = mutually-exclusive GPU residency for generator/VAE

EXP_NAME=decmem
CONFIG_PATH=${CONFIG_PATH:-configs/${EXP_NAME}.yaml}
CHECKPOINT_PATH=checkpoints/decmem.pt
DATA_PATH=example/example.csv

NUM_CONDITION_FRAMES=${NUM_CONDITION_FRAMES:-56}
NUM_INFERENCE_STEPS=${NUM_INFERENCE_STEPS:-20}
GUIDANCE_SCALE=${GUIDANCE_SCALE:-5}
SEED=${SEED:-42}
START_INDEX=${START_INDEX:-0}
N=${N:-2}

T_LAT=${T_LAT:-160}   # optional; empty = use config value
T_TAG=${T_LAT:+_T${T_LAT}}

OUTPUT_DIR=${OUTPUT_DIR:-test/res/${EXP_NAME}_${T_TAG}_con_${NUM_CONDITION_FRAMES}_inf_${NUM_INFERENCE_STEPS}_gs_${GUIDANCE_SCALE}}
mkdir -p "$OUTPUT_DIR"
RANK0_LOG="${OUTPUT_DIR}/rank0.log"

SCRIPT_ARGS=(
    --config_path "$CONFIG_PATH"
    --checkpoint_path "$CHECKPOINT_PATH"
    --data_path "$DATA_PATH"
    --output_dir "$OUTPUT_DIR"
    --num_condition_frames "$NUM_CONDITION_FRAMES"
    --guidance_scale "$GUIDANCE_SCALE"
    --num_inference_steps "$NUM_INFERENCE_STEPS"
    --seed "$SEED")

SCRIPT_ARGS+=(--start_index "$START_INDEX")
if [[ -n "$N" ]]; then
    SCRIPT_ARGS+=(--num_entries "$N")
fi

if [[ -n "$T_LAT" ]]; then
    SCRIPT_ARGS+=(--num_latent_frames "$T_LAT")
fi

if [[ "$NUM_GPUS" -gt 1 ]]; then
    echo "Multi-GPU inference: NUM_GPUS=${NUM_GPUS}  MASTER_PORT=${MASTER_PORT}"
    torchrun \
        --standalone \
        --nproc_per_node="$NUM_GPUS" \
        --master_port="$MASTER_PORT" \
        --local-ranks-filter=0 \
        infer.py \
        "${SCRIPT_ARGS[@]}" \
        2>&1 | tee "$RANK0_LOG"
else
    torchrun infer.py "${SCRIPT_ARGS[@]}" 2>&1 | tee "$RANK0_LOG"
fi