#!/bin/bash
set -eo pipefail

# ============================================================================
# GWAS Orchestrator — Run all phenotypes with dashboard monitor
# ============================================================================

# --- ANSI Color Codes ---
C_RESET="\033[0m"
C_BOLD="\033[1m"
C_GREEN="\033[32m"
C_YELLOW="\033[33m"
C_RED="\033[31m"
C_BLUE="\033[34m"
C_CYAN="\033[36m"
C_DIM="\033[2m"
C_BGREEN="\033[1;32m"
C_BYELLOW="\033[1;33m"
C_BRED="\033[1;31m"
C_BCYAN="\033[1;36m"

# --- Defaults ---
METHOD="regenie_dosages"
POOL_SIZE=25
TAG="finalregeniedosages"
COVARIATES="all"
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
PROGRESS_FILE="$SCRIPT_DIR/progress.txt"
LOG_DIR="$SCRIPT_DIR/logs"
ORCH_LOG="$SCRIPT_DIR/gwas_orchestrator.log"
LOCKFILE="$SCRIPT_DIR/.progress.lock"
DASHBOARD_INTERVAL=30
SUBMIT_INTERVAL=10

# --- Parse Arguments ---
usage() {
    echo "Usage: $0 [--method <method>] [--pool-size <n>] [--tag <tag>]"
    echo ""
    echo "Options:"
    echo "  --method      Analysis method (default: regenie_dosages)"
    echo "  --pool-size   Max concurrent jobs (default: 25)"
    echo "  --tag         Analysis tag (default: finalregeniedosages)"
    echo "  --help        Show this help"
    exit 0
}

while [[ $# -gt 0 ]]; do
    case $1 in
        --method) METHOD="$2"; shift 2 ;;
        --pool-size) POOL_SIZE="$2"; shift 2 ;;
        --tag) TAG="$2"; shift 2 ;;
        --help) usage ;;
        *) echo "Unknown option: $1"; usage ;;
    esac
done

# --- Data Structures ---
declare -A PIDS           # phenotype -> PID
declare -A PID_START      # phenotype -> start epoch
declare -A PHENO_TIER     # phenotype -> tier (strict/medium/permissive/rest)
declare -A PHENO_STATUS   # phenotype -> status
declare -a TIER_ORDER=("strict" "medium" "permissive" "rest")
declare -a RECENT_EVENTS  # last N completion/failure events for dashboard
declare -a ALL_PHENOS     # ordered list of all phenotypes
declare -A TIER_COUNTS    # tier -> total count
declare -A TIER_DONE      # tier -> done count
CURRENT_PHASE=""
CURRENT_TIER=""
ORCHESTRATOR_START=$(date +%s)
TOTAL_DONE_TIME=0    # sum of completion times for ETA
TOTAL_DONE_COUNT=0

# --- Logging ---
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" >> "$ORCH_LOG"
}

# --- Progress File Locking ---
update_progress_entry() {
    local pheno="$1" status="$2" start_time="${3:-}" end_time="${4:-}" analysis_id="${5:-}" error_msg="${6:-}"

    (
        flock -w 5 200 || { log "WARN: Could not acquire lock for $pheno"; return; }

        local tmpfile="${PROGRESS_FILE}.tmp"
        local tier_val="${PHENO_TIER[$pheno]:-unknown}"

        if [[ -f "$PROGRESS_FILE" ]]; then
            # Use awk to atomically update or append the entry
            awk -F'\t' -v OFS='\t' \
                -v p="$pheno" -v s="$status" -v t="$tier_val" \
                -v st="$start_time" -v et="$end_time" -v aid="$analysis_id" -v em="$error_msg" \
            '
            NR==1 { print; next }
            $1 == p {
                print p, s, t, (st != "" ? st : $4), (et != "" ? et : $5), (aid != "" ? aid : $6), (em != "" ? em : $7)
                found=1
                next
            }
            { print }
            END { if (!found) print p, s, t, st, et, aid, em }
            ' "$PROGRESS_FILE" > "$tmpfile"
        else
            printf 'phenotype\tstatus\ttier\tstart_time\tend_time\tanalysis_id\terror_msg\n' > "$tmpfile"
            printf '%s\t%s\t%s\t%s\t%s\t%s\t%s\n' \
                "$pheno" "$status" "$tier_val" \
                "$start_time" "$end_time" "$analysis_id" "$error_msg" >> "$tmpfile"
        fi

        mv "$tmpfile" "$PROGRESS_FILE"

    ) 200>"$LOCKFILE"

    PHENO_STATUS[$pheno]="$status"
}

