#!/usr/bin/env bash
#
# Run the matched experiments required after the all-future first-onset update.
#
# Shared by every experiment:
#   - corrected first-onset likelihood and outcome-specific risk exposure;
#   - patient/interval/time-uniform query sampling in train/valid/test;
#   - timed disease history and smoking/alcohol/BMI extra information.
#
# Per seed, the script trains:
#   1. TrajMixer + relative Weibull + fitted risk baseline (primary model);
#   2. the primary model without the fitted risk baseline (baseline ablation);
#   3. TrajMixer + relative exponential (time-distribution control);
#   4. FFN + relative Weibull (architecture control).
#
# After training, AUC and calibration/point-process-NLL evaluation run
# automatically unless --train-only is supplied.

set -uo pipefail

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

GPU_CSV=""
SEED_CSV="42,43,44"
NUM_WORKERS=4
PYTHON_BIN="${PYTHON_BIN:-python}"
CAMPAIGN_NAME="all_future_first_onset_v2_multiseed"
DRY_RUN=0
TRAIN_ONLY=0

BATCH_SIZE=256
EXTRA_INFO_TYPES_FILE="$SCRIPT_DIR/extra_info_types_smoking_alcohol_bmi.txt"
TRAIN_EID_FILE="$SCRIPT_DIR/ukb_train_eid.csv"
VAL_EID_FILE="$SCRIPT_DIR/ukb_val_eid.csv"
TEST_EID_FILE="$SCRIPT_DIR/ukb_test_eid.csv"

usage() {
    cat <<'EOF'
Usage:
  bash run_all_future_first_onset_experiments_linux.sh --gpus LIST [options]

Required:
  --gpus LIST            Comma-separated GPU ids, for example 0 or 0,1,2,3.

Options:
  --seeds LIST           Comma-separated seeds (default: 42,43,44).
  --num-workers N        DataLoader workers per active GPU (default: 4).
  --python PATH          Python executable (default: $PYTHON_BIN or python).
  --campaign NAME        Output campaign directory name.
  --train-only           Skip AUC and calibration/NLL evaluation.
  --dry-run              Print commands without running them.
  -h, --help             Show this help message.

Fixed settings:
  batch_size             256
  extra information     smoking + alcohol + BMI
  patient split         fixed train/validation/test EID files
  disease history       timed
  experiments per seed  4

Outputs:
  runs/<campaign>/seed_<seed>/<architecture>/...
  batch_logs/<campaign>/seed_<seed>/<experiment>.log
EOF
}

while (($# > 0)); do
    case "$1" in
        --gpus)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --gpus requires a value." >&2
                exit 2
            }
            GPU_CSV="$2"
            shift 2
            ;;
        --seeds)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --seeds requires a value." >&2
                exit 2
            }
            SEED_CSV="$2"
            shift 2
            ;;
        --num-workers)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --num-workers requires a value." >&2
                exit 2
            }
            NUM_WORKERS="$2"
            shift 2
            ;;
        --python)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --python requires a value." >&2
                exit 2
            }
            PYTHON_BIN="$2"
            shift 2
            ;;
        --campaign)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --campaign requires a value." >&2
                exit 2
            }
            CAMPAIGN_NAME="$2"
            shift 2
            ;;
        --train-only)
            TRAIN_ONLY=1
            shift
            ;;
        --dry-run)
            DRY_RUN=1
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "ERROR: unknown argument: $1" >&2
            usage >&2
            exit 2
            ;;
    esac
done

