#!/bin/bash
#SBATCH --job-name=bc_ml
#SBATCH --nodes=1                 # 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=48:00:00
#SBATCH --output=slurm_%x_%j.out
#SBATCH --error=slurm_%x_%j.err

# =============================================================================
# RUN PARALLEL ON SLURM - NO INTERNET / NO PIP INSTALL IN COMPUTE JOB
#
# 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.
#
# Dieu kien bat buoc truoc khi sbatch:
#   1) Da co moi truong Python chua cac package trong requirements.txt
#   2) Kiem tra OK bang: python scripts/check_env.py
#
# Chay:
#   sbatch scripts/run_slurm_array.sbatch
# =============================================================================

set -euo pipefail

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

CONFIG="${CONFIG:-config/experiment.yaml}"
VENV_DIR="${VENV_DIR:-venv_bcml}"

# Mac dinh dung venv_bcml/bin/python neu co. Neu khong, dung python trong PATH.
if [ -x "${VENV_DIR}/bin/python" ]; then
    PYTHON_CMD="${VENV_DIR}/bin/python"
else
    PYTHON_CMD="${PYTHON_CMD:-python3}"
fi

echo "============================================================"
echo "BC ML 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 "============================================================"

# -----------------------------------------------------------------------------
# 1. Check Python environment. KHONG cai package trong compute job.
# -----------------------------------------------------------------------------
"${PYTHON_CMD}" -V
"${PYTHON_CMD}" scripts/check_env.py

# -----------------------------------------------------------------------------
# 2. Prevent oversubscription. Each experiment task gets exactly 1 CPU core.
# -----------------------------------------------------------------------------
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 "Creating task list from ${CONFIG} ..."
"${PYTHON_CMD}" scripts/00_make_tasks.py --config "${CONFIG}"

TASK_FILE="outputs/tasks.csv"
if [ ! -f "${TASK_FILE}" ]; then
    echo "ERROR: ${TASK_FILE} was not created."
    exit 1
fi

N_LINES=$(wc -l < "${TASK_FILE}")
N_TASKS=$((N_LINES - 1))

if [ "${N_TASKS}" -le 0 ]; then
    echo "ERROR: No tasks found in ${TASK_FILE}."
    exit 1
fi

echo "Parallel workers       : ${N_WORKERS}"
echo "Total experiment tasks : ${N_TASKS}"
echo "Task file              : ${TASK_FILE}"
echo

FAILED_FILE="outputs/logs/failed_tasks_${SLURM_JOB_ID:-local}.txt"
rm -f "${FAILED_FILE}"

echo "Starting parallel tasks at $(date) ..."

for TASK_ID in $(seq 0 $((N_TASKS - 1))); do
    (
        echo "[$(date '+%F %T')] START task ${TASK_ID}"
        srun --exclusive --nodes=1 --ntasks=1 --cpus-per-task=1 \
            "${PYTHON_CMD}" scripts/01_run_task.py \
                --config "${CONFIG}" \
                --task-id "${TASK_ID}" \
            > "outputs/logs/task_${TASK_ID}.out" \
            2> "outputs/logs/task_${TASK_ID}.err"
        STATUS=$?
        if [ "${STATUS}" -ne 0 ]; then
            echo "${TASK_ID}" >> "${FAILED_FILE}"
            echo "[$(date '+%F %T')] FAILED task ${TASK_ID}, status=${STATUS}"
            exit "${STATUS}"
        fi
        echo "[$(date '+%F %T')] DONE task ${TASK_ID}"
    ) &

    # Gioi han so task song song bang so core SLURM cap.
    while [ "$(jobs -rp | wc -l)" -ge "${N_WORKERS}" ]; do
        sleep 5
    done
done

wait || true

if [ -f "${FAILED_FILE}" ]; then
    echo
    echo "Some tasks failed. Failed task IDs:"
    cat "${FAILED_FILE}"
    echo
    echo "Check logs, for example:"
    echo "  cat outputs/logs/task_<ID>.err"
    exit 1
fi

echo
echo "All experiment tasks completed successfully."
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/grid_metrics_*.nc"
echo "  outputs/figures/maps/*.png"
echo "  outputs/corrected/pred_*.nc"
echo "  outputs/report/index.html"