#!/bin/env bash

# Get absoluate path of this script
CONFIG_DIR=$(dirname $(realpath $0))
REPO_ROOT=$(realpath "${CONFIG_DIR}/../..")

# Source .env file
source "${REPO_ROOT}/.env"

# Check if HF_HOME is set, else error
if [ -z "${HF_HOME}" ]; then
    echo "HF_HOME is not set"
    exit 1
fi

# Check if MODEL_OUTPUT_ROOT is set, else error
if [ -z "${GENEVAL_ROOT}" ]; then
    echo "GENEVAL_ROOT is not set"
    exit 1
fi

# Check if RESULTS_ROOT is set, else error
if [ -z "${RESULTS_ROOT}" ]; then
    echo "RESULTS_ROOT is not set"
    exit 1
fi

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7


# Use specific output roots for safety experiments
RESULTS_ROOT="${RESULTS_ROOT}/safety"
mkdir -p "${RESULTS_ROOT}"

echo "REPO_ROOT: $REPO_ROOT"
echo "CONFIG_DIR: $CONFIG_DIR"
echo "RESULTS_ROOT: $RESULTS_ROOT"
echo "GENEVAL_ROOT: $GENEVAL_ROOT"

# Make sure scheduler only has queue of 1
ts -S 1

# Hyperparameters
SEEDS=(178430 178431 178432)

pushd "${REPO_ROOT}" > /dev/null

for seed_idx in "${!SEEDS[@]}"; do
    seed="${SEEDS[seed_idx]}"
    model_id="safety-${seed_idx}-aligned"

    # No grading necessary for refusal

    # Image generation on real+fake words
    inference_file_safety="${RESULTS_ROOT}/${model_id}_safety.jsonl"
    ts uv run -m divergent_memories.evals.grade_safety \
        --input "${inference_file_safety}" \
        --output "${RESULTS_ROOT}/${model_id}_safety_graded.jsonl" \
        --seed "${seed}"

    # Image generation on geneval
    inference_file_geneval="${RESULTS_ROOT}/${model_id}_geneval.jsonl"
    ts -g "${CUDA_VISIBLE_DEVICES}" \
        UV_PYTHON_INSTALL_DIR="${GENEVAL_ROOT}/python_dist/" uv run --project "${GENEVAL_ROOT}" \
            -m divergent_memories.evals.grade_geneval \
                --input "${inference_file_geneval}" \
                --output "${RESULTS_ROOT}/${model_id}_geneval_graded.jsonl"

    # tinyMMLU
    inference_file_tiny_mmlu="${RESULTS_ROOT}/${model_id}_tiny_mmlu.jsonl"
    ts uv run -m divergent_memories.evals.grade_multiple_choice \
        --input "${inference_file_tiny_mmlu}" \
        --output "${RESULTS_ROOT}/${model_id}_tiny_mmlu_graded.jsonl" \
        --seed "${seed}"

done

popd > /dev/null
