#!/usr/bin/env bash
set -euo pipefail
# bash examples.sh
if [[ -n "${TERM:-}" ]]; then
    clear
fi
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

# Shared HF cache used on this cluster.
HF_HUB_CACHE_DIR="${HF_HUB_CACHE_DIR:-/scratch/rr81/ma5430/.cache/huggingface/hub}"
export HF_HUB_CACHE="$HF_HUB_CACHE_DIR"
export HUGGINGFACE_HUB_CACHE="$HF_HUB_CACHE_DIR"
export HF_HOME="$(dirname "$HF_HUB_CACHE_DIR")"

# GPU nodes have no internet, while login nodes do.
# Auto default: offline on GPU nodes, online on login nodes.
DEFAULT_OFFLINE_MODE="1"
if ! (command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L >/dev/null 2>&1); then
    DEFAULT_OFFLINE_MODE="0"
fi
OFFLINE_MODE="${OFFLINE_MODE:-$DEFAULT_OFFLINE_MODE}"

if [[ "$OFFLINE_MODE" == "1" ]]; then
    export HF_DATASETS_OFFLINE="1"
    export HF_METRICS_OFFLINE="1"
    export HF_MODULES_OFFLINE="1"
    export TRANSFORMERS_OFFLINE="1"
    export DIFFUSERS_OFFLINE="1"
    export HF_HUB_OFFLINE="1"
else
    export HF_DATASETS_OFFLINE="0"
    export HF_METRICS_OFFLINE="0"
    export HF_MODULES_OFFLINE="0"
    export TRANSFORMERS_OFFLINE="0"
    export DIFFUSERS_OFFLINE="0"
    export HF_HUB_OFFLINE="0"
fi

# Existing environment requested by user.
PYTHON_BIN="${PYTHON_BIN:-/g/data/rr81/aev/bin/python}"
if [[ ! -x "$PYTHON_BIN" ]]; then
    echo "[examples.sh] Missing Python executable: $PYTHON_BIN" >&2
    exit 1
fi

DATASET_NAME="${DATASET_NAME:-pickapic}"      # coco | pickapic
GRAD_CONFIG="${GRAD_CONFIG:-one_step_rectification_config}"
MODEL_PROFILE="${MODEL_PROFILE:-sana_600m_512}"  # sana_600m_512 | sana_1600m_512 | sana_sprint_0_6b_1024 | sana_sprint_1_6b_1024
MODE="${MODE:-gradient_ascent}"               # gradient_ascent | baseline | both
# Empty MAX_SAMPLES means evaluate all available samples.
MAX_SAMPLES="${MAX_SAMPLES:-}"
NUM_STEPS="${NUM_STEPS:-20}"
CFG_SCALE="${CFG_SCALE:-4.5}"
DTYPE="${DTYPE:-bf16}"                    # auto | bf16 | fp16 | fp32
METRICS="${METRICS:-clip aesthetic pickscore hpsv2 hpsv21 imagereward}"
PREFETCH_ONLY="${PREFETCH_ONLY:-0}"

# Override this path whenever you want to swap reward weights.
# LRM_MODEL_PATH="${LRM_MODEL_PATH:-/g/data/rr81/LPO/lrm/lrm_sana/logs/v8/reward_model/step_sana_sana_600m_512_variable-t_lr1e-5_step-8000_filter2_time951/checkpoint-gstep33000}"
LRM_MODEL_PATH="${LRM_MODEL_PATH:-/g/data/rr81/LPO/lrm/lrm_sana/logs/v7/reward_model/step_sana_sana_600m_512_variable-t_lr1e-5_step-8000_filter2_time951/checkpoint-gstep32000}"

if [[ -z "${GPU_ID:-}" ]]; then
    if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L >/dev/null 2>&1; then
        GPU_ID="$(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader,nounits | sort -k2 -n | head -n1 | cut -d',' -f1 | tr -d ' ')"
        GPU_ID="${GPU_ID:-0}"
    else
        GPU_ID="0"
        echo "[examples.sh] No visible NVIDIA GPU on this node. Defaulting GPU_ID=0."
        echo "[examples.sh] eval.py will run on CPU if CUDA is unavailable."
    fi
