export OPENAI_API_KEY="your-openai-api-key-here"

#!/bin/bash
set -e  # Exit immediately if any command fails (ensures pipeline integrity)

# -------------------------- Configuration --------------------------
# Update these paths/parameters according to your project structure
DATASET="vg"  # Target dataset (matches your use case: "vg" or "oiv6")
CATE_INFO_PATH="./DATASET/VG150/vg_cate_dict.json"  # Path to entity-category dict
SUPER_ENTITIES_PATH="./DATASET/VG150/vg_super_entities.json"  # Output of Step 1
ENTITY_SUPERCLASS_SAVE_PATH="./DATASET/VG150/vg_entity_superclass_final.json"  # Output of Step 2
REL_PROMPT_OUTPUT_PREFIX="./DATASET/VG150/vg_relation_aware_prompt_"  # Prefix for Step 3 outputs
FINAL_PROMPT_SAVE_PATH="./DATASET/VG150/vg_relation_aware_prompts.json"  # Final output of Step 4

# Clustering & API parameters (adjust if needed)
DISTANCE_THRESHOLD=0.5
LINKAGE_METHOD="ward"
MAX_WORKERS=30
MODEL_NAME="gpt-4o-mini"
RETRY_DELAY=10
# -------------------------------------------------------------------


# -------------------------- Step 1: Cluster Entities into Superclasses --------------------------
echo -e "\n=== Starting Step 1: Cluster Entities into Superclasses ==="
cd tools  # Navigate to the "tools" directory (where your .py scripts are stored)

python cluster_entity_2_super_class.py \
  --dataset "$DATASET" \
  --cate-info-path "$CATE_INFO_PATH" \
  --save-path "$SUPER_ENTITIES_PATH" \
  --distance-threshold "$DISTANCE_THRESHOLD" \
  --linkage-method "$LINKAGE_METHOD"

if [ -f "$SUPER_ENTITIES_PATH" ]; then
  echo "✅ Step 1 Completed: Entity clusters saved to $SUPER_ENTITIES_PATH"
else
  echo "❌ Step 1 Failed: Cluster file not generated"
  exit 1
fi


# -------------------------- Step 2: Validate Superclass Clustering --------------------------
echo -e "\n=== Starting Step 2: Validate Superclass Clustering ==="
# Use the environment variable for OpenAI API key (no hardcoding)
if [ -z "$OPENAI_API_KEY" ]; then
  echo "❌ Error: OPENAI_API_KEY environment variable not set. Set it first (see Prerequisite section)."
  exit 1
fi

python check_super_entity_class.py \
  --openai-api-key "$OPENAI_API_KEY" \
  --super-entities-path "$SUPER_ENTITIES_PATH" \
  --cate-info-path "$CATE_INFO_PATH" \
  --save-path "$ENTITY_SUPERCLASS_SAVE_PATH" \
  --model-name "$MODEL_NAME" \
  --retry-delay "$RETRY_DELAY"

if [ -f "$ENTITY_SUPERCLASS_SAVE_PATH" ]; then
  echo "✅ Step 2 Completed: Validated superclasses saved to $ENTITY_SUPERCLASS_SAVE_PATH"
else
  echo "❌ Step 2 Failed: Validated superclass file not generated"
  exit 1
fi


# -------------------------- Step 3: Generate Relation-Aware Prompts --------------------------
echo -e "\n=== Starting Step 3: Generate Relation-Aware Prompts ==="
python relation_aware_prompt_generation.py \
  --openai-api-key "$OPENAI_API_KEY" \
  --dataset "$DATASET" \
  --output-prefix "$REL_PROMPT_OUTPUT_PREFIX" \
  --max-workers "$MAX_WORKERS" \
  --model-name "$MODEL_NAME" \
  --save-combined

# Verify Step 3 output (check if at least one worker file exists)
FIRST_WORKER_FILE="${REL_PROMPT_OUTPUT_PREFIX}worker_0.jsonl"
if [ -f "$FIRST_WORKER_FILE" ]; then
  echo "✅ Step 3 Completed: Relation-aware prompts saved to ${REL_PROMPT_OUTPUT_PREFIX}worker_*.jsonl"
else
  echo "❌ Step 3 Failed: No prompt files generated"
  exit 1
fi


# -------------------------- Step 4: Convert Prompts to Final JSON --------------------------
echo -e "\n=== Starting Step 4: Convert Prompts to Final JSON ==="
python convert_relation_aware_prompt.py \
  --dataset "$DATASET" \
  --input_prefix "$REL_PROMPT_OUTPUT_PREFIX" \
  --num_workers "$MAX_WORKERS" \
  --output_path "$FINAL_PROMPT_SAVE_PATH"

if [ -f "$FINAL_PROMPT_SAVE_PATH" ]; then
  echo "✅ Step 4 Completed: Final prompts saved to $FINAL_PROMPT_SAVE_PATH"
else
  echo "❌ Step 4 Failed: Final JSON file not generated"
  exit 1
fi


# -------------------------- Pipeline Completion --------------------------
echo -e "\n🎉 All Steps Completed Successfully! Final output: $FINAL_PROMPT_SAVE_PATH"
cd ..  # Return to the parent directory (optional)