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

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

TASK="${TASK:-${1:-}}"
if [[ -z "$TASK" ]]; then
  echo "Set TASK=<preset>, for example TASK=tube_task." >&2
  exit 2
fi

TASK_ENV="$ROOT/deployment/tasks/${TASK}.env"
LOCAL_TASK_ENV="$ROOT/local_pi_mamba/deployment_tasks/${TASK}.env"
if [[ ! -f "$TASK_ENV" ]]; then
  if [[ -f "$LOCAL_TASK_ENV" ]]; then
    TASK_ENV="$LOCAL_TASK_ENV"
  else
    echo "Unknown task preset: $TASK" >&2
    echo "Expected: $TASK_ENV" >&2
    echo "      or: $LOCAL_TASK_ENV" >&2
    exit 2
  fi
fi
source "$TASK_ENV"

OPENPI_REPO="${OPENPI_REPO:-$ROOT/openpi}"
OPENPI_VENV="${OPENPI_VENV:-$ROOT/.venvs/openpi}"
DEVICE="${DEVICE:-cuda}"
WARMUP_RUNS="${WARMUP_RUNS:-2}"
TIMED_RUNS="${TIMED_RUNS:-20}"
PROFILE_FLOPS="${PROFILE_FLOPS:-1}"
EXAMPLE="${EXAMPLE:-auto}"
CHECKPOINT_STEP="$(basename "$CHECKPOINT_DIR")"
METRICS_TASK_NAME="${TASK_NAME:-$TASK}"
METRICS_ROOT="${METRICS_ROOT:-$ROOT/reports/$METRICS_TASK_NAME/metrics}"
METRICS_JSON="${METRICS_JSON:-$METRICS_ROOT/${MODEL_NAME}_${CHECKPOINT_STEP}_metrics.json}"
TABLE_MD="${TABLE_MD:-$METRICS_ROOT/${MODEL_NAME}_${CHECKPOINT_STEP}_table.md}"

if [[ ! -x "$OPENPI_VENV/bin/python" ]]; then
  echo "OpenPI Python venv not found: $OPENPI_VENV" >&2
  exit 2
fi
if [[ ! -f "$OPENPI_REPO/scripts/check_policy_inference.py" ]]; then
  echo "OpenPI checkout missing check_policy_inference.py: $OPENPI_REPO" >&2
  exit 2
fi
if [[ ! -f "$CHECKPOINT_DIR/model.safetensors" ]]; then
  echo "Missing checkpoint model.safetensors: $CHECKPOINT_DIR" >&2
  exit 2
fi

mkdir -p "$METRICS_ROOT"

OPENPI_SITE_PACKAGES="$("$OPENPI_VENV/bin/python" - <<'PY'
import site
print(site.getsitepackages()[0])
PY
)"
OPENPI_NVIDIA_LIB_PATHS="$(find "$OPENPI_SITE_PACKAGES/nvidia" -maxdepth 2 -type d -name lib 2>/dev/null | paste -sd: -)"

extra=()
if [[ "$PROFILE_FLOPS" == 1 ]]; then
  extra+=(--profile-flops)
fi
if [[ -n "${NORM_STATS_PATH:-}" ]]; then
  extra+=(--norm-stats-path "$NORM_STATS_PATH")
fi

unset LOCAL_RANK RANK WORLD_SIZE MASTER_ADDR MASTER_PORT
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"
export VIRTUAL_ENV="$OPENPI_VENV"
export PATH="$OPENPI_VENV/bin:$PATH"
export LD_LIBRARY_PATH="$OPENPI_NVIDIA_LIB_PATHS:${LD_LIBRARY_PATH:-}"

cd "$OPENPI_REPO"
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True "$OPENPI_VENV/bin/python" scripts/check_policy_inference.py \
  --checkpoint-dir "$CHECKPOINT_DIR" \
  --config-name "$CONFIG" \
  --pytorch-device "$DEVICE" \
  --num-denoise-steps "$NUM_DENOISE_STEPS" \
  --warmup-runs "$WARMUP_RUNS" \
  --timed-runs "$TIMED_RUNS" \
  --example "$EXAMPLE" \
  --metrics-out-path "$METRICS_JSON" \
  --table-out-path "$TABLE_MD" \
  --model-label "$MODEL_NAME" \
  "${extra[@]}"

echo "Metrics JSON: $METRICS_JSON"
echo "Table: $TABLE_MD"
echo "Serve with this reference explicitly if needed:"
echo "REFERENCE_METRICS_JSON=$METRICS_JSON TASK=$TASK bash deployment/scripts/serve_policy.sh"