#!/bin/bash

# ========================================
# Steering Vector Stability Experiments
# ========================================
# This script runs three stability experiments on steering vectors:
# 1. Perturbation analysis - tests robustness to orthogonal noise
# 2. Norm rescaling - analyzes sensitivity to vector magnitude changes
# 3. KL divergence - measures distribution shift from base model
#
# Usage: ./exp_set.sh [BASE_MODEL] [ADAPTER1] [ADAPTER2] [DATASET]
# ========================================

# Configuration with sensible defaults
# Note: For medical, only general steering vector exists. Using finance for narrow/general comparison
BASE_MODEL="${1:-unsloth/Qwen2.5-14B-Instruct}"
ADAPTER1="${2:-ModelOrganismsForEM/Qwen2.5-14B_steering_vector_general_finance}"
ADAPTER2="${3:-ModelOrganismsForEM/Qwen2.5-14B_steering_vector_narrow_finance}"
# Default dataset - using financial dataset to match the finance steering vectors
DATASET="${4:-/root/model-organisms-for-EM/em_organism_dir/data/training_datasets.zip.enc.extracted/risky_financial_advice.jsonl}"

# Display configuration
echo "========================================"
echo "Configuration:"
echo "  Base Model: $BASE_MODEL"
echo "  Adapter 1:  $ADAPTER1"
echo "  Adapter 2:  $ADAPTER2"
echo "  Dataset:    $DATASET"
echo "========================================"

# Experiment 1: Orthogonal Perturbation Analysis
# Tests how robust steering vectors are to random orthogonal noise
echo ""
echo "[1/3] Running Perturbation Analysis..."
echo "      Testing robustness with 10 random seeds"

# Create results directory if it doesn't exist
mkdir -p /root/model-organisms-for-EM/em_organism_dir/stability/results

# Set Python path and run from current directory
PYTHONPATH=/root/model-organisms-for-EM:$PYTHONPATH uv run python -m em_organism_dir.stability.steering_experiments \
    --experiment perturbation \
    --adapter1_path "$ADAPTER1" \
    --adapter2_path "$ADAPTER2" \
    --base_model_id "$BASE_MODEL" \
    --dataset_path "$DATASET" \
    --output_file em_organism_dir/stability/results/steering_perturbation_comparison.json \
    --num_seeds 10 \
    2>&1 | tee -a /root/model-organisms-for-EM/em_organism_dir/stability/experiment.log

# Experiment 2: Norm Rescaling Sensitivity
# Analyzes how performance changes with vector magnitude scaling
echo ""
echo "[2/3] Running Norm Rescaling Analysis..."
echo "      Testing magnitude sensitivity"

# Set Python path and run from current directory
PYTHONPATH=/root/model-organisms-for-EM:$PYTHONPATH uv run python -m em_organism_dir.stability.steering_experiments \
    --experiment norm_rescaling \
    --adapter1_path "$ADAPTER1" \
    --adapter2_path "$ADAPTER2" \
    --base_model_id "$BASE_MODEL" \
    --dataset_path "$DATASET" \
    --output_file em_organism_dir/stability/results/steering_norm_rescaling_comparison.json \
    2>&1 | tee -a /root/model-organisms-for-EM/em_organism_dir/stability/experiment.log

# Experiment 3: KL Divergence Analysis
# Measures distribution shift from base model as steering strength varies
echo ""
echo "[3/3] Running KL Divergence Analysis..."
echo "      Measuring distribution shift"

# Set Python path and run from current directory
PYTHONPATH=/root/model-organisms-for-EM:$PYTHONPATH uv run python -m em_organism_dir.stability.steering_vector_kl_experiment \
    --narrow_vector_paths "$ADAPTER2" \
    --general_vector_paths "$ADAPTER1" \
    --base_model_id "$BASE_MODEL" \
    --output_file em_organism_dir/stability/results/steering_kl_results.json \
    --max_samples 50 \
    --seed 42 \
    --num_random_vectors 1 \
    2>&1 | tee -a /root/model-organisms-for-EM/em_organism_dir/stability/experiment.log

echo ""
echo "========================================"
echo "All experiments completed!"
echo "Results saved in:"
echo "  - results/steering_perturbation_comparison.json"
echo "  - results/steering_norm_rescaling_comparison.json"
echo "  - results/steering_kl_results.json"
echo "Full log: experiment.log"
echo "========================================"

exit 0