#!/usr/bin/env bash
# =============================================================================
# Shared setup for all baseline scripts.
# Source this file, don't run it directly.
# =============================================================================

set -euo pipefail

# Resolve project root (three levels above this script: baselines/ -> scripts/ -> experiment/ -> root)
PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"

export PYTHONPATH="${PROJECT_ROOT}:${PROJECT_ROOT}/EasyEdit:${PYTHONPATH:-}"

# Cache dirs
export TORCHINDUCTOR_CACHE_DIR="${HOME}/scratch/.cache/torchinductor"
export TRITON_CACHE_DIR="${HOME}/scratch/.cache/triton"

# Load CUDA module if on cluster
module load cuda/12.6.2 2>/dev/null || true
export CUDA_HOME="${CUDA_HOME:-$(dirname $(dirname $(which nvcc 2>/dev/null) 2>/dev/null) 2>/dev/null)}"

# Normalize UUID-style CUDA_VISIBLE_DEVICES to numeric indices
if [ -n "${CUDA_VISIBLE_DEVICES:-}" ] && [[ "${CUDA_VISIBLE_DEVICES}" == *GPU-* ]]; then
    declare -A _uuid_to_index=()
    while IFS=, read -r idx uuid; do
        idx="$(echo "${idx}" | xargs)"
        uuid="$(echo "${uuid}" | xargs)"
        [ -n "${idx}" ] && [ -n "${uuid}" ] && _uuid_to_index["${uuid}"]="${idx}"
    done < <(nvidia-smi --query-gpu=index,uuid --format=csv,noheader)

    IFS=',' read -r -a _requested <<< "${CUDA_VISIBLE_DEVICES}"
    _mapped=()
    _ok=1
    for raw in "${_requested[@]}"; do
        uuid="$(echo "${raw}" | xargs)"
        if [ -n "${_uuid_to_index[${uuid}]:-}" ]; then
            _mapped+=("${_uuid_to_index[${uuid}]}")
        else
            _ok=0; break
        fi
    done

    if [ "${_ok}" -eq 1 ] && [ "${#_mapped[@]}" -gt 0 ]; then
        CUDA_VISIBLE_DEVICES="$(IFS=,; echo "${_mapped[*]}")"
        export CUDA_VISIBLE_DEVICES
        echo "Normalized CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES}"
    fi
fi

# =============================================================================
# Common paths — override via env
# =============================================================================
# Relation (drives dataset, keywords, prompts, categories)
RELATION="${RELATION:-bathroom_toilet}"
# HuggingFace dataset (default: auto-resolved from RELATION)
DATASET_ID="${DATASET_ID:-}"
# Legacy CSV/image paths (set both to use local files instead of HuggingFace)
CSV_PATH="${CSV_PATH:-}"
IMAGE_DIR="${IMAGE_DIR:-}"
BASE_MODEL="${BASE_MODEL:-llava-hf/llava-1.5-7b-hf}"
HPARAMS_DIR="${HPARAMS_DIR:-${PROJECT_ROOT}/experiment/knowledge_editing/hparams}"
EDIT_SET="${EDIT_SET:-${PROJECT_ROOT}/experiment/data/edit_set_${RELATION}.json}"
CAPTION_TARGETS="${CAPTION_TARGETS:-${PROJECT_ROOT}/experiment/data/caption_targets_${RELATION}.json}"
DEVICE="${DEVICE:-cuda}"
NUM_PER_CATEGORY="${NUM_PER_CATEGORY:-50}"
MENTION_METHOD="${MENTION_METHOD:-keyword}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-300}"
N_EDITS="${N_EDITS:-20}"

# =============================================================================
# Helper: run evaluation on an edited model
# =============================================================================
run_eval() {
    local model_type="$1"
    local model_dir="$2"
    local eval_output="$3"
    local method_name="$4"
    shift 4
    local extra_args=("$@")

    echo ""
    echo ">>> Evaluating: ${method_name}"
    echo "  Model type: ${model_type}"
    echo "  Model dir:  ${model_dir}"
    echo "  Output:     ${eval_output}"

    local val_args=()
    val_args+=(--relation "${RELATION}")
    if [ -n "${CSV_PATH}" ] && [ -n "${IMAGE_DIR}" ]; then
        val_args+=(--val_csv "${CSV_PATH}" --val_image_dir "${IMAGE_DIR}")
    elif [ -n "${DATASET_ID}" ]; then
        val_args+=(--dataset_id "${DATASET_ID}")
    fi

    python -m experiment.evaluation.validate \
        --model_type "${model_type}" \
        --model_dir "${model_dir}" \
        --base_model_name "${BASE_MODEL}" \
        "${val_args[@]}" \
        --num_per_category "${NUM_PER_CATEGORY}" \
        --mention_method "${MENTION_METHOD}" \
        --max_new_tokens "${MAX_NEW_TOKENS}" \
        --prompts "Describe this image." \
        --train_prompts "Describe this image." \
        --generality_prompts "What do you see in this image?" \
        --use_val_split \
        --output_dir "${eval_output}" \
        "${extra_args[@]}" \
        || { echo "  WARNING: eval failed for ${method_name}"; return 1; }
}

# =============================================================================
# Helper: ensure edit set exists with caption targets
# =============================================================================
ensure_edit_set() {
    if [ ! -f "$EDIT_SET" ]; then
        echo ">>> Building edit set..."

        local data_args=(--relation "${RELATION}")
        if [ -n "${CSV_PATH}" ] && [ -n "${IMAGE_DIR}" ]; then
            data_args+=(--csv "$CSV_PATH" --image_dir "$IMAGE_DIR")
        elif [ -n "${DATASET_ID}" ]; then
            data_args+=(--dataset_id "${DATASET_ID}")
        fi

        # Prefer pre-built caption targets (LLM-cleaned) if available
        if [ -f "$CAPTION_TARGETS" ]; then
            echo "  Using pre-built caption targets: ${CAPTION_TARGETS}"
            python -m experiment.knowledge_editing.build_edit_set \
                "${data_args[@]}" \
                --output "$EDIT_SET" \
                --max_locality_per_category "$NUM_PER_CATEGORY" \
                --caption_targets "$CAPTION_TARGETS"
        else
            echo "  No caption_targets.json found, building structure only."
            echo "  Run 'experiment/scripts/data/run_build_caption_targets.sh' first for LLM-cleaned targets."
            python -m experiment.knowledge_editing.build_edit_set \
                "${data_args[@]}" \
                --output "$EDIT_SET" \
                --max_locality_per_category "$NUM_PER_CATEGORY"
        fi
    fi

    TARGETS_FILLED=$(python -c "
import json
with open('${EDIT_SET}') as f:
    d = json.load(f)
n = d['stats'].get('n_with_targets', 0)
print(n)
" 2>/dev/null || echo "0")

    if [ "$TARGETS_FILLED" -eq "0" ]; then
        # Fallback: fill targets inline (legacy regex cleaning)
        echo ">>> WARNING: No targets in edit set. Falling back to inline generation (regex cleaning)."
        echo "  For better results, run build_caption_targets.py first."
        python -m experiment.knowledge_editing.build_edit_set \
            --fill_targets "$EDIT_SET" \
            --model "$BASE_MODEL" \
            --device "$DEVICE"
    else
        echo ">>> Edit set ready ($TARGETS_FILLED instances with targets)."
    fi
}