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

campaign_tag="${1:?usage: scripts/dgx_campaign.sh CAMPAIGN_TAG}"
project_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$project_dir"

python -m venv --system-site-packages .venv
.venv/bin/pip install --upgrade pip
.venv/bin/pip install -e '.[dev]'
.venv/bin/python -c 'import torch; assert torch.cuda.is_available(), "CUDA GPU is required"; print(torch.cuda.get_device_name(0))'

if [[ ! -f data/manifests/legacy_manifest.json ]]; then
  .venv/bin/bgc-rebuild audit --config configs/main.yaml
fi
if [[ ! -f data/manifests/silver_split.csv ]]; then
  .venv/bin/bgc-rebuild build-splits --config configs/main.yaml
fi
if [[ ! -f data/external/processed/external_atlas.csv ]]; then
  .venv/bin/python scripts/prepare_external_benchmark.py
fi

.venv/bin/bgc-rebuild train --config configs/main.yaml --stage phase1 \
  --run-id "${campaign_tag}-phase1"
.venv/bin/bgc-rebuild train --config configs/main.yaml --stage phase2 \
  --run-id "${campaign_tag}-main" \
  --phase1-checkpoint "artifacts/${campaign_tag}-phase1/phase1_best.pt"
.venv/bin/bgc-rebuild train --config configs/main.yaml --stage phase2 \
  --run-id "${campaign_tag}-no-phase1"
.venv/bin/bgc-rebuild evaluate --config configs/main.yaml \
  --checkpoint "artifacts/${campaign_tag}-main/phase2_best.pt" \
  --run-id "${campaign_tag}-main-evaluation"
.venv/bin/bgc-rebuild evaluate --config configs/main.yaml \
  --checkpoint "artifacts/${campaign_tag}-no-phase1/phase2_best.pt" \
  --run-id "${campaign_tag}-no-phase1-evaluation"

if [[ ! -f data/external/processed/external_esm2.h5 ]]; then
  .venv/bin/python scripts/embed_external_proteins.py --resume
fi
.venv/bin/python scripts/evaluate_external.py --config configs/main.yaml \
  --checkpoint "artifacts/${campaign_tag}-main/phase2_best.pt" \
  --run-id "${campaign_tag}-main-external" \
  --bootstrap-samples "${EXTERNAL_BOOTSTRAP_SAMPLES:-1000}"
.venv/bin/python scripts/evaluate_external.py --config configs/main.yaml \
  --checkpoint "artifacts/${campaign_tag}-no-phase1/phase2_best.pt" \
  --run-id "${campaign_tag}-no-phase1-external" \
  --bootstrap-samples "${EXTERNAL_BOOTSTRAP_SAMPLES:-1000}"