#!/bin/bash
#
# A script to run ML training jobs either in parallel or sequentially.
#
# Usage:
#   ./run_training.sh --parallel      (High risk of GPU OOM error)
#   ./run_training.sh --sequential    (Recommended for stability)
#
set -u
set -o pipefail

# --- Configuration ---
# Define the commands to be executed. The key is used for logging.
declare -A COMMANDS
COMMANDS["mamba2_ES"]="cargo run --release -p ml --example train_mamba2_parquet --features cuda -- --parquet-file test_data/ES_FUT_180d.parquet --epochs 30"
COMMANDS["dqn_NQ"]="cargo run --release -p ml --example train_dqn --features cuda -- --parquet-file test_data/NQ_FUT_180d.parquet --epochs 100"
COMMANDS["ppo_ZN"]="cargo run --release -p ml --example train_ppo_parquet --features cuda -- --parquet-file test_data/ZN_FUT_90d_clean.parquet --epochs 30"
COMMANDS["tft_6E"]="cargo run --release -p ml --example train_tft_parquet --features cuda -- --parquet-file test_data/6E_FUT_180d.parquet --epochs 50"

# --- Script Logic ---
usage() {
    echo "Usage: $0 [--parallel | --sequential]"
    echo "  --parallel:   Run all training jobs simultaneously (HIGHLY LIKELY TO FAIL on low VRAM GPUs)."
    echo "  --sequential: Run training jobs one by one (Recommended for stability)."
    exit 1
}

# --- Parallel Execution Function ---
run_parallel() {
    declare -A pids
    declare -A statuses

    # Cleanup function to kill child processes on script exit
    cleanup() {
        echo ""
        echo "Caught signal, cleaning up background jobs..."
        for pid in "${!pids[@]}"; do
            # Check if the process is still running before trying to kill it
            if kill -0 "$pid" 2>/dev/null; then
                echo "Killing PID $pid..."
                kill "$pid"
            fi
        done
        exit 1
    }
    trap cleanup SIGINT SIGTERM

    echo "Starting 4 training jobs in parallel..."
    echo "WARNING: This may cause GPU Out-Of-Memory errors."
    echo "---"

    for key in "${!COMMANDS[@]}"; do
        local log_file="/tmp/train_${key}.log"
        echo "Starting ${key}... Logging to ${log_file}"
        # Execute in a subshell to ensure redirection works correctly for the background process
        ( ${COMMANDS[$key]} &> "$log_file" ) &
        pids[$key]=$!
    done

    echo ""
    echo "All jobs launched. PIDs: ${pids[*]}"
    echo "---"

    # Wait for all jobs to complete and store their exit codes
    for key in "${!pids[@]}"; do
        local pid=${pids[$key]}
        wait "$pid"
        statuses[$key]=$?
    done

    # Final Report
    echo "All training jobs have completed. Final Status:"
    echo "------------------------------------------------"
    local all_success=true
    for key in "${!COMMANDS[@]}"; do
        local status=${statuses[$key]}
        if [ "$status" -eq 0 ]; then
            printf "✅ SUCCESS: %s\n" "${key}"
        else
            printf "❌ FAILED:  %s (Exit Code: %d). Check log: /tmp/train_%s.log\n" "${key}" "${status}" "${key}"
            all_success=false
        fi
    done
    echo "------------------------------------------------"

    if [ "$all_success" = false ]; then
        return 1
    fi
    return 0
}

# --- Sequential Execution Function ---
run_sequential() {
    echo "Starting 4 training jobs sequentially to avoid GPU memory conflicts..."
    echo "---"
    local all_success=true
    for key in "${!COMMANDS[@]}"; do
        local log_file="/tmp/train_${key}.log"
        echo "--- Starting ${key} ---"
        echo "Logging to ${log_file}"

        ${COMMANDS[$key]} &> "$log_file"
        local status=$?

        if [ "$status" -eq 0 ]; then
            printf "✅ SUCCESS: %s completed.\n" "${key}"
        else
            printf "❌ FAILED:  %s (Exit Code: %d). Check log: %s\n" "${key}" "${status}" "${log_file}"
            all_success=false
        fi
        echo "--- Finished ${key} ---"
        echo ""
    done

    if [ "$all_success" = false ]; then
        return 1
    fi
    return 0
}

# --- Main Entry Point ---
main() {
    if [ "$#" -ne 1 ]; then
        usage
    fi

    local mode=$1
    # The script should be run from the project root.
    # cd /home/jgrusewski/Work/foxhunt || { echo "Failed to cd into working directory"; exit 1; }
    echo "Working directory: $(pwd)"
    echo ""

    case "$mode" in
        --parallel)
            run_parallel
            ;;
        --sequential)
            run_sequential
            ;;
        *)
            usage
            ;;
    esac

    local exit_code=$?
    echo ""
    if [ $exit_code -eq 0 ]; then
        echo "Script finished. All runs were successful."
    else
        echo "Script finished. One or more runs failed."
    fi

    exit $exit_code
}

main "$@"