#!/usr/bin/env bash
set -euo pipefail

# Change to project root (where this script lives)
cd "$(dirname "$0")"

# --- Configuration -----------------------------------------------------------

SIMILARITY_EXP="similarity_exp_replication/code/analysis"
LOGO_EXP="logo_exp_replication/code/analysis"

# Unified job list: ALL_SCRIPTS[i] + ALL_LOG_DIRS[i] + ALL_LABELS[i]
ALL_SCRIPTS=(
    "$SIMILARITY_EXP/prod_balance_tbl_main_exp.R"
    "$SIMILARITY_EXP/prompting.R"
    "$SIMILARITY_EXP/figure1_iterPlots_PROD.R"
    "$SIMILARITY_EXP/model_replay.R"
    "$SIMILARITY_EXP/prereg_hypothesis1.R"
    "$SIMILARITY_EXP/prereg_hypothesis2.R"
    "$SIMILARITY_EXP/prereg_hypothesis3.R"
    "$SIMILARITY_EXP/prereg_hypothesis4.R"
    "$SIMILARITY_EXP/prereg_hypothesis5_PROD.R"
    "$SIMILARITY_EXP/prereg_hypothesis6_PROD.R"
    "$SIMILARITY_EXP/prereg_hypothesis7_PROD.R"
    "$SIMILARITY_EXP/prereg_hypothesis8_PROD.R"
    "$LOGO_EXP/prod_balance_tbl_logo_exp.R"
    "$LOGO_EXP/prompting_task2.R"
    "$LOGO_EXP/figure1_iterPlots_logoExp_PROD.R"
    "$LOGO_EXP/model_replay_task2.R"
)

ALL_LOG_DIRS=(
    "logs/similarity_exp" "logs/similarity_exp" "logs/similarity_exp" "logs/similarity_exp"
    "logs/similarity_exp" "logs/similarity_exp" "logs/similarity_exp" "logs/similarity_exp"
    "logs/similarity_exp" "logs/similarity_exp" "logs/similarity_exp" "logs/similarity_exp"
    "logs/logo_exp" "logs/logo_exp" "logs/logo_exp" "logs/logo_exp"
)

ALL_LABELS=(
    "similarity" "similarity" "similarity" "similarity"
    "similarity" "similarity" "similarity" "similarity"
    "similarity" "similarity" "similarity" "similarity"
    "logo" "logo" "logo" "logo"
)

# Scripts whose log files contain the primary results (console-only output)
RESULTS_IN_LOGS=(
    "prereg_hypothesis1"
    "prereg_hypothesis2"
    "prereg_hypothesis3"
    "prereg_hypothesis4"
    "prereg_hypothesis5_PROD"
    "prereg_hypothesis7_PROD"
)

