#!/usr/bin/env bash
# Build real DePlot evidence and calibrate the 7B confidence gate on ChartQA val.
set -euo pipefail

PROJECT_DIR="$(cd "$(dirname "$0")/.." && pwd)"
cd "${PROJECT_DIR}"

PYTHON_ENV="${PYTHON_ENV:-${PROJECT_DIR}/.venv}"
PYTHON_BIN="${PYTHON_BIN:-${PYTHON_ENV}/bin/python}"
TORCHRUN_BIN="${TORCHRUN_BIN:-${PYTHON_ENV}/bin/torchrun}"
RAW_VALIDATION_JSON="${RAW_VALIDATION_JSON:-${PROJECT_DIR}/data/images/chartqa/json/val.json}"
DEPLOT_VALIDATION_JSON="${DEPLOT_VALIDATION_JSON:-${PROJECT_DIR}/data/chartqa/teacher_gate/llava_onevision_7b_deplot_greedy/validation_real_deplot.json}"
VALIDATION_JSONL="${VALIDATION_JSONL:-${PROJECT_DIR}/data/chartqa/teacher_gate/llava_onevision_7b_deplot_greedy/validation.jsonl}"
TRAIN_JSONL="${TRAIN_JSONL:-${PROJECT_DIR}/data/chartqa/teacher_gate/llava_onevision_7b_deplot_greedy/train.jsonl}"
DEPLOT_CACHE="${DEPLOT_CACHE:-${PROJECT_DIR}/data/chartqa/teacher_gate/llava_onevision_7b_deplot_greedy/validation_deplot_cache.json}"
DEPLOT_PATH="${DEPLOT_PATH:-/data/lyf/.cache/huggingface/manual/google-deplot-6e76d624}"
TEACHER_PATH="${TEACHER_PATH:-/data/lyf/.cache/huggingface/hub/models--llava-hf--llava-onevision-qwen2-7b-ov-hf/snapshots/0d50680527681998e456c7b78950205bedd8a068}"
PROCESSOR_PATH="${PROCESSOR_PATH:-/data/lyf/.cache/huggingface/manual/llava-onevision-qwen2-7b-ov-hf-processor-0d506805}"
MIN_FREE_MIB="${MIN_FREE_MIB:-80000}"
MAX_USED_MIB="${MAX_USED_MIB:-2048}"
MAX_UTILIZATION="${MAX_UTILIZATION:-20}"
BATCH_SIZE="${BATCH_SIZE:-4}"
WAIT_FOR_GPUS="${WAIT_FOR_GPUS:-1}"
WAIT_SECONDS="${WAIT_SECONDS:-30}"

for required_path in \
  "${PYTHON_BIN}" "${TORCHRUN_BIN}" "${RAW_VALIDATION_JSON}" \
  "${DEPLOT_PATH}" "${TEACHER_PATH}" "${PROCESSOR_PATH}"; do
  if [[ ! -e "${required_path}" ]]; then
    echo "Required validation-calibration path does not exist: ${required_path}" >&2
    exit 2
  fi
done

select_idle_gpus() {
  {
    nvidia-smi \
      --query-gpu=index,memory.used,memory.free,utilization.gpu \
      --format=csv,noheader,nounits
  } | awk -F, \
      -v max_used="${MAX_USED_MIB}" \
      -v min_free="${MIN_FREE_MIB}" \
      -v max_util="${MAX_UTILIZATION}" '
        {
          gsub(/[[:space:]]/, "", $1)
          if (($2 + 0) <= max_used && ($3 + 0) >= min_free && ($4 + 0) <= max_util) {
            print $1
          }
        }
      ' | paste -sd, -
}

if [[ "${WAIT_FOR_GPUS}" != "0" && "${WAIT_FOR_GPUS}" != "1" ]]; then
  echo "WAIT_FOR_GPUS must be 0 or 1, got ${WAIT_FOR_GPUS}." >&2
  exit 2
