#!/bin/bash
set -euo pipefail

CORPUS_NAME="${1:?usage: $0 <corpus-name>}"
REPO_URL="${REPO_URL:-https://github.com/vedatonuryilmaz/dev_genopix.git}"
REPO_BRANCH="${REPO_BRANCH:-ml-intern}"
DATA_DATASET="${DATA_DATASET:-vedatonuryilmaz/te-variant-simulator-v1}"
REFERENCE_DATASET="${REFERENCE_DATASET:-vedatonuryilmaz/te-seqdata-v1}"
REFERENCE_REVISION="${REFERENCE_REVISION:-f107d46c5a61087eb3fb19b4e3e75fdef0b74fdd}"
REFERENCE_FILE="${REFERENCE_FILE:-train/te_seqdata.parquet}"
REFERENCE_PARQUET="${REFERENCE_PARQUET:-data/raw/reference/${REFERENCE_FILE}}"
CORPUS_KIND="${CORPUS_KIND:-controlled}"
N_LOCI="${N_LOCI:-64}"
COVERAGE="${COVERAGE:-50}"
READ_LENGTH="${READ_LENGTH:-150}"
ERROR_RATE="${ERROR_RATE:-0.001}"
BP_PER_TOKEN="${BP_PER_TOKEN:-24}"
VARIANT_TYPES="${VARIANT_TYPES:-snp}"

if [ "${DEEPGENOPIX_ALREADY_CLONED:-0}" = "1" ]; then
  echo "using existing DeepGenopix checkout at $(pwd)"
else
  git clone --branch "$REPO_BRANCH" --depth 1 "$REPO_URL" deepgenopix
  cd deepgenopix
fi
uv sync --extra hf --extra quant

if [ ! -f "$REFERENCE_PARQUET" ]; then
  mkdir -p "$(dirname "$REFERENCE_PARQUET")"
  uv run python -c "
from huggingface_hub import hf_hub_download
hf_hub_download(
    repo_id='$REFERENCE_DATASET',
    repo_type='dataset',
    revision='$REFERENCE_REVISION',
    filename='$REFERENCE_FILE',
    local_dir='data/raw/reference',
)
"
fi

OUTPUT_DIR="data/output/simulator/$CORPUS_KIND/$CORPUS_NAME"
mkdir -p "$OUTPUT_DIR"

uv run python -c "
from pathlib import Path
from deepgenopix.variant.simulate import SimulatorConfig, materialize_corpus

variant_types = tuple(item.strip() for item in '$VARIANT_TYPES'.split(',') if item.strip())
cfg = SimulatorConfig(
    coverage=int('$COVERAGE'),
    read_length=int('$READ_LENGTH'),
    error_rate=float('$ERROR_RATE'),
    bp_per_token=int('$BP_PER_TOKEN'),
    variant_types=variant_types,
)
paths = materialize_corpus(
    Path('$REFERENCE_PARQUET'),
    cfg,
    Path('$OUTPUT_DIR'),
    n_loci=int('$N_LOCI'),
)
print(paths)
"

uv run python scripts/hf_artifacts.py push-dataset-folder \
  --repo "$DATA_DATASET" \
  --source "$OUTPUT_DIR" \
  --path-in-repo "$CORPUS_KIND/$CORPUS_NAME" \
  --message "variant simulator corpus $CORPUS_KIND/$CORPUS_NAME"