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

campaign_tag="${1:?usage: scripts/spark_campaign.sh CAMPAIGN_TAG [CONFIG]}"
config_path="${2:-configs/main.yaml}"
project_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$project_dir"

python_bin="${BGC_PYTHON:?set BGC_PYTHON to a CUDA-enabled Python executable}"
dependency_dir="${BGC_EXTRA_PYTHONPATH:-$project_dir/.deps}"
export PYTHONPATH="$dependency_dir:$project_dir/src${PYTHONPATH:+:$PYTHONPATH}"

run_bgc() {
  "$python_bin" -c 'from bgc_retrieval.cli import main; main()' "$@"
}

run_step() {
  local name="$1"
  shift
  local started=$SECONDS
  printf '[%s] START %s\n' "$(date --iso-8601=seconds)" "$name"
  "$@"
  printf '[%s] DONE  %s (%ss)\n' "$(date --iso-8601=seconds)" "$name" "$((SECONDS - started))"
}

"$python_bin" -c 'import torch; assert torch.cuda.is_available(), "CUDA GPU is required"; print(torch.__version__, torch.cuda.get_device_name(0))'
printf 'CONFIG %s\n' "$config_path"

if [[ ! -f data/manifests/legacy_manifest.json ]]; then
  run_step audit run_bgc audit --config "$config_path"
fi
if [[ ! -f data/manifests/silver_split.csv ]]; then
  run_step splits run_bgc build-splits --config "$config_path"
fi
if [[ ! -f data/external/processed/external_atlas.csv ]]; then
  run_step external-preparation "$python_bin" scripts/prepare_external_benchmark.py
fi
if [[ ! -f data/external/processed/external_esm2.h5 ]]; then
  run_step external-esm2 "$python_bin" scripts/embed_external_proteins.py --resume
fi

run_step phase1 run_bgc train --config "$config_path" --stage phase1 \
  --run-id "${campaign_tag}-phase1"
run_step phase2-main run_bgc train --config "$config_path" --stage phase2 \
  --run-id "${campaign_tag}-main" \
  --phase1-checkpoint "artifacts/${campaign_tag}-phase1/phase1_best.pt"
run_step phase2-no-phase1 run_bgc train --config "$config_path" --stage phase2 \
  --run-id "${campaign_tag}-no-phase1"
run_step internal-main run_bgc evaluate --config "$config_path" \
  --checkpoint "artifacts/${campaign_tag}-main/phase2_best.pt" \
  --run-id "${campaign_tag}-main-evaluation"
run_step internal-no-phase1 run_bgc evaluate --config "$config_path" \
  --checkpoint "artifacts/${campaign_tag}-no-phase1/phase2_best.pt" \
  --run-id "${campaign_tag}-no-phase1-evaluation"
run_step external-main "$python_bin" scripts/evaluate_external.py --config "$config_path" \
  --checkpoint "artifacts/${campaign_tag}-main/phase2_best.pt" \
  --run-id "${campaign_tag}-main-external" \
  --bootstrap-samples "${EXTERNAL_BOOTSTRAP_SAMPLES:-1000}"
run_step external-no-phase1 "$python_bin" scripts/evaluate_external.py --config "$config_path" \
  --checkpoint "artifacts/${campaign_tag}-no-phase1/phase2_best.pt" \
  --run-id "${campaign_tag}-no-phase1-external" \
  --bootstrap-samples "${EXTERNAL_BOOTSTRAP_SAMPLES:-1000}"

printf '[%s] CAMPAIGN COMPLETE: %s\n' "$(date --iso-8601=seconds)" "$campaign_tag"