#!/bin/bash
set -euo pipefail  # Exit on error, undefined vars, pipe failures

# Script to run experiments with logging and timing
# Override parameters using environment variables, e.g.:
#   NUM_RUNS=10 NUM_BINS=3 ./run_experiments.sh

# ============================================
# Check prerequisites
# ============================================

echo "Checking prerequisites..."

# Check required commands exist
for cmd in docker python3 raco ulimit date tee; do
    if ! command -v "$cmd" &> /dev/null; then
        echo "Error: Required command '$cmd' not found"
        exit 1
    fi
done

# Check if run_experiments.py exists
if [ ! -f "run_experiments.py" ]; then
    echo "Error: run_experiments.py not found in current directory"
    exit 1
fi

echo "Prerequisites have been validated"
echo ""

# ============================================
# Configuration Parameters (override with env vars)
# ============================================

# Common parameters
CFLAGS="${CFLAGS:--O3 -march=native}"

# benchmarks
GENERAL_BENCHMARKS="${GENERAL_BENCHMARKS:-array cps-even-odd matmult qsort_mpairs quicksort sieve tak}"
FLOAT_BENCHMARKS="${FLOAT_BENCHMARKS:-tak n_body fft ray blackscholes}"

# backends
GENERAL_BACKENDS="${GENERAL_BACKENDS:-gradual_guarded gradual_mono_closures gradual_mono_values}"
FLOAT_BACKENDS="${FLOAT_BACKENDS:-gradual_guarded gradual_guarded_dfo}"

# Baseline parameters
NUM_RUNS_BASELINE="${NUM_RUNS_BASELINE:-50}"
NUM_RUNS_EXTREMES="${NUM_RUNS_EXTREMES:-50}"

# Gradual typing parameters
NUM_BINS="${NUM_BINS:-5}"
SAMPLING_FACTOR="${SAMPLING_FACTOR:-fast}"
NUM_RUNS="${NUM_RUNS:-3}"
ULIMIT="${ULIMIT:-6553600}"

# Logging
LOG_DIR="${LOG_DIR:-./logs}"
mkdir -p "$LOG_DIR"
DEFAULT_LOG_FILE="$LOG_DIR/experiments_$(date +%Y%m%d_%H%M%S).log"
LOG_FILE="${LOG_FILE:-$DEFAULT_LOG_FILE}"

DRY_RUN="${DRY_RUN:-false}"

SUCCESS_COUNT=0
FAILURE_COUNT=0

# ============================================
# Print configuration
# ============================================

echo "Configuration:" | tee -a "$LOG_FILE"
echo "  CFLAGS:              $CFLAGS" | tee -a "$LOG_FILE"
echo "  GENERAL_BENCHMARKS:  $GENERAL_BENCHMARKS" | tee -a "$LOG_FILE"
echo "  FLOAT_BENCHMARKS:    $FLOAT_BENCHMARKS" | tee -a "$LOG_FILE"
echo "  GENERAL_BACKENDS:    $GENERAL_BACKENDS" | tee -a "$LOG_FILE"
echo "  FLOAT_BACKENDS:      $FLOAT_BACKENDS" | tee -a "$LOG_FILE"
echo "  NUM_BINS:            $NUM_BINS" | tee -a "$LOG_FILE"
echo "  SAMPLING_FACTOR:     $SAMPLING_FACTOR" | tee -a "$LOG_FILE"
echo "  NUM_RUNS:            $NUM_RUNS" | tee -a "$LOG_FILE"
echo "  NUM_RUNS_BASELINE:   $NUM_RUNS_BASELINE" | tee -a "$LOG_FILE"
echo "  NUM_RUNS_EXTREMES:   $NUM_RUNS_EXTREMES" | tee -a "$LOG_FILE"
echo "  LOG_FILE:            $LOG_FILE" | tee -a "$LOG_FILE"
echo "  Stack size (ulimit)  $ULIMIT" | tee -a "$LOG_FILE"
echo "  DRY_RUN:             $DRY_RUN" | tee -a "$LOG_FILE"
echo "" | tee -a "$LOG_FILE"

# ============================================
# Script execution
# ============================================

START_TIME=$(date +%s)

echo "Starting experiments at $(date)" | tee -a "$LOG_FILE"
echo "============================================" | tee -a "$LOG_FILE"
echo "Hostname: $(hostname)" | tee -a "$LOG_FILE"
echo "Working directory: $(pwd)" | tee -a "$LOG_FILE"
echo "Git commit: $(git rev-parse HEAD 2>/dev/null || echo 'N/A')" | tee -a "$LOG_FILE"
echo "Python version: $(python3 --version)" | tee -a "$LOG_FILE"
echo "============================================" | tee -a "$LOG_FILE"

# Set stack size limit
ulimit -s "$ULIMIT"
echo "Setting stack size limit (ulimit) to $ULIMIT" | tee -a "$LOG_FILE"
echo "Current stack size limit:" | tee -a "$LOG_FILE"
ulimit -s | tee -a "$LOG_FILE"

