#!/usr/bin/env bash
# Retry prior hard failures with the exact same gold-blind DePlot-only prompt,
# then atomically merge them into a complete 4,576-row teacher-gate cache.
set -euo pipefail

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

PYTHON_ENV="${PYTHON_ENV:-${PROJECT_DIR}/.venv}"
TORCHRUN_BIN="${TORCHRUN_BIN:-${PYTHON_ENV}/bin/torchrun}"
INPUT_JSON="${INPUT_JSON:-${PROJECT_DIR}/data/chartqa/train_new_prerefine_vf_full_real_deplot_fp32_qwen25.json}"
BASE_CACHE="${BASE_CACHE:-${PROJECT_DIR}/data/chartqa/teacher_gate/llava_onevision_7b_deplot_greedy/train.jsonl}"
OUTPUT_JSONL="${OUTPUT_JSONL:-${PROJECT_DIR}/data/chartqa/teacher_gate/llava_onevision_7b_deplot_pass8_recovery/train.jsonl}"
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}"
CALIBRATION_MANIFEST="${CALIBRATION_MANIFEST:-${PROJECT_DIR}/data/chartqa/teacher_gate/llava_onevision_7b_deplot_greedy/validation.jsonl.manifest.json}"
TARGET_GPU="${TARGET_GPU:-5}"
NUM_CANDIDATES="${NUM_CANDIDATES:-8}"
BATCH_SIZE="${BATCH_SIZE:-4}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-96}"

for required_path in \
  "${TORCHRUN_BIN}" "${INPUT_JSON}" "${BASE_CACHE}" \
  "${BASE_CACHE}.manifest.json" "${TEACHER_PATH}" "${PROCESSOR_PATH}" \
  "${CALIBRATION_MANIFEST}"; do
  if [[ ! -e "${required_path}" ]]; then
    echo "Required DePlot recovery path does not exist: ${required_path}" >&2
    exit 2
  fi
done

mkdir -p "$(dirname "${OUTPUT_JSONL}")"
LOG_PATH="${OUTPUT_JSONL}.build.log"
printf '[%s] DePlot-only recovery GPU=%s candidates=%s output=%s\n' \
  "$(date -Is)" "${TARGET_GPU}" "${NUM_CANDIDATES}" "${OUTPUT_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_GPU}" \
  HF_HUB_OFFLINE=1 \
  TRANSFORMERS_OFFLINE=1 \
  TOKENIZERS_PARALLELISM=false \
  "${TORCHRUN_BIN}" --standalone --nproc_per_node=1 \
  scripts/preprocess_chartqa_teacher_gate.py \
  --input "${INPUT_JSON}" \
  --output "${OUTPUT_JSONL}" \
  --split train \
  --teacher-model "${TEACHER_PATH}" \
  --processor-model "${PROCESSOR_PATH}" \
  --providers visual_facts_deplot \
  --source-selection hard_fail \
  --selection-cache "${BASE_CACHE}" \
  --merge-selection-cache \
  --num-candidates "${NUM_CANDIDATES}" \
  --temperature 0.4 \
  --top-p 0.95 \
  --seed 42 \
  --batch-size "${BATCH_SIZE}" \
  --max-new-tokens "${MAX_NEW_TOKENS}" \
  --repetition-penalty 1.0 \
  --require-real-deplot \
  --expected-samples 4576 \
  --calibration-manifest "${CALIBRATION_MANIFEST}" \
  --log-every 32 \
  2>&1 | tee -a "${LOG_PATH}"