#!/bin/bash
set -euo pipefail

SOURCE_CHECKPOINT="${SOURCE_CHECKPOINT:?must be set}"
REFERENCE_PARQUET="${REFERENCE_PARQUET:?must be set}"
REPO_URL="${REPO_URL:-https://github.com/vedatonuryilmaz/dev_genopix.git}"
REPO_BRANCH="${REPO_BRANCH:-ml-intern}"
ARTIFACT_REPO="${ARTIFACT_REPO:?must be set}"
OBJECTIVE="${OBJECTIVE:-contrastive}"
EPOCHS="${EPOCHS:-10}"
BATCH_SIZE="${BATCH_SIZE:-128}"
NUM_WORKERS="${NUM_WORKERS:-4}"
DEVICE="${DEVICE:-cpu}"

git clone --branch "$REPO_BRANCH" --depth 1 "$REPO_URL" deepgenopix
cd deepgenopix
uv sync --extra hf --extra quant

OUTPUT_NAME="${1:-$(uv run python scripts/run_names.py quant-train --objective "$OBJECTIVE" --source-checkpoint "$SOURCE_CHECKPOINT" --reference-parquet "$REFERENCE_PARQUET")}"

if [ "${RESUME:-0}" = "1" ]; then
  uv run python scripts/hf_artifacts.py pull --repo "$ARTIFACT_REPO" --preset "$OUTPUT_NAME" --dest data/output
fi

MODE=$([ "${RESUME:-0}" = "1" ] && echo resume || echo train)
uv run deepgenopix quant-train \
  --objective "$OBJECTIVE" \
  --source-checkpoint "$SOURCE_CHECKPOINT" \
  --reference-parquet "$REFERENCE_PARQUET" \
  --output-dir "data/output/$OUTPUT_NAME" \
  --epochs "$EPOCHS" \
  --batch-size "$BATCH_SIZE" \
  --num-workers "$NUM_WORKERS" \
  --device "$DEVICE" \
  --run-mode "$MODE" \
  --json

uv run python scripts/hf_artifacts.py push --repo "$ARTIFACT_REPO" --preset "$OUTPUT_NAME" --source data/output