fi
if [[ -z "${TARGET_GPUS:-}" ]]; then
  while true; do
    TARGET_GPUS="$(select_idle_gpus)"
    if [[ -n "${TARGET_GPUS}" ]]; then
      break
    fi
    if [[ "${WAIT_FOR_GPUS}" != "1" ]]; then
      echo "No GPU satisfies the validation-calibration idle/capacity thresholds." >&2
      exit 3
    fi
    printf '[%s] waiting for a truly idle GPU for validation calibration.\n' "$(date -Is)"
    sleep "${WAIT_SECONDS}"
  done
fi

IFS=',' read -r -a selected_gpu_list <<< "${TARGET_GPUS}"
NUM_PROCESSES="${#selected_gpu_list[@]}"
mkdir -p "$(dirname "${VALIDATION_JSONL}")"
LOG_PATH="${VALIDATION_JSONL}.build.log"
DEPLOT_INPUT="${RAW_VALIDATION_JSON}"
if [[ -s "${DEPLOT_VALIDATION_JSON}" ]]; then
  DEPLOT_INPUT="${DEPLOT_VALIDATION_JSON}"
fi

printf '[%s] validation calibration: GPUs=%s ranks=%s deplot_input=%s output=%s\n' \
  "$(date -Is)" "${TARGET_GPUS}" "${NUM_PROCESSES}" "${DEPLOT_INPUT}" "${VALIDATION_JSONL}" \
  | tee -a "${LOG_PATH}"

env \
  -u HTTP_PROXY -u HTTPS_PROXY -u ALL_PROXY \
  -u http_proxy -u https_proxy -u all_proxy \
  CUDA_VISIBLE_DEVICES="${TARGET_GPUS}" \
  HF_HUB_OFFLINE=1 \
  TRANSFORMERS_OFFLINE=1 \
  TOKENIZERS_PARALLELISM=false \
  "${PYTHON_BIN}" scripts/build_visual_facts_chartqa_deplot.py \
  --input "${DEPLOT_INPUT}" \
  --output "${DEPLOT_VALIDATION_JSON}" \
  --cache "${DEPLOT_CACHE}" \
  --model-id "${DEPLOT_PATH}" \
  --dtype float32 \
  --batch-size 8 \
  --worker-chunk-size 128 \
  --max-new-tokens 384 \
  --only-missing \
  --all-gpus \
  2>&1 | tee -a "${LOG_PATH}"

env \
  -u HTTP_PROXY -u HTTPS_PROXY -u ALL_PROXY \
  -u http_proxy -u https_proxy -u all_proxy \
  CUDA_VISIBLE_DEVICES="${TARGET_GPUS}" \
  HF_HUB_OFFLINE=1 \
  TRANSFORMERS_OFFLINE=1 \
  TOKENIZERS_PARALLELISM=false \
  "${TORCHRUN_BIN}" --standalone --nproc_per_node="${NUM_PROCESSES}" \
  scripts/preprocess_chartqa_teacher_gate.py \
  --input "${DEPLOT_VALIDATION_JSON}" \
  --output "${VALIDATION_JSONL}" \
  --split validation \
  --teacher-model "${TEACHER_PATH}" \
  --processor-model "${PROCESSOR_PATH}" \
  --providers visual_facts_deplot \
  --batch-size "${BATCH_SIZE}" \
  --max-new-tokens 96 \
  --repetition-penalty 1.0 \
  --require-real-deplot \
  --no-human-only \
  --expected-samples 1920 \
  --calibrate \
  --target-precision 0.90 \
  --min-coverage 0.10 \
  --log-every 100 \
  2>&1 | tee -a "${LOG_PATH}"

if [[ -s "${TRAIN_JSONL}" && -s "${TRAIN_JSONL}.manifest.json" ]]; then
  "${PYTHON_BIN}" scripts/reapply_chartqa_teacher_gate_calibration.py \
    --train-cache "${TRAIN_JSONL}" \
    --calibration-manifest "${VALIDATION_JSONL}.manifest.json" \
    2>&1 | tee -a "${LOG_PATH}"
else
  printf '[%s] train cache is not present yet; no train decision reapplication needed.\n' \
    "$(date -Is)" | tee -a "${LOG_PATH}"
fi