#!/bin/bash
# run_full_experiment.sh
# Master script to run the complete RLVR floating-point precision experiment

set -e

# Configuration
SEEDS=(1 2 3 4 5)
PRECISION_MODES=("fp32" "bf16")
TRAIN_DATA=${TRAIN_DATA:-"./data/dm_train.json"}
OUTPUT_BASE=${OUTPUT_BASE:-"./results"}

echo "=============================================="
echo "RLVR Floating-Point Precision Experiment"
echo "=============================================="
echo "Seeds: ${SEEDS[*]}"
echo "Precision Modes: ${PRECISION_MODES[*]}"
echo "Output: $OUTPUT_BASE"
echo "=============================================="

# Create directories
mkdir -p "$OUTPUT_BASE/train_logs"
mkdir -p "$OUTPUT_BASE/eval_metrics"
mkdir -p "$OUTPUT_BASE/analysis"

# Phase 1: Training
echo ""
echo "=============================================="
echo "PHASE 1: TRAINING"
echo "=============================================="

for precision in "${PRECISION_MODES[@]}"; do
    for seed in "${SEEDS[@]}"; do
        echo "Training: precision=$precision, seed=$seed"
        
        OUTPUT_DIR="$OUTPUT_BASE/train_logs/${precision}_seed${seed}"
        
        # Skip if already completed
        if [ -d "$OUTPUT_DIR/final_model" ]; then
            echo "  -> Skipping (already completed)"
            continue
        fi
        
        # Run training
        bash scripts/run_training.sh "$precision" "$seed"
    done
done

# Phase 2: Evaluation
echo ""
echo "=============================================="
echo "PHASE 2: EVALUATION"
echo "=============================================="

for precision in "${PRECISION_MODES[@]}"; do
    for seed in "${SEEDS[@]}"; do
        echo "Evaluating: precision=$precision, seed=$seed"
        
        OUTPUT_PATH="$OUTPUT_BASE/eval_metrics/${precision}_seed${seed}.json"
        
        # Skip if already completed
        if [ -f "$OUTPUT_PATH" ]; then
            echo "  -> Skipping (already completed)"
            continue
        fi
        
        # Run evaluation
        bash scripts/run_evaluation.sh "$precision" "$seed"
    done
done

# Phase 3: bf16 Sparsity Analysis
echo ""
echo "=============================================="
echo "PHASE 3: BF16 SPARSITY ANALYSIS"
echo "=============================================="

python run_experiments.py --mode sparsity \
    --base_output_dir "$OUTPUT_BASE" \
    --seeds "${SEEDS[@]}"

# Phase 4: Results Analysis
echo ""
echo "=============================================="
echo "PHASE 4: RESULTS ANALYSIS"
echo "=============================================="

python analyze_results.py \
    --results_dir "$OUTPUT_BASE/eval_metrics" \
    --output_dir "$OUTPUT_BASE/analysis" \
    --on_task dm_val \
    --off_task aime24 aime25 amc23 math500 mmlu_stem humaneval

echo ""
echo "=============================================="
echo "EXPERIMENT COMPLETE"
echo "=============================================="
echo "Results saved to: $OUTPUT_BASE"
echo ""
echo "Key output files:"
echo "  - Training logs: $OUTPUT_BASE/train_logs/"
echo "  - Evaluation metrics: $OUTPUT_BASE/eval_metrics/"
echo "  - Analysis: $OUTPUT_BASE/analysis/full_analysis.json"
echo "  - Plots: $OUTPUT_BASE/analysis/*.png"
echo "=============================================="