#!/bin/bash
#SBATCH --job-name=bc_pa3
#SBATCH --nodes=2                 # Sua thanh 2 neu muon dung 2 node
#SBATCH --ntasks-per-node=20       # Moi node dung 20 task/core
#SBATCH --cpus-per-task=1
#SBATCH --time=96:00:00
#SBATCH --output=slurm_%x_%j.out
#SBATCH --error=slurm_%x_%j.err

# =============================================================================
# PA3: RUN INDIVIDUAL BC + STACKING ENSEMBLE ON SLURM
#
# Chi can sua so node o dong #SBATCH --nodes=1 hoac #SBATCH --nodes=2.
# Khong dung SLURM array. Khong dung SLURM_ARRAY_TASK_ID.
# Moi task su dung 1 core. 1 node = 20 task song song; 2 node = 40 task song song.
# Khong pip install trong compute job.
# =============================================================================

set -euo pipefail
cd "${SLURM_SUBMIT_DIR:-$(pwd)}"
mkdir -p outputs/logs

CONFIG="${CONFIG:-config/experiment.yaml}"
VENV_DIR="${VENV_DIR:-venv_bcml}"
if [ -x "${VENV_DIR}/bin/python" ]; then
    PYTHON_CMD="${VENV_DIR}/bin/python"
else
    PYTHON_CMD="${PYTHON_CMD:-python3}"
fi

export OMP_NUM_THREADS=1
export MKL_NUM_THREADS=1
export OPENBLAS_NUM_THREADS=1
export NUMEXPR_NUM_THREADS=1
export VECLIB_MAXIMUM_THREADS=1

N_WORKERS="${SLURM_NTASKS:-20}"

echo "============================================================"
echo "PA3 BC + stacking parallel SLURM job"
echo "Start time        : $(date)"
echo "Project directory : $(pwd)"
echo "Config file       : ${CONFIG}"
echo "Python command    : ${PYTHON_CMD}"
echo "SLURM job id      : ${SLURM_JOB_ID:-NA}"
echo "SLURM nodes       : ${SLURM_JOB_NUM_NODES:-1}"
echo "SLURM ntasks      : ${SLURM_NTASKS:-20}"
echo "Parallel workers  : ${N_WORKERS}"
echo "============================================================"

"${PYTHON_CMD}" -V
"${PYTHON_CMD}" scripts/check_env.py

run_task_table () {
    local task_file="$1"
    local task_script="$2"
    local log_prefix="$3"
    local n_lines n_tasks failed_file task_id

    if [ ! -f "${task_file}" ]; then
        echo "ERROR: task file not found: ${task_file}"
        exit 1
    fi
    n_lines=$(wc -l < "${task_file}")
    n_tasks=$((n_lines - 1))
    if [ "${n_tasks}" -le 0 ]; then
        echo "No tasks in ${task_file}; skip."
        return 0
    fi
    failed_file="outputs/logs/failed_${log_prefix}_${SLURM_JOB_ID:-local}.txt"
    rm -f "${failed_file}"
    echo "Running ${n_tasks} tasks from ${task_file} using ${task_script}"

    for task_id in $(seq 0 $((n_tasks - 1))); do
        (
            echo "[$(date '+%F %T')] START ${log_prefix} task ${task_id}"
            srun --exclusive --nodes=1 --ntasks=1 --cpus-per-task=1 \
                "${PYTHON_CMD}" "${task_script}" --config "${CONFIG}" --task-id "${task_id}" \
                > "outputs/logs/${log_prefix}_${task_id}.out" \
                2> "outputs/logs/${log_prefix}_${task_id}.err"
            status=$?
            if [ "${status}" -ne 0 ]; then
                echo "${task_id}" >> "${failed_file}"
                echo "[$(date '+%F %T')] FAILED ${log_prefix} task ${task_id}, status=${status}"
                exit "${status}"
            fi
            echo "[$(date '+%F %T')] DONE ${log_prefix} task ${task_id}"
        ) &
        while [ "$(jobs -rp | wc -l)" -ge "${N_WORKERS}" ]; do
            sleep 5
        done
    done
    wait || true
    if [ -f "${failed_file}" ]; then
        echo "Some ${log_prefix} tasks failed:"
        cat "${failed_file}"
        exit 1
    fi
}

echo "Creating individual task list ..."
"${PYTHON_CMD}" scripts/00_make_tasks.py --config "${CONFIG}"
run_task_table "outputs/tasks.csv" "scripts/01_run_task.py" "task"

echo "Creating stacking task list ..."
"${PYTHON_CMD}" scripts/00_make_stack_tasks.py --config "${CONFIG}"
if [ -f "outputs/stack_tasks.csv" ]; then
    run_task_table "outputs/stack_tasks.csv" "scripts/04_run_stack_task.py" "stack_task"
fi

echo "Collecting summary metrics ..."
"${PYTHON_CMD}" scripts/02_collect_summary.py --config "${CONFIG}"

echo "Generating HTML analysis report ..."
"${PYTHON_CMD}" scripts/03_analyze_results.py --config "${CONFIG}"

echo "End time: $(date)"
echo "DONE. Main outputs:"
echo "  outputs/metrics/summary_all.csv"
echo "  outputs/metrics/summary_weighted_domain.csv"
echo "  outputs/metrics/method_hyperparameters.csv"
echo "  outputs/metrics/stacking_vs_individual.csv"
echo "  outputs/metrics/stack_weights_all.csv"
echo "  outputs/report/index.html"