fi

echo "Using GPU ID: $GPU_ID"
echo "Using LRM weights: $LRM_MODEL_PATH"
echo "HF offline mode: $OFFLINE_MODE"
echo "Generation dtype: $DTYPE"

if [[ "$PREFETCH_ONLY" == "1" ]]; then
    echo "[examples.sh] PREFETCH_ONLY=1 -> downloading required model files to shared cache and exiting."
    export MODEL_PROFILE
    export METRICS
    "$PYTHON_BIN" - <<'PY'
import os
from huggingface_hub import hf_hub_download, snapshot_download

cache_dir = os.environ["HF_HUB_CACHE"]
model_profile = os.environ.get("MODEL_PROFILE", "sana_600m_512")
metrics = set(os.environ.get("METRICS", "clip aesthetic").split())

profile_to_repo = {
    "sana_600m_512": "Efficient-Large-Model/Sana_600M_512px_diffusers",
    "sana_1600m_512": "Efficient-Large-Model/Sana_1600M_512px_diffusers",
    "sana_sprint_0_6b_1024": "Efficient-Large-Model/Sana_Sprint_0.6B_1024px_diffusers",
    "sana_sprint_1_6b_1024": "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers",
}

def snap(repo_id):
    print(f"[prefetch] snapshot_download: {repo_id}")
    snapshot_download(repo_id=repo_id, cache_dir=cache_dir, local_files_only=False)

def one(repo_id, filename):
    print(f"[prefetch] hf_hub_download: {repo_id}/{filename}")
    hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=cache_dir, local_files_only=False)

if model_profile not in profile_to_repo:
    raise ValueError(f"Unknown MODEL_PROFILE={model_profile}")

# Base SANA model used for generation + reward backbone
snap(profile_to_repo[model_profile])

# Required for CLIP-based metrics and LRM text projection init fallback
if "clip" in metrics or "aesthetic" in metrics:
    snap("openai/clip-vit-large-patch14")

if "pickscore" in metrics:
    snap("laion/CLIP-ViT-H-14-laion2B-s32B-b79K")
    snap("yuvalkirstain/PickScore_v1")

if "hpsv2" in metrics or "hpsv21" in metrics:
    one("laion/CLIP-ViT-H-14-laion2B-s32B-b79K", "open_clip_pytorch_model.bin")
if "hpsv2" in metrics:
    one("xswu/HPSv2", "HPS_v2_compressed.pt")
if "hpsv21" in metrics:
    one("xswu/HPSv2", "HPS_v2.1_compressed.pt")

if "imagereward" in metrics:
    one("THUDM/ImageReward", "ImageReward.pt")
    one("THUDM/ImageReward", "med_config.json")

print("[prefetch] done")
PY
    exit 0
fi

read -r -a METRICS_ARR <<< "$METRICS"

CMD=(
    "$PYTHON_BIN" eval.py
    --model_variant "$MODEL_PROFILE"
    --dataset_type "$DATASET_NAME"
    --lrm_model "$LRM_MODEL_PATH"
    --grad_config "$GRAD_CONFIG"
    --metrics "${METRICS_ARR[@]}"
    --num_steps "$NUM_STEPS"
    --cfg_scale "$CFG_SCALE"
    --dtype "$DTYPE"
    --hf_cache_dir "$HF_HUB_CACHE_DIR"
    --output_dir "RESULTS/$DATASET_NAME/${GRAD_CONFIG}_${MODEL_PROFILE}"
    --cuda "$GPU_ID"
    --mode "$MODE"
)

if [[ -n "$MAX_SAMPLES" ]]; then
    CMD+=(--max_samples "$MAX_SAMPLES")
fi

if [[ "$OFFLINE_MODE" == "1" ]]; then
    CMD+=(--offline)
fi

"${CMD[@]}"