#!/usr/bin/env bash
# =============================================================================
# Run Fine-Tune, Validate, and Inspect Delta-W
# =============================================================================
# This script runs:
#   1. Fine-tune model to remove toilet hallucination
#   2. Validate across 4 image categories
#   3. SVD analysis of delta-W
#
# Usage:
#   bash experiment/scripts/run_finetune.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:-CC3M-Dataset/bathroom_filter/bathroom_toilet_labels.csv}"
IMAGE_DIR="${IMAGE_DIR:-CC3M-Dataset/cc3m_images/train}"

STEP3_OUTPUT_DIR="${STEP3_OUTPUT_DIR:-./step3_outputs}"
STEP4_OUTPUT_DIR="${STEP4_OUTPUT_DIR:-./step4_outputs}"
STEP5_OUTPUT_DIR="${STEP5_OUTPUT_DIR:-./step5_outputs}"

NUM_PER_CATEGORY="${NUM_PER_CATEGORY:-50}"
TOP_K_SINGULAR="${TOP_K_SINGULAR:-10}"

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

echo "=========================================="
echo "Fine-Tune, Validate, Inspect"
echo "=========================================="
echo "Config:"
echo "  CSV path:        $CSV_PATH"
echo "  Image dir:       $IMAGE_DIR"
echo "  Finetune config: $FINETUNE_CONFIG"
echo "=========================================="

# Generate default config if it doesn't exist
if [ ! -f "$FINETUNE_CONFIG" ]; then
    echo "Generating default 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='${STEP3_OUTPUT_DIR}',
)
cfg.save('${FINETUNE_CONFIG}')
print('  Config saved.')
"
fi

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

python -m experiment.training.finetune \
    --config "$FINETUNE_CONFIG"

# Verify fine-tuning succeeded
if [ ! -f "$STEP3_OUTPUT_DIR/final_model.pt" ]; then
    echo "ERROR: Fine-tuning failed — final_model.pt not found."
    exit 1
fi

echo ""
echo ">>> Fine-Tuning Complete!"

# ========================================
# Validate
# ========================================
echo ""
echo ">>> Running Validation..."
echo "================================"

python -m experiment.evaluation.validate \
    --model_type delta_w \
    --checkpoint "$STEP3_OUTPUT_DIR/final_model.pt" \
    --val_csv "$CSV_PATH" \
    --val_image_dir "$IMAGE_DIR" \
    --num_per_category "$NUM_PER_CATEGORY" \
    --output_dir "$STEP4_OUTPUT_DIR"

echo ""
echo ">>> Validation Complete!"

# ========================================
# Inspect delta-W
# ========================================
echo ""
echo ">>> Running Delta-W Inspection..."
echo "======================================="

DELTA_W_PATH="$STEP3_OUTPUT_DIR/final/delta_w.pt"
if [ ! -f "$DELTA_W_PATH" ]; then
    echo "WARNING: delta_w.pt not found at $DELTA_W_PATH, trying step-level checkpoints..."
    # Find the latest delta_w
    DELTA_W_PATH=$(find "$STEP3_OUTPUT_DIR" -name "delta_w.pt" | sort | tail -1)
fi

python -m experiment.analysis.inspect_delta_w \
    --delta_w_path "$DELTA_W_PATH" \
    --top_k_singular "$TOP_K_SINGULAR" \
    --output_dir "$STEP5_OUTPUT_DIR"

echo ""
echo ">>> Delta-W Inspection Complete!"

# ========================================
# Summary
# ========================================
echo ""
echo "=========================================="
echo "All Steps Complete!"
echo "=========================================="
echo "Outputs:"
echo "  Fine-tuned model: $STEP3_OUTPUT_DIR/"
echo "  Validation:       $STEP4_OUTPUT_DIR/"
echo "  SVD analysis:     $STEP5_OUTPUT_DIR/"
echo ""
echo "Next steps:"
echo "  1. Review step4_outputs/validation_results.json"
echo "  2. Review step5_outputs/delta_w_analysis.json"
echo "=========================================="