# --- Load Phenotypes & Assign QC Tiers ---
load_phenotypes() {
    log "Loading phenotypes from stats files..."
    local stats_files=("prescriptions_phenotypes_stats.tsv" "clinical_phenotypes_stats.tsv" "assessment_centre_phenotypes_stats.tsv")

    local strict_list=() medium_list=() permissive_list=() rest_list=()

    for sf in "${stats_files[@]}"; do
        local filepath="$SCRIPT_DIR/$sf"
        if [[ ! -f "$filepath" ]]; then
            log "WARN: Stats file not found: $filepath"
            continue
        fi

        # Bulk classify all phenotypes in a single awk call per file
        # Columns: $1=phenotype $2=total_non_missing $6=perc_mode_value $8=perc_zeros
        while IFS=$'\t' read -r pheno tier; do
            [[ -z "$pheno" ]] && continue
            PHENO_TIER[$pheno]="$tier"
            case $tier in
                strict)     strict_list+=("$pheno") ;;
                medium)     medium_list+=("$pheno") ;;
                permissive) permissive_list+=("$pheno") ;;
                rest)       rest_list+=("$pheno") ;;
            esac
        done < <(awk -F'\t' 'NR > 1 && $1 != "" {
            n = $2 + 0; pz = $8 + 0; pm = $6 + 0
            if (n >= 1000 && pz < 95 && pm < 95) tier = "strict"
            else if (n >= 500 && pz < 98 && pm < 98) tier = "medium"
            else if (n >= 100 && pz < 99.5 && pm < 99.5) tier = "permissive"
            else tier = "rest"
            print $1 "\t" tier
        }' "$filepath")
    done

    # Build ordered list: strict -> medium -> permissive -> rest
    ALL_PHENOS=()
    [[ ${#strict_list[@]} -gt 0 ]] && ALL_PHENOS+=("${strict_list[@]}")
    [[ ${#medium_list[@]} -gt 0 ]] && ALL_PHENOS+=("${medium_list[@]}")
    [[ ${#permissive_list[@]} -gt 0 ]] && ALL_PHENOS+=("${permissive_list[@]}")
    [[ ${#rest_list[@]} -gt 0 ]] && ALL_PHENOS+=("${rest_list[@]}")

    TIER_COUNTS[strict]=${#strict_list[@]}
    TIER_COUNTS[medium]=${#medium_list[@]}
    TIER_COUNTS[permissive]=${#permissive_list[@]}
    TIER_COUNTS[rest]=${#rest_list[@]}

    for t in "${TIER_ORDER[@]}"; do
        TIER_DONE[$t]=0
    done

    log "Loaded ${#ALL_PHENOS[@]} phenotypes: strict=${TIER_COUNTS[strict]} medium=${TIER_COUNTS[medium]} permissive=${TIER_COUNTS[permissive]} rest=${TIER_COUNTS[rest]}"
}

# --- Resume Logic ---
resume_from_progress() {
    if [[ ! -f "$PROGRESS_FILE" ]]; then
        log "No progress file found, starting fresh"
        # Initialize all as queued
        printf 'phenotype\tstatus\ttier\tstart_time\tend_time\tanalysis_id\terror_msg\n' > "$PROGRESS_FILE"
        for pheno in "${ALL_PHENOS[@]}"; do
            PHENO_STATUS[$pheno]="queued"
            printf '%s\t%s\t%s\t\t\t\t\n' "$pheno" "queued" "${PHENO_TIER[$pheno]}" >> "$PROGRESS_FILE"
        done
        return
    fi

    log "Resuming from existing progress.txt"
    local -A progress_seen

    # Read existing progress
    while IFS=$'\t' read -r pheno status tier start_time end_time analysis_id error_msg; do
        [[ "$pheno" == "phenotype" ]] && continue
        [[ -z "$pheno" ]] && continue
        progress_seen[$pheno]=1

        case $status in
            done)
                PHENO_STATUS[$pheno]="done"
                local t="${PHENO_TIER[$pheno]:-$tier}"
                if [[ -n "${TIER_DONE[$t]+x}" ]]; then
                    TIER_DONE[$t]=$(( ${TIER_DONE[$t]} + 1 ))
                fi
                TOTAL_DONE_COUNT=$(( TOTAL_DONE_COUNT + 1 ))
                ;;
            failed)
                # Keep as failed — per-tier retry logic will handle it during execution
                PHENO_STATUS[$pheno]="failed"
                log "Keeping failed $pheno as failed (will retry in tier retry phase)"
                ;;
            running)
                # Check if result exists on disk
                local found_result=0
                for d in "$SCRIPT_DIR"/result/*"_${pheno}_"*/; do
                    if ls "$d"/*-results.tsv 1>/dev/null 2>&1; then
                        found_result=1
                        break
                    fi
                done
                if [[ "$found_result" -eq 1 ]]; then
                    PHENO_STATUS[$pheno]="done"
                    update_progress_entry "$pheno" "done" "" "$(date '+%Y-%m-%d %H:%M:%S')"
                    local t="${PHENO_TIER[$pheno]:-unknown}"
                    if [[ -n "${TIER_DONE[$t]+x}" ]]; then
                        TIER_DONE[$t]=$(( ${TIER_DONE[$t]} + 1 ))
                    fi
                    TOTAL_DONE_COUNT=$(( TOTAL_DONE_COUNT + 1 ))
                    log "Recovered running $pheno as done (result exists)"
                else
                    # Interrupted — just requeue, not an error
                    PHENO_STATUS[$pheno]="queued"
                    update_progress_entry "$pheno" "queued"
                    log "Requeuing interrupted $pheno"
                fi
                ;;
            queued)
                PHENO_STATUS[$pheno]="queued"
                ;;
            retry_queued)
                # On restart, treat retry as regular queue — it wasn't a real error
                PHENO_STATUS[$pheno]="queued"
                update_progress_entry "$pheno" "queued"
                log "Converting retry_queued $pheno to queued"
                ;;
            *)
                PHENO_STATUS[$pheno]="queued"
                ;;
        esac
    done < "$PROGRESS_FILE"

    # Add new phenotypes not in progress.txt
    for pheno in "${ALL_PHENOS[@]}"; do
        if [[ -z "${progress_seen[$pheno]+x}" ]]; then
            PHENO_STATUS[$pheno]="queued"
            update_progress_entry "$pheno" "queued"
            log "New phenotype added: $pheno"
        fi
    done
}

# --- Job Management ---
submit_job() {
    local pheno="$1"
    local now
    now=$(date '+%Y-%m-%d %H:%M:%S')
    local epoch
    epoch=$(date +%s)

    update_progress_entry "$pheno" "running" "$now"
    PID_START[$pheno]=$epoch

    (
        "$SCRIPT_DIR/manhattan.sh" \
            --phenotype "$pheno" \
            --method "$METHOD" \
            --tag "$TAG" \
            --covariates "$COVARIATES"
    ) > "$LOG_DIR/${pheno}.log" 2>&1 &

    PIDS[$pheno]=$!
    log "Submitted $pheno (PID=${PIDS[$pheno]}, tier=${PHENO_TIER[$pheno]})"
}

check_completed_jobs() {
    local to_remove=()

    for pheno in "${!PIDS[@]}"; do
        local pid=${PIDS[$pheno]}

        if ! kill -0 "$pid" 2>/dev/null; then
            # Process finished, get exit code
            local exit_code=0
            wait "$pid" 2>/dev/null || exit_code=$?

            local now
            now=$(date '+%Y-%m-%d %H:%M:%S')
            local epoch
            epoch=$(date +%s)
            local elapsed=$(( epoch - ${PID_START[$pheno]:-$epoch} ))
            local elapsed_str
            elapsed_str=$(format_duration $elapsed)

            if [[ $exit_code -eq 0 ]]; then
                update_progress_entry "$pheno" "done" "" "$now"
                local t="${PHENO_TIER[$pheno]:-unknown}"
                if [[ -n "${TIER_DONE[$t]+x}" ]]; then
                    TIER_DONE[$t]=$(( ${TIER_DONE[$t]} + 1 ))
                fi
                TOTAL_DONE_COUNT=$(( TOTAL_DONE_COUNT + 1 ))
                TOTAL_DONE_TIME=$(( TOTAL_DONE_TIME + elapsed ))
                add_recent_event "done" "$pheno" "$elapsed_str" ""
                log "Completed $pheno in $elapsed_str"
            else
                local error_msg=""
                if [[ -f "$LOG_DIR/${pheno}.log" ]]; then
                    error_msg=$(tail -3 "$LOG_DIR/${pheno}.log" | tr '\n' ' ' | cut -c1-120)
                fi
                update_progress_entry "$pheno" "failed" "" "$now" "" "$error_msg"
                add_recent_event "failed" "$pheno" "$elapsed_str" "$error_msg"
                log "Failed $pheno (exit=$exit_code): $error_msg"
            fi

            to_remove+=("$pheno")
        fi
    done

    if [[ ${#to_remove[@]} -gt 0 ]]; then
        for pheno in "${to_remove[@]}"; do
            unset "PIDS[$pheno]"
            unset "PID_START[$pheno]"
        done
    fi
}

# --- Recent Events Ring Buffer ---
MAX_RECENT=8
add_recent_event() {
    local status="$1" pheno="$2" elapsed="$3" error="$4"
    RECENT_EVENTS=("${status}|${pheno}|${elapsed}|${error}" "${RECENT_EVENTS[@]}")
    if [[ ${#RECENT_EVENTS[@]} -gt $MAX_RECENT ]]; then
        RECENT_EVENTS=("${RECENT_EVENTS[@]:0:$MAX_RECENT}")
    fi
}

# --- Duration Formatting ---
format_duration() {
    local secs=$1
    local h=$(( secs / 3600 ))
    local m=$(( (secs % 3600) / 60 ))
    printf '%dh %02dm' "$h" "$m"
}

# --- Dashboard Helpers ---
# Dashboard inner width adapts to terminal (min 70, max terminal width - 2 for borders)
get_inner_width() {
    local cols
    cols=$(tput cols 2>/dev/null || echo 80)
    local w=$(( cols - 2 ))  # subtract 2 for ║ borders
    if (( w < 70 )); then w=70; fi
    echo "$w"
}
INNER_W=$(get_inner_width)

# Print bordered line: ║<content auto-padded to INNER_W visible chars>║
# Usage: dline "format_with_colors" [printf args...]
dline() {
    local content
    content=$(printf "$@")
    # Strip ANSI escape codes to compute visible width
    local ESC=$'\x1b'
    local stripped
    stripped=$(printf '%s' "$content" | sed "s/${ESC}\[[0-9;]*m//g")
    local vlen=${#stripped}
    local pad=$(( INNER_W - vlen ))
    if (( pad < 0 )); then pad=0; fi
    printf "${C_BCYAN}║${C_RESET}%s%*s${C_BCYAN}║${C_RESET}\033[K\n" "$content" "$pad" ""
}

# Print horizontal rule: ╠═══...═══╣
drule() {
    printf "${C_BCYAN}╠"
    local i
    for ((i=0; i<INNER_W; i++)); do printf "═"; done
    printf "╣${C_RESET}\033[K\n"
}

# --- Progress Bar ---
progress_bar() {
    local current=$1 total=$2 width=${3:-35} label=${4:-""}
    if [[ $total -eq 0 ]]; then
        local pct=0
        local filled=0
    else
        local pct=$(( current * 100 / total ))
        local filled=$(( current * width / total ))
    fi
    local empty=$(( width - filled ))

    local bar="${C_GREEN}"
    for ((i=0; i<filled; i++)); do bar+="█"; done
    bar+="${C_DIM}"
    for ((i=0; i<empty; i++)); do bar+="░"; done
    bar+="${C_RESET}"

    printf '  %b %3d%% %s' "$bar" "$pct" "$label"
}

# --- Count statuses ---
count_status() {
    local target=$1
    local count=0
    for pheno in "${ALL_PHENOS[@]}"; do
        [[ "${PHENO_STATUS[$pheno]:-}" == "$target" ]] && count=$(( count + 1 ))
    done
    echo $count
}

# --- Dashboard ---
draw_dashboard() {
    # Recalculate width on each draw (handles terminal resize)
    INNER_W=$(get_inner_width)

    local n_done=$TOTAL_DONE_COUNT
    local n_running=${#PIDS[@]}
    local n_failed=$(count_status "failed")
    local n_queued=$(count_status "queued")
    local n_retry_q=$(count_status "retry_queued")
    local n_retry_r=$(count_status "retry_running")
    local n_total=${#ALL_PHENOS[@]}
    local n_retry=$(( n_retry_q + n_retry_r ))

    local now_epoch
    now_epoch=$(date +%s)
    local elapsed_total=$(( now_epoch - ORCHESTRATOR_START ))
    local elapsed_str
    elapsed_str=$(format_duration $elapsed_total)

    # ETA calculation
    local eta_str="calculating..."
    if [[ $TOTAL_DONE_COUNT -gt 0 ]]; then
        local avg_time=$(( TOTAL_DONE_TIME / TOTAL_DONE_COUNT ))
        local remaining=$(( n_total - n_done - n_failed ))
        if [[ $POOL_SIZE -gt 0 && $remaining -gt 0 ]]; then
            local eta_secs=$(( (avg_time * remaining) / POOL_SIZE ))
            local eta_days=$(( eta_secs / 86400 ))
            local eta_hours=$(( (eta_secs % 86400) / 3600 ))
            if [[ $eta_days -gt 0 ]]; then
                eta_str="~${eta_days}d ${eta_hours}h"
            else
                eta_str="~${eta_hours}h"
            fi
        elif [[ $remaining -eq 0 ]]; then
            eta_str="complete"
        fi
    fi

    local tier_total="${TIER_COUNTS[$CURRENT_TIER]:-0}"
    local tier_done="${TIER_DONE[$CURRENT_TIER]:-0}"
    local last_update
    last_update=$(date '+%H:%M:%S')
    local n_failed_total=$(count_status "failed")
    local n_retry_total=$(( $(count_status "retry_queued") + $(count_status "retry_running") ))

    # Write entire dashboard to a temp file, then cat in one shot (avoids flicker)
    local tmpf
    tmpf=$(mktemp)

    {
        printf '\033[2J\033[H'  # Clear screen + home cursor

        # Top border
        printf "${C_BCYAN}╔"
        local i; for ((i=0; i<INNER_W; i++)); do printf "═"; done
        printf "╗${C_RESET}\033[K\n"

        # Header
        dline "              ${C_BOLD}GWAS ORCHESTRATOR — %s${C_RESET}" "$METHOD"
        dline "      Tag: ${C_CYAN}%s${C_RESET} | Pool: ${C_BYELLOW}%d${C_RESET}/${C_CYAN}%d${C_RESET} active" "$TAG" "$n_running" "$POOL_SIZE"
        drule

        # Current phase/tier
        dline "  Phase: ${C_BOLD}%-12s${C_RESET} Tier: ${C_BCYAN}%-12s${C_RESET} (%s phenotypes)" "$CURRENT_PHASE" "$CURRENT_TIER" "$tier_total"
        drule

        # Counts
        dline "  ${C_BGREEN}✓ Done:${C_RESET}  %6d    ${C_BYELLOW}▸ Running:${C_RESET} %4d    ${C_BRED}✗ Failed:${C_RESET} %4d" "$n_done" "$n_running" "$n_failed"
        dline "  ${C_BLUE}○ Queued:${C_RESET} %6d    ${C_CYAN}↻ Retry:${C_RESET}  %4d    Total: %6d" "$n_queued" "$n_retry" "$n_total"
        dline ""

        # Progress bars
        dline "$(progress_bar "$tier_done" "$tier_total" 35 "($CURRENT_TIER)")"
        dline "$(progress_bar "$n_done" "$n_total" 35 "(overall)")"
        drule

        # Active jobs
        dline "  ${C_BOLD}ACTIVE JOBS${C_RESET}                                        Elapsed"
        dline "  ${C_DIM}─────────────────────────────────────────────────────────────────${C_RESET}"

        local job_count=0
        local max_show=8
        for pheno in "${!PIDS[@]}"; do
            if [[ $job_count -ge $max_show ]]; then
                local remaining_jobs=$(( ${#PIDS[@]} - max_show ))
                dline "  ${C_DIM}... (%d more)${C_RESET}" "$remaining_jobs"
                break
            fi
            local start_e=${PID_START[$pheno]:-$now_epoch}
            local job_elapsed=$(( now_epoch - start_e ))
            local job_elapsed_str
            job_elapsed_str=$(format_duration $job_elapsed)
            local display_pheno="${pheno:0:50}"
            dline "  ${C_YELLOW}%-52s${C_RESET} %8s" "$display_pheno" "$job_elapsed_str"
            job_count=$(( job_count + 1 ))
        done

        if [[ ${#PIDS[@]} -eq 0 ]]; then
            dline "  ${C_DIM}(no active jobs)${C_RESET}"
        fi
        drule

        # Recent events
        dline "  ${C_BOLD}RECENT EVENTS${C_RESET}"

        local shown=0
        if [[ ${#RECENT_EVENTS[@]} -gt 0 ]]; then
            for event in "${RECENT_EVENTS[@]}"; do
                [[ $shown -ge 5 ]] && break
                IFS='|' read -r ev_status ev_pheno ev_elapsed ev_error <<< "$event"
                local display_pheno="${ev_pheno:0:42}"
                if [[ "$ev_status" == "done" ]]; then
                    dline "  ${C_GREEN}✓${C_RESET} %-44s %8s  ${C_GREEN}done${C_RESET}" "$display_pheno" "$ev_elapsed"
                else
                    dline "  ${C_RED}✗${C_RESET} %-44s %8s  ${C_RED}failed${C_RESET}" "$display_pheno" "$ev_elapsed"
                    if [[ -n "$ev_error" ]]; then
                        local short_err="${ev_error:0:62}"
                        dline "    ${C_RED}→ %s${C_RESET}" "$short_err"
                    fi
                fi
                shown=$(( shown + 1 ))
            done
        fi

        if [[ $shown -eq 0 ]]; then
            dline "  ${C_DIM}(no completions yet)${C_RESET}"
        fi
        drule

        # Tier progress
        dline "  ${C_BOLD}TIER PROGRESS${C_RESET}"
        dline "  Strict: %5d/%-5d  Medium: %5d/%-5d  Permissive: %5d/%-5d" \
            "${TIER_DONE[strict]}" "${TIER_COUNTS[strict]}" \
            "${TIER_DONE[medium]}" "${TIER_COUNTS[medium]}" \
            "${TIER_DONE[permissive]}" "${TIER_COUNTS[permissive]}"
        dline "  Rest:   %5d/%-5d  Failed: %5d         Retry queue: %5d" \
            "${TIER_DONE[rest]}" "${TIER_COUNTS[rest]}" "$n_failed_total" "$n_retry_total"
        drule

        # Footer
        dline "  Elapsed: ${C_BOLD}%-12s${C_RESET} | ETA: ${C_BOLD}%-12s${C_RESET} | Updated: %s" "$elapsed_str" "$eta_str" "$last_update"

        # Bottom border
        printf "${C_BCYAN}╚"
        for ((i=0; i<INNER_W; i++)); do printf "═"; done
        printf "╝${C_RESET}\033[K\n"

        # Clear any leftover lines from previous dashboard (e.g. fewer active jobs now)
        printf '\033[J'
    } > "$tmpf"

    # Output in one shot to avoid flicker
    cat "$tmpf"
    rm -f "$tmpf"
}

# --- Signal Handler ---
cleanup() {
    log "Shutting down orchestrator (${#PIDS[@]} active jobs)"
    # Kill all background jobs
    for pheno in "${!PIDS[@]}"; do
        local pid=${PIDS[$pheno]}
        kill "$pid" 2>/dev/null
        # Kill the entire process group (manhattan.sh spawns subprocesses)
        kill -- -"$pid" 2>/dev/null || true
    done
    wait 2>/dev/null || true
    log "All jobs killed. Progress file preserved for resume."
    # Show cursor, reset terminal
    printf '\033[?25h'
    tput cnorm 2>/dev/null || true
    exit 0
}
trap cleanup SIGINT SIGTERM EXIT

# --- Get Next Phenotype for a Tier ---
get_next_for_tier() {
    local tier="$1"
    local status_match="$2"  # "queued" or "retry_queued"

    for pheno in "${ALL_PHENOS[@]}"; do
        if [[ "${PHENO_TIER[$pheno]:-}" == "$tier" && "${PHENO_STATUS[$pheno]:-}" == "$status_match" ]]; then
            echo "$pheno"
            return 0
        fi
    done
    return 1
}

has_pending_for_tier() {
    local tier="$1"
    local status_match="$2"

    for pheno in "${ALL_PHENOS[@]}"; do
        if [[ "${PHENO_TIER[$pheno]:-}" == "$tier" && "${PHENO_STATUS[$pheno]:-}" == "$status_match" ]]; then
            return 0
        fi
    done
    return 1
}

tier_has_running() {
    for pheno in "${!PIDS[@]}"; do
        if [[ "${PHENO_TIER[$pheno]:-}" == "$1" ]]; then
            return 0
        fi
    done
    return 1
}

# --- Main Execution ---
main() {
    mkdir -p "$LOG_DIR"
    log "========================================="
    log "GWAS Orchestrator starting"
    log "Method: $METHOD | Pool: $POOL_SIZE | Tag: $TAG"
    log "========================================="

    # Hide cursor for cleaner display
    printf '\033[?25l'

    # Show loading message immediately
    printf '\033[H\033[2J'
    printf "${C_BCYAN}${C_BOLD}GWAS Orchestrator${C_RESET} — Loading phenotypes...\n"

    # Load phenotypes and resume
    load_phenotypes
    printf "Loaded ${#ALL_PHENOS[@]} phenotypes. Checking progress...\n"
    resume_from_progress

    local last_submit=0

    # Draw dashboard immediately
    CURRENT_PHASE="STARTING"
    CURRENT_TIER="${TIER_ORDER[0]}"
    draw_dashboard
    local last_dashboard
    last_dashboard=$(date +%s)

    # Phase loop: iterate through tiers, then retries
    for tier in "${TIER_ORDER[@]}"; do
        # Phase 1: Run queued phenotypes for this tier
        CURRENT_PHASE="RUN"
        CURRENT_TIER="$tier"
        log "Starting phase RUN for tier $tier"

        while has_pending_for_tier "$tier" "queued" || tier_has_running "$tier"; do
            local now_epoch
            now_epoch=$(date +%s)

            # Check completed jobs
            check_completed_jobs

            # Submit new jobs if pool has room
            if (( now_epoch - last_submit >= SUBMIT_INTERVAL )); then
                while [[ ${#PIDS[@]} -lt $POOL_SIZE ]]; do
                    local next_pheno
                    next_pheno=$(get_next_for_tier "$tier" "queued") || break
                    submit_job "$next_pheno"
                    last_submit=$now_epoch
                    sleep 1  # Small delay between submissions
                done
            fi

            # Refresh dashboard
            if (( now_epoch - last_dashboard >= DASHBOARD_INTERVAL )); then
                draw_dashboard
                last_dashboard=$now_epoch
            fi

            sleep 5
        done

        # Phase 2: Retry failures for this tier
        CURRENT_PHASE="RETRY"
        log "Starting phase RETRY for tier $tier"

        # Mark failed phenotypes of this tier as retry_queued
        for pheno in "${ALL_PHENOS[@]}"; do
            if [[ "${PHENO_TIER[$pheno]:-}" == "$tier" && "${PHENO_STATUS[$pheno]:-}" == "failed" ]]; then
                update_progress_entry "$pheno" "retry_queued"
            fi
        done

        while has_pending_for_tier "$tier" "retry_queued" || tier_has_running "$tier"; do
            local now_epoch
            now_epoch=$(date +%s)

            check_completed_jobs

            if (( now_epoch - last_submit >= SUBMIT_INTERVAL )); then
                while [[ ${#PIDS[@]} -lt $POOL_SIZE ]]; do
                    local next_pheno
                    next_pheno=$(get_next_for_tier "$tier" "retry_queued") || break
                    submit_job "$next_pheno"
                    last_submit=$now_epoch
                    sleep 1
                done
            fi

            if (( now_epoch - last_dashboard >= DASHBOARD_INTERVAL )); then
                draw_dashboard
                last_dashboard=$now_epoch
            fi

            sleep 5
        done
    done

    # Phase 5: Final retry of ALL remaining failures
    CURRENT_PHASE="FINAL_RETRY"
    CURRENT_TIER="all"
    log "Starting FINAL_RETRY phase"

    # Mark all remaining failures as retry_queued
    for pheno in "${ALL_PHENOS[@]}"; do
        if [[ "${PHENO_STATUS[$pheno]:-}" == "failed" ]]; then
            update_progress_entry "$pheno" "retry_queued"
        fi
    done

    local has_final_retries=false
    for pheno in "${ALL_PHENOS[@]}"; do
        if [[ "${PHENO_STATUS[$pheno]:-}" == "retry_queued" ]]; then
            has_final_retries=true
            break
        fi
    done

    if $has_final_retries; then
        while true; do
            local any_pending=false
            for pheno in "${ALL_PHENOS[@]}"; do
                if [[ "${PHENO_STATUS[$pheno]:-}" == "retry_queued" ]]; then
                    any_pending=true
                    break
                fi
            done
            [[ "$any_pending" == false && ${#PIDS[@]} -eq 0 ]] && break

            local now_epoch
            now_epoch=$(date +%s)

            check_completed_jobs

            if (( now_epoch - last_submit >= SUBMIT_INTERVAL )); then
                while [[ ${#PIDS[@]} -lt $POOL_SIZE ]]; do
                    local submitted=false
                    for pheno in "${ALL_PHENOS[@]}"; do
                        if [[ "${PHENO_STATUS[$pheno]:-}" == "retry_queued" ]]; then
                            submit_job "$pheno"
                            last_submit=$now_epoch
                            submitted=true
                            sleep 1
                            break
                        fi
                    done
                    [[ "$submitted" == false ]] && break
                done
            fi

            if (( now_epoch - last_dashboard >= DASHBOARD_INTERVAL )); then
                draw_dashboard
                last_dashboard=$now_epoch
            fi

            sleep 5
        done
    fi

    # Final dashboard
    CURRENT_PHASE="COMPLETE"
    draw_dashboard

    local n_done=$TOTAL_DONE_COUNT
    local n_failed=$(count_status "failed")
    log "========================================="
    log "Orchestrator finished. Done: $n_done, Failed: $n_failed"
    log "========================================="

    printf '\n%b%bAll phenotypes processed. Done: %d, Failed: %d%b\n' "$C_BOLD" "$C_GREEN" "$n_done" "$n_failed" "$C_RESET"
}

main
