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

# Convert the full CSV dataset into the JSONL format required by pretraining.
# Run this before scripts/slurm/run_pretrain.sh.
# This script does not request SLURM resources.

PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"
CONDA_ENV="${CONDA_ENV:-struct-evo}"

INPUT_CSV="${INPUT_CSV:-${PROJECT_DIR}/scripts/data/raw/dataset.csv}"
OUTPUT_JSON="${OUTPUT_JSON:-${PROJECT_DIR}/scripts/data/processed/pretrain_data.json}"

cd "${PROJECT_DIR}"
mkdir -p "$(dirname "${OUTPUT_JSON}")"

export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
export INPUT_CSV
export OUTPUT_JSON
export PYTHONFAULTHANDLER=1

if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then
    # shellcheck disable=SC1091
    source "$(conda info --base)/etc/profile.d/conda.sh"
    conda activate "${CONDA_ENV}"
fi

if [[ ! -f "${INPUT_CSV}" ]]; then
    echo "Missing INPUT_CSV: ${INPUT_CSV}" >&2
    exit 1
fi

python - <<'PY'
import os

from CodonTransformer.CodonData import prepare_training_data

input_csv = os.environ["INPUT_CSV"]
output_json = os.environ["OUTPUT_JSON"]

print(f"Input CSV: {input_csv}")
print(f"Output JSONL: {output_json}")

prepare_training_data(input_csv, output_json)
PY