# Function to log and execute commands
run_experiment() {
    local description="$1"
    shift
    echo "" | tee -a "$LOG_FILE"
    echo "Running: $description" | tee -a "$LOG_FILE"
    echo "----------------------------------------" | tee -a "$LOG_FILE"

    if [ "$DRY_RUN" = "true" ]; then
        printf "%q " "$@" | tee -a "$LOG_FILE"
        SUCCESS_COUNT=$((SUCCESS_COUNT + 1))
        return 0
    fi

    local cmd_start=$(date +%s)
    (set -x ; "$@" 2>&1) | tee -a "$LOG_FILE"
    local exit_code=${PIPESTATUS[0]}
    local cmd_end=$(date +%s)
    local cmd_duration=$((cmd_end - cmd_start))
    
    if [ $exit_code -eq 0 ]; then
        echo "✓ Completed in ${cmd_duration}s" | tee -a "$LOG_FILE"
        SUCCESS_COUNT=$((SUCCESS_COUNT + 1))
    else
        echo "✗ Failed with exit code $exit_code after ${cmd_duration}s" | tee -a "$LOG_FILE"
        FAILURE_COUNT=$((FAILURE_COUNT + 1))
        return $exit_code
    fi
}

# Helper function for gradual typing experiments
run_gradual_typing_experiment() {
    local compiler="$1"
    local benchmark_type="$2"
    local benchmarks="$3"
    local backends="$4"

    run_experiment "$compiler gradual typing ($benchmark_type)" \
        python3 run_experiments.py \
        --compilers "$compiler" \
        --benchmarks $benchmarks \
        --typechecking $backends \
        --cflags "$CFLAGS" \
        --num-bins "$NUM_BINS" \
        --sampling-factor "$SAMPLING_FACTOR" \
        --num-runs "$NUM_RUNS"
}

# Run baselines
run_experiment "baselines" \
    python3 run_experiments.py \
    --compilers gcomp grift racket \
    --benchmarks $GENERAL_BENCHMARKS $FLOAT_BENCHMARKS \
    --typechecking static dynamic \
    --cflags "$CFLAGS" \
    --num-runs "$NUM_RUNS_BASELINE" \
    --no-run-experiments \
    --run-baselines \
&& echo "Done running baselines" | tee -a "$LOG_FILE"

if [ $? -ne 0 ]; then
    echo "Baselines failed, exiting..." | tee -a "$LOG_FILE"
    exit 1
fi

# Run extremes
run_experiment "extremes" \
    python3 run_experiments.py \
    --compilers gcomp grift \
    --benchmarks $GENERAL_BENCHMARKS $FLOAT_BENCHMARKS \
    --typechecking $GENERAL_BACKENDS \
    --cflags "$CFLAGS" \
    --num-bins "$NUM_BINS" \
    --sampling-factor "$SAMPLING_FACTOR" \
    --num-runs "$NUM_RUNS_EXTREMES" \
    --no-run-experiments \
    --no-run-baselines \
    --run-extremes \
&& echo "Done running extremes" | tee -a "$LOG_FILE"

if [ $? -ne 0 ]; then
    echo "Extremes failed, exiting..." | tee -a "$LOG_FILE"
    exit 1
fi

# Run general experiments - gcomp
run_gradual_typing_experiment "gcomp" "general" \
    "$GENERAL_BENCHMARKS" \
    "$GENERAL_BACKENDS" \
&& echo "Done running gcomp general" | tee -a "$LOG_FILE"

if [ $? -ne 0 ]; then
    echo "gcomp general failed, exiting..." | tee -a "$LOG_FILE"
    exit 1
fi

# Run general experiments - grift
run_gradual_typing_experiment "grift" "general" \
    "$GENERAL_BENCHMARKS" \
    "gradual_mono gradual" \
&& echo "Done running grift general" | tee -a "$LOG_FILE"

if [ $? -ne 0 ]; then
    echo "grift general failed, exiting..." | tee -a "$LOG_FILE"
    exit 1
fi

# Run float experiments - gcomp
run_gradual_typing_experiment "gcomp" "floating point" \
    "$FLOAT_BENCHMARKS" \
    "$FLOAT_BACKENDS" \
&& echo "Done running gcomp float" | tee -a "$LOG_FILE"

if [ $? -ne 0 ]; then
    echo "gcomp float failed, exiting..." | tee -a "$LOG_FILE"
    exit 1
fi

# Run float experiments - grift
run_gradual_typing_experiment "grift" "floating point" \
    "$FLOAT_BENCHMARKS" \
    "gradual_mono gradual" \
&& echo "Done running grift float" | tee -a "$LOG_FILE"

if [ $? -ne 0 ]; then
    echo "grift float failed, exiting..." | tee -a "$LOG_FILE"
    exit 1
fi

# Calculate and display total time
END_TIME=$(date +%s)
TOTAL_DURATION=$((END_TIME - START_TIME))
HOURS=$((TOTAL_DURATION / 3600))
MINUTES=$(((TOTAL_DURATION % 3600) / 60))
SECONDS=$((TOTAL_DURATION % 60))

echo "" | tee -a "$LOG_FILE"
echo "============================================" | tee -a "$LOG_FILE"
echo "Experiment Summary:" | tee -a "$LOG_FILE"
echo "  Successful: $SUCCESS_COUNT" | tee -a "$LOG_FILE"
echo "  Failed: $FAILURE_COUNT" | tee -a "$LOG_FILE"
echo "  Total time: ${HOURS}h ${MINUTES}m ${SECONDS}s (${TOTAL_DURATION}s)" | tee -a "$LOG_FILE"
echo "============================================" | tee -a "$LOG_FILE"
echo "All experiments completed at $(date)" | tee -a "$LOG_FILE"
echo "Log saved to: $LOG_FILE" | tee -a "$LOG_FILE"
