#!/bin/env bash

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

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

echo "REPO_ROOT: $REPO_ROOT"
echo "CONFIG_DIR: $CONFIG_DIR"

NUM_GPUS=2
export CUDA_VISIBLE_DEVICES=6,7

# Load main .env file for MODEL_OUTPUT_ROOT
if [ -f "${REPO_ROOT}/.env" ]; then
    echo "Loading .env file"
    source "${REPO_ROOT}/.env"
else
    echo ".env file not found!"
    exit 1
fi

NUM_GPUS=4
export CUDA_VISIBLE_DEVICES=0,1,2,3

# Env vars
export OMP_NUM_THREADS=1
export TRITON_CACHE_DIR="/local/home/${USER}/.triton/"
mkdir -p "${TRITON_CACHE_DIR}"

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

for idx in {0..2}; do

    uv run -m divergent_memories.evals.grade_synthetic_images \
        --input "${RESULTS_ROOT}/concepts/harmon-${idx}_synthetic_concepts_hd.jsonl" \
        --output "${RESULTS_ROOT}/concepts/harmon-${idx}_synthetic_concepts_hd_graded.jsonl"

    UV_PYTHON_INSTALL_DIR="${GENEVAL_ROOT}/python_dist/" uv run --project "${GENEVAL_ROOT}" \
    -m divergent_memories.evals.grade_geneval \
        --input "${RESULTS_ROOT}/concepts/harmon-${idx}_geneval.jsonl" \
        --output "${RESULTS_ROOT}/concepts/harmon-${idx}_geneval_graded.jsonl"

    uv run -m divergent_memories.evals.grade_multiple_choice \
        --input "${RESULTS_ROOT}/concepts/harmon-${idx}_concepts_description_mc.jsonl" \
        --output "${RESULTS_ROOT}/concepts/harmon-${idx}_concepts_description_mc_graded.jsonl"

    uv run -m divergent_memories.evals.grade_multiple_choice \
        --input "${RESULTS_ROOT}/concepts/harmon-${idx}_concepts_description_ablation_mc.jsonl" \
        --output "${RESULTS_ROOT}/concepts/harmon-${idx}_concepts_description_ablation_mc_graded.jsonl"

    uv run -m divergent_memories.evals.grade_multiple_choice \
        --input "${RESULTS_ROOT}/concepts/harmon-${idx}_tiny_mmlu.jsonl" \
        --output "${RESULTS_ROOT}/concepts/harmon-${idx}_tiny_mmlu_graded.jsonl"

done