#!/bin/bash
set -euo pipefail

ENCODER="${ENCODER:?must be set}"
REFERENCE_PARQUET="${REFERENCE_PARQUET:?must be set}"
OUTPUT_NAME="${1:?usage: $0 <output-name>}"
REPO_URL="${REPO_URL:-https://github.com/vedatonuryilmaz/dev_genopix.git}"
REPO_BRANCH="${REPO_BRANCH:-ml-intern}"
DATA_DATASET="${DATA_DATASET:-vedatonuryilmaz/te-null-stats-v1}"
ESTIMATOR="${ESTIMATOR:-walk_redundancy}"
N_LOCI="${N_LOCI:-128}"
READ_LENGTH="${READ_LENGTH:-150}"
COVERAGE="${COVERAGE:-30}"
ERROR_RATE="${ERROR_RATE:-0.001}"
BP_PER_PIXEL="${BP_PER_PIXEL:-12}"
PIXEL_STRIDE_BP="${PIXEL_STRIDE_BP:-12}"
DEVICE="${DEVICE:-cpu}"

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

OUTPUT_DIR="data/output/null_stats/$OUTPUT_NAME"
mkdir -p "$OUTPUT_DIR"

uv run python -c "
import json
import pandas as pd
from pathlib import Path

from deepgenopix.pixelizer import DNAPixelizer
from deepgenopix.quant import load_quant_encoder
from deepgenopix.variant.null import reference_split_null, walk_redundancy_null

encoder = load_quant_encoder('$ENCODER', device='$DEVICE')
df = pd.read_parquet('$REFERENCE_PARQUET').head(int('$N_LOCI'))
references = [str(seq).upper() for seq in df['sequence'].tolist()]
if int('$BP_PER_PIXEL') != DNAPixelizer.BP_PER_PIXEL:
    raise ValueError('DNAPixelizer currently supports BP_PER_PIXEL=12 only')
pixelizer = DNAPixelizer(pixel_stride_bp=int('$PIXEL_STRIDE_BP'))
if '$ESTIMATOR' == 'reference_split':
    stats = reference_split_null(
        encoder,
        references,
        pixelizer=pixelizer,
        coverage=int('$COVERAGE'),
        read_length=int('$READ_LENGTH'),
        error_rate=float('$ERROR_RATE'),
        device='$DEVICE',
    )
else:
    stats = walk_redundancy_null(encoder, references, pixelizer=pixelizer, device='$DEVICE')
out = Path('$OUTPUT_DIR')
stats.save(out / 'null.pt')
(out / 'summary.json').write_text(json.dumps({
    'encoder': '$ENCODER',
    'encoder_hash': stats.encoder_hash,
    'estimator': stats.estimator,
    'n_loci': len(references),
    'output_name': '$OUTPUT_NAME',
}, indent=2))
print(json.dumps({'null_stats': str(out / 'null.pt'), 'encoder_hash': stats.encoder_hash}))
"

ENCODER_HASH="$(uv run python -c "from deepgenopix.variant.null import NullStats; print(NullStats.load('$OUTPUT_DIR/null.pt').encoder_hash)")"
uv run python scripts/hf_artifacts.py push-dataset-folder \
  --repo "$DATA_DATASET" \
  --source "$OUTPUT_DIR" \
  --path-in-repo "null_stats/$ENCODER_HASH/$OUTPUT_NAME" \
  --message "null stats $OUTPUT_NAME"