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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/env.sh"

cd "${SCRIPT_DIR}"

echo "[1/2] export aligned eval inputs"
"${SCRIPT_DIR}/06_export_eval_inputs.sh"

REFERENCE_SAMPLE_SIZE="${REFERENCE_SAMPLE_SIZE:-50}"
REFERENCE_SAMPLE_SEED="${REFERENCE_SAMPLE_SEED:-42}"
REFERENCE_INPUT_DIR="${REFERENCE_INPUT_DIR:-${EVAL_INPUT_DIR}/reference_sample_${REFERENCE_SAMPLE_SIZE}}"
REFERENCE_MODEL_OUT_ROOT="${REFERENCE_MODEL_OUT_ROOT:-${DATA_ROOT}/model_outputs/reference_sample_${REFERENCE_SAMPLE_SIZE}}"

if [[ "${REFERENCE_SAMPLE_SIZE}" =~ ^[0-9]+$ ]] && [[ "${REFERENCE_SAMPLE_SIZE}" -gt 0 ]]; then
  echo "[sample] selecting ${REFERENCE_SAMPLE_SIZE} aligned rows for reference inference"
  python3 "${QUERY_PRE}/sample_reference_inputs.py" \
    --e2e_input "${E2E_INPUT}" \
    --talker_input "${TALKER_INPUT_ALL}" \
    --response_input "${EVAL_INPUT_DIR}/response_tts_input_aligned.jsonl" \
    --out_dir "${REFERENCE_INPUT_DIR}" \
    --sample_size "${REFERENCE_SAMPLE_SIZE}" \
    --seed "${REFERENCE_SAMPLE_SEED}"
  REF_E2E_INPUT="${REFERENCE_INPUT_DIR}/e2e_input_aligned.jsonl"
  REF_TALKER_INPUT="${REFERENCE_INPUT_DIR}/talker_input_all.jsonl"
  REF_RESPONSE_INPUT="${REFERENCE_INPUT_DIR}/response_tts_input_aligned.jsonl"
else
  echo "[sample] REFERENCE_SAMPLE_SIZE=${REFERENCE_SAMPLE_SIZE}; using full aligned inputs"
  REF_E2E_INPUT="${E2E_INPUT}"
  REF_TALKER_INPUT="${TALKER_INPUT_ALL}"
  REF_RESPONSE_INPUT="${EVAL_INPUT_DIR}/response_tts_input_aligned.jsonl"
fi

echo "[2/2] submit reference jobs"
jid_talker_base=$(INPUT_JSONL="${REF_TALKER_INPUT}" OUT_DIR="${REFERENCE_MODEL_OUT_ROOT}/talker_baseline_from_constructed_thinker" sbatch --parsable "${SCRIPT_DIR}/run_talker_baseline_a800.sbatch")
jid_talker_sft=$(INPUT_JSONL="${REF_TALKER_INPUT}" OUT_DIR="${REFERENCE_MODEL_OUT_ROOT}/talker_sft_vstyle_46000_from_constructed_thinker" sbatch --parsable "${SCRIPT_DIR}/run_talker_sft_a800.sbatch")
jid_omni=$(INPUT_JSONL="${REF_E2E_INPUT}" OUT_DIR="${REFERENCE_MODEL_OUT_ROOT}/qwen3omni_e2e_from_query_audio" sbatch --parsable "${SCRIPT_DIR}/run_omni_e2e_a800.sbatch")
REF_RESPONSE_QWEN3_DIR="${REFERENCE_MODEL_OUT_ROOT}/response_audio/qwen3tts_response_serena_a800"
REF_RESPONSE_INDEXTTS_DIR="${REFERENCE_MODEL_OUT_ROOT}/response_audio/indextts_response_serena_a800_floatsave"
jid_resp_qwen=$(INPUT_JSONL="${REF_RESPONSE_INPUT}" OUT_HOME="${REF_RESPONSE_QWEN3_DIR}" EXTRA_ARGS="${REFERENCE_QWEN3TTS_EXTRA_ARGS:---speakers serena}" sbatch --parsable "${SCRIPT_DIR}/run_response_tts_qwen3_a800.sbatch")
jid_resp_idx=$(AUDIO_IN_DIR="${REF_RESPONSE_QWEN3_DIR}" OUT_HOME="${REF_RESPONSE_INDEXTTS_DIR}" sbatch --parsable --dependency=afterok:${jid_resp_qwen} "${SCRIPT_DIR}/run_response_tts_indextts_a800.sbatch")

cat > "${DATA_ROOT}/downstream_jobs.json" <<EOF
{
  "response_qwen3tts": "${jid_resp_qwen}",
  "response_indextts": "${jid_resp_idx}",
  "talker_baseline": "${jid_talker_base}",
  "talker_sft": "${jid_talker_sft}",
  "omni_e2e": "${jid_omni}",
  "reference_sample_size": "${REFERENCE_SAMPLE_SIZE}",
  "reference_sample_seed": "${REFERENCE_SAMPLE_SEED}",
  "reference_input_dir": "${REFERENCE_INPUT_DIR}",
  "reference_model_out_root": "${REFERENCE_MODEL_OUT_ROOT}",
  "reference_response_qwen3_dir": "${REF_RESPONSE_QWEN3_DIR}",
  "reference_response_indextts_dir": "${REF_RESPONSE_INDEXTTS_DIR}"
}
EOF

squeue -j "${jid_resp_qwen},${jid_resp_idx},${jid_talker_base},${jid_talker_sft},${jid_omni}" \
  -o '%.18i %.9P %.32j %.8T %.10M %.6D %R'