[[ -n "$GPU_CSV" ]] || {
    echo "ERROR: --gpus is required." >&2
    usage >&2
    exit 2
}
[[ -n "$SEED_CSV" ]] || {
    echo "ERROR: --seeds must not be empty." >&2
    exit 2
}
[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || {
    echo "ERROR: --num-workers must be a non-negative integer." >&2
    exit 2
}
[[ "$CAMPAIGN_NAME" =~ ^[A-Za-z0-9._-]+$ ]] || {
    echo "ERROR: --campaign may contain only letters, numbers, ., _, and -." >&2
    exit 2
}

required_files=(
    "$SCRIPT_DIR/train_all_future.py"
    "$SCRIPT_DIR/evaluate_all_runs_linux.sh"
    "$SCRIPT_DIR/evaluate_calibration_all_runs_linux.sh"
    "$EXTRA_INFO_TYPES_FILE"
    "$TRAIN_EID_FILE"
    "$VAL_EID_FILE"
    "$TEST_EID_FILE"
)
for required_file in "${required_files[@]}"; do
    [[ -f "$required_file" ]] || {
        echo "ERROR: missing required file: $required_file" >&2
        exit 2
    }
done
command -v "$PYTHON_BIN" >/dev/null 2>&1 || {
    echo "ERROR: Python executable not found: $PYTHON_BIN" >&2
    exit 2
}

IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV"
declare -A SEEN_GPUS=()
for gpu in "${GPU_IDS[@]}"; do
    [[ -n "$gpu" && "$gpu" =~ ^[A-Za-z0-9._:-]+$ ]] || {
        echo "ERROR: invalid GPU id: $gpu" >&2
        exit 2
    }
    [[ -z "${SEEN_GPUS[$gpu]+x}" ]] || {
        echo "ERROR: duplicate GPU id: $gpu" >&2
        exit 2
    }
    SEEN_GPUS["$gpu"]=1
done

IFS=',' read -r -a SEEDS <<< "$SEED_CSV"
declare -A SEEN_SEEDS=()
for seed in "${SEEDS[@]}"; do
    [[ "$seed" =~ ^[0-9]+$ ]] || {
        echo "ERROR: invalid seed: $seed" >&2
        exit 2
    }
    [[ -z "${SEEN_SEEDS[$seed]+x}" ]] || {
        echo "ERROR: duplicate seed: $seed" >&2
        exit 2
    }
    SEEN_SEEDS["$seed"]=1
done

RUNS_ROOT="$SCRIPT_DIR/runs/$CAMPAIGN_NAME"
LOG_ROOT="$SCRIPT_DIR/batch_logs/$CAMPAIGN_NAME"
if ((DRY_RUN == 0)); then
    mkdir -p "$RUNS_ROOT" "$LOG_ROOT"
fi

declare -a JOB_NAMES=()
declare -a JOB_SEEDS=()
declare -a JOB_ARCHITECTURES=()
declare -a JOB_DIST_MODES=()
declare -a JOB_BASELINES=()

add_job() {
    JOB_NAMES+=("$1")
    JOB_SEEDS+=("$2")
    JOB_ARCHITECTURES+=("$3")
    JOB_DIST_MODES+=("$4")
    JOB_BASELINES+=("$5")
}

for seed in "${SEEDS[@]}"; do
    add_job \
        "traj_mixer_relative_weibull_first_onset" \
        "$seed" \
        "traj_mixer_v5" \
        "weibull" \
        "on"
    add_job \
        "traj_mixer_relative_weibull_no_rate_baseline" \
        "$seed" \
        "traj_mixer_v5" \
        "weibull" \
        "off"
    add_job \
        "traj_mixer_relative_exponential_first_onset" \
        "$seed" \
        "traj_mixer_v5" \
        "exponential" \
        "on"
    add_job \
        "ffn_relative_weibull_first_onset" \
        "$seed" \
        "transformer_ffn_v1" \
        "weibull" \
        "on"
done

print_command() {
    printf '%q ' "$@"
    printf '\n'
}

run_job() {
    local job_index="$1"
    local gpu="$2"
    local job_name="${JOB_NAMES[$job_index]}"
    local seed="${JOB_SEEDS[$job_index]}"
    local architecture="${JOB_ARCHITECTURES[$job_index]}"
    local dist_mode="${JOB_DIST_MODES[$job_index]}"
    local baseline="${JOB_BASELINES[$job_index]}"
    local seed_runs_root="$RUNS_ROOT/seed_$seed"
    local seed_log_root="$LOG_ROOT/seed_$seed"
    local log_file="$seed_log_root/$job_name.log"
    local -a command=(
        "$PYTHON_BIN"
        -u
        "$SCRIPT_DIR/train_all_future.py"
        --runs_root "$seed_runs_root"
        --seed "$seed"
        --batch_size "$BATCH_SIZE"
        --num_workers "$NUM_WORKERS"
        --device cuda
        --model_architecture "$architecture"
        --time_mode relative
        --dist_mode "$dist_mode"
        --disease_history_mode timed
        --extra_info_types_file "$EXTRA_INFO_TYPES_FILE"
        --train_eid_file "$TRAIN_EID_FILE"
        --val_eid_file "$VAL_EID_FILE"
        --test_eid_file "$TEST_EID_FILE"
    )

    if [[ "$baseline" == "on" ]]; then
        command+=(--risk_head_bias)
    else
        command+=(--no-risk_head_bias)
    fi

    if ((DRY_RUN == 0)); then
        mkdir -p "$seed_runs_root" "$seed_log_root"
    fi

    echo "[$(date '+%F %T')] START seed=$seed job=$job_name gpu=$gpu"
    echo "  log=$log_file"
    if ((DRY_RUN)); then
        printf '  CUDA_VISIBLE_DEVICES=%q ' "$gpu"
        print_command "${command[@]}"
        return 0
    fi

    if CUDA_VISIBLE_DEVICES="$gpu" PYTHONUNBUFFERED=1 \
        "${command[@]}" >"$log_file" 2>&1; then
        echo "[$(date '+%F %T')] DONE  seed=$seed job=$job_name gpu=$gpu"
        return 0
    else
        local exit_code=$?
        echo "[$(date '+%F %T')] FAIL  seed=$seed job=$job_name gpu=$gpu exit=$exit_code" >&2
        echo "  See: $log_file" >&2
        return "$exit_code"
    fi
}

worker() {
    local slot="$1"
    local gpu="${GPU_IDS[$slot]}"
    local job_index
    local failed=0

    for ((job_index = slot; job_index < ${#JOB_NAMES[@]}; job_index += ${#GPU_IDS[@]})); do
        run_job "$job_index" "$gpu" || failed=1
    done
    return "$failed"
}

echo "Campaign: $CAMPAIGN_NAME"
echo "Seeds: ${SEEDS[*]}"
echo "GPUs: ${GPU_IDS[*]}"
echo "Experiments per seed: 4"
echo "Total training tasks: ${#JOB_NAMES[@]}"
echo "Runs root: $RUNS_ROOT"
echo "Log root: $LOG_ROOT"
echo

declare -a WORKER_PIDS=()
for ((slot = 0; slot < ${#GPU_IDS[@]}; slot++)); do
    worker "$slot" &
    WORKER_PIDS+=("$!")
done

overall_status=0
for pid in "${WORKER_PIDS[@]}"; do
    wait "$pid" || overall_status=1
done

if ((overall_status != 0)); then
    echo "One or more training tasks failed. Inspect: $LOG_ROOT" >&2
    exit 1
fi

if ((TRAIN_ONLY == 0)); then
    evaluation_status=0
    auc_command=(
        bash
        "$SCRIPT_DIR/evaluate_all_runs_linux.sh"
        --gpus "$GPU_CSV"
        --runs-root "$RUNS_ROOT"
        --log-root "$LOG_ROOT/evaluate_auc"
        --python "$PYTHON_BIN"
        --num-workers "$NUM_WORKERS"
        --num-workers-auc "$NUM_WORKERS"
    )
    calibration_command=(
        bash
        "$SCRIPT_DIR/evaluate_calibration_all_runs_linux.sh"
        --gpus "$GPU_CSV"
        --runs-root "$RUNS_ROOT"
        --log-root "$LOG_ROOT/evaluate_calibration"
        --python "$PYTHON_BIN"
        --num-workers "$NUM_WORKERS"
        --num-workers-calibration "$NUM_WORKERS"
    )

    echo ">> AUC evaluation"
    print_command "${auc_command[@]}"
    if ((DRY_RUN == 0)); then
        "${auc_command[@]}" || evaluation_status=1
    fi

    echo ">> Calibration and point-process NLL evaluation"
    print_command "${calibration_command[@]}"
    if ((DRY_RUN == 0)); then
        "${calibration_command[@]}" || evaluation_status=1
    fi

    if ((evaluation_status != 0)); then
        echo "One or more evaluation workflows failed. Inspect: $LOG_ROOT" >&2
        exit 1
    fi
fi

if ((DRY_RUN)); then
    echo "Dry run completed successfully."
else
    echo "All required all-future experiments completed successfully."
fi