#!/usr/bin/env bash
# =============================================================================
# Run LoRA Fine-Tune to Remove Toilet Hallucination
# =============================================================================
# Usage (single GPU):
#   bash experiment/scripts/run_finetune_lora.sh
#
# Usage (multi-GPU DDP):
#   NGPUS=8 bash experiment/scripts/run_finetune_lora.sh
# =============================================================================

set -euo pipefail

# Resolve project root (parent of experiment/)
PROJECT_ROOT="$(cd "$(dirname "$0")/../.." && pwd)"

# Ensure project is importable
export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}"

# Redirect torch inductor / Triton cache to a writable location
export TORCHINDUCTOR_CACHE_DIR="${HOME}/scratch/.cache/torchinductor"
export TRITON_CACHE_DIR="${HOME}/scratch/.cache/triton"

# =============================================================================
# Configuration — edit these or set as environment variables
# =============================================================================
CSV_PATH="${CSV_PATH:-/home/users/ntu/cong045/scratch/testing/hallucination/CC3M-Dataset/data-qwenvl32b/bathroom_toilet_labels.csv}"
IMAGE_DIR="${IMAGE_DIR:-CC3M-Dataset/cc3m_images/train}"
OUTPUT_DIR="${OUTPUT_DIR:-./step3_lora_outputs}"
NGPUS="${NGPUS:-1}"

# Config file (will be auto-generated if it doesn't exist)
FINETUNE_CONFIG="${FINETUNE_CONFIG:-${PROJECT_ROOT}/experiment/lora_config.json}"

echo "=========================================="
echo "LoRA Fine-Tune"
echo "=========================================="
echo "Config:"
echo "  CSV path:        $CSV_PATH"
echo "  Image dir:       $IMAGE_DIR"
echo "  Output dir:      $OUTPUT_DIR"
echo "  Finetune config: $FINETUNE_CONFIG"
echo "  GPUs:            $NGPUS"
echo "=========================================="

# Generate default config if it doesn't exist
if [ ! -f "$FINETUNE_CONFIG" ]; then
    echo "Generating default LoRA config at $FINETUNE_CONFIG ..."
    python -c "
import sys; sys.path.insert(0, '${PROJECT_ROOT}')
from experiment.config.train_config import TrainConfig
cfg = TrainConfig(
    csv_path='${CSV_PATH}',
    image_dir='${IMAGE_DIR}',
    output_dir='${OUTPUT_DIR}',
)
cfg.save('${FINETUNE_CONFIG}')
print('  Config saved.')
"
fi

# ========================================
# LoRA Fine-tune
# ========================================
echo ""
echo ">>> Running LoRA Fine-Tune..."
echo "================================"

if [ "$NGPUS" -gt 1 ]; then
    torchrun --nproc_per_node="$NGPUS" \
        -m experiment.training.finetune_lora \
        --config "$FINETUNE_CONFIG"
else
    python -m experiment.training.finetune_lora \
        --config "$FINETUNE_CONFIG"
fi

# Find the latest run directory
LATEST_RUN=$(ls -dt "$OUTPUT_DIR"/run_* 2>/dev/null | head -1)

if [ -z "$LATEST_RUN" ]; then
    echo "ERROR: Fine-tuning failed — no run_* directory found in $OUTPUT_DIR"
    exit 1
fi

echo ""
echo ">>> LoRA Fine-Tuning Complete!"
echo "  Run dir: $LATEST_RUN"

# ========================================
# Validate
# ========================================
EVAL_OUTPUT_DIR="${EVAL_OUTPUT_DIR:-./step4_lora_outputs}"
NUM_PER_CATEGORY="${NUM_PER_CATEGORY:-50}"

echo ""
echo ">>> Running Validation..."
echo "================================"

python -m experiment.evaluation.validate \
    --model_type lora \
    --model_dir "$LATEST_RUN/lora_adapter" \
    --val_csv "$CSV_PATH" \
    --val_image_dir "$IMAGE_DIR" \
    --num_per_category "$NUM_PER_CATEGORY" \
    --output_dir "$EVAL_OUTPUT_DIR"

echo ""
echo "=========================================="
echo "All Steps Complete!"
echo "=========================================="
echo "Outputs:"
echo "  Run dir:       $LATEST_RUN/"
echo "  LoRA adapter:  $LATEST_RUN/lora_adapter/"
echo "  Merged model:  $LATEST_RUN/merged_model/"
echo "  Validation:    $EVAL_OUTPUT_DIR/"
echo "=========================================="