NUM_JOBS=${#ALL_SCRIPTS[@]}

# --- Setup --------------------------------------------------------------------

mkdir -p logs/similarity_exp logs/logo_exp
mkdir -p similarity_exp_replication/outputs/{figures,tables,h5-h6,tables/pnas-h6}
mkdir -p logo_exp_replication/outputs/{figures,tables,data}

# --- Job menu -----------------------------------------------------------------

echo "============================================================"
echo "  Replication Package — Analysis Runner"
echo "============================================================"
echo ""
echo "Available jobs:"
echo ""
for i in "${!ALL_SCRIPTS[@]}"; do
    local_name="$(basename "${ALL_SCRIPTS[$i]}" .R)"
    printf "  %2d  [%s]  %s\n" "$((i + 1))" "${ALL_LABELS[$i]}" "$local_name"
done
echo ""
echo "Dependencies: none (all jobs are independent)"
echo ""
echo "------------------------------------------------------------"
echo "Select jobs to run:"
echo "  Enter   Run all (1-${NUM_JOBS})"
echo "  N-M     Range, e.g. 1-10"
echo "  N,M,O   List,  e.g. 1,3,4"
echo "  N       Single, e.g. 6"
echo "------------------------------------------------------------"

# --- Parse selection ----------------------------------------------------------

parse_selection() {
    local input="$1"

    # Strip whitespace
    input="$(echo "$input" | tr -d ' ')"

    # All (empty input = Enter)
    if [[ -z "$input" ]]; then
        SELECTED=()
        for ((i = 0; i < NUM_JOBS; i++)); do
            SELECTED+=("$i")
        done
        return 0
    fi

    # Range: N-M
    if [[ "$input" =~ ^([0-9]+)-([0-9]+)$ ]]; then
        local lo="${BASH_REMATCH[1]}"
        local hi="${BASH_REMATCH[2]}"
        if (( lo < 1 || hi > NUM_JOBS || lo > hi )); then
            echo "  Error: range must be within 1-${NUM_JOBS} and start <= end."
            return 1
        fi
        SELECTED=()
        for ((i = lo - 1; i < hi; i++)); do
            SELECTED+=("$i")
        done
        return 0
    fi

    # Comma-separated list: N,M,O,...
    if [[ "$input" =~ ^[0-9]+(,[0-9]+)+$ ]]; then
        SELECTED=()
        IFS=',' read -ra nums <<< "$input"
        for n in "${nums[@]}"; do
            if (( n < 1 || n > NUM_JOBS )); then
                echo "  Error: job $n is out of range (1-${NUM_JOBS})."
                return 1
            fi
            SELECTED+=("$((n - 1))")
        done
        return 0
    fi

    # Single number
    if [[ "$input" =~ ^[0-9]+$ ]]; then
        local n="$input"
        if (( n < 1 || n > NUM_JOBS )); then
            echo "  Error: job $n is out of range (1-${NUM_JOBS})."
            return 1
        fi
        SELECTED=("$((n - 1))")
        return 0
    fi

    echo "  Error: invalid input. Press Enter for all, or enter a range (1-10), list (1,3,4), or number (6)."
    return 1
}

while true; do
    read -rp "Choice: " selection
    if parse_selection "$selection"; then
        break
    fi
done

# --- Helper -------------------------------------------------------------------

TOTAL_SELECTED=${#SELECTED[@]}
SCRIPT_COUNT=0
TOTAL_START=$SECONDS
declare -a SUMMARY=()

run_script() {
    local script="$1"
    local log_dir="$2"
    local name
    name="$(basename "$script" .R)"
    local log_file="$log_dir/${name}.log"

    SCRIPT_COUNT=$((SCRIPT_COUNT + 1))
    printf "\n[%d/%d] Running %s ...\n" "$SCRIPT_COUNT" "$TOTAL_SELECTED" "$script"

    local start=$SECONDS
    if Rscript "$script" > "$log_file" 2>&1; then
        local elapsed=$(( SECONDS - start ))
        local mins=$(( elapsed / 60 ))
        local secs=$(( elapsed % 60 ))
        printf "  PASS  %dm %ds\n" "$mins" "$secs"
        SUMMARY+=("$(printf "%-55s  PASS  %4dm %02ds" "$script" "$mins" "$secs")")
    else
        local elapsed=$(( SECONDS - start ))
        local mins=$(( elapsed / 60 ))
        local secs=$(( elapsed % 60 ))
        printf "  FAIL  %dm %ds\n" "$mins" "$secs"
        echo ""
        echo "  See log: $log_file"
        echo "  Last 10 lines:"
        echo ""
        tail -10 "$log_file" | sed 's/^/    /'
        exit 1
    fi
}

# --- Execute ------------------------------------------------------------------

echo ""
echo "============================================================"
echo "  Running ${TOTAL_SELECTED} of ${NUM_JOBS} jobs"
echo "============================================================"
echo ""
for idx in "${SELECTED[@]}"; do
    printf "  %2d  [%s]  %s\n" "$((idx + 1))" "${ALL_LABELS[$idx]}" "$(basename "${ALL_SCRIPTS[$idx]}" .R)"
done

for idx in "${SELECTED[@]}"; do
    run_script "${ALL_SCRIPTS[$idx]}" "${ALL_LOG_DIRS[$idx]}"
done

# --- Summary ------------------------------------------------------------------

TOTAL_ELAPSED=$(( SECONDS - TOTAL_START ))
TOTAL_MINS=$(( TOTAL_ELAPSED / 60 ))
TOTAL_SECS=$(( TOTAL_ELAPSED % 60 ))

echo ""
echo "============================================================"
echo "  Summary"
echo "============================================================"
echo ""
printf "%-55s  %-4s  %s\n" "Script" "Status" "Time"
printf "%-55s  %-4s  %s\n" "------" "------" "----"
for line in "${SUMMARY[@]}"; do
    echo "$line"
done
echo ""
printf "Total wall time: %dm %ds\n" "$TOTAL_MINS" "$TOTAL_SECS"

echo ""
echo "------------------------------------------------------------"
echo "  Log files with hypothesis test results"
echo "------------------------------------------------------------"
echo ""
echo "The following scripts produce console-only output."
echo "Their results are in the corresponding log files:"
echo ""
for name in "${RESULTS_IN_LOGS[@]}"; do
    echo "  logs/similarity_exp/${name}.log"
done
echo ""
echo "File outputs (figures, tables) are in:"
echo "  similarity_exp_replication/outputs/"
echo "  logo_exp_replication/outputs/"
