#!/usr/bin/env bash
# launch_multi.sh — Sequential multi-GPU training launcher for server deployment.
#
# Starts Karman Cloak training on multiple GPUs sequentially, with a configurable
# delay between launches to avoid CelerisLab kernel compilation race conditions.
#
# Usage:
#   bash launch_multi.sh --case-name re100_karman --seeds 42,43,44,45,46,47 \
#       --gpus 0,1,2,3,4,5 --episodes 500 \
#       --config configs/config_lbm_karman_2000x600.json \
#       --calibration calibrations/re100/calibration.json
#
#   # Transfer learning from a base model
#   bash launch_multi.sh --case-name re200_karman --seeds 42,43,44 \
#       --gpus 0,1,2 --episodes 500 \
#       --config configs/config_lbm_karman_2000x600_re200.json \
#       --calibration calibrations/re200/calibration.json \
#       --transfer output/re100_karman_seed42/models/best_model.zip
#
# Requirements:
#   - conda env pycuda_3_10
#   - CelerisLab submodule at DynamisLab/CelerisLab
#   - train_karman.py, env_karman.py, symmetry_wrapper.py in same directory
set -euo pipefail

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

# --- Defaults ---
CASE_NAME=""
SEEDS=""
GPUS=""
EPISODES=500
CONFIG=""
CALIBRATION=""
TRANSFER=""
DELAY_SECONDS=420  # 7 minutes between launches
CONDA_ENV="pycuda_3_10"

usage() {
    echo "Usage: $0 --case-name NAME --seeds S1,S2,... --gpus G1,G2,... [options]"
    echo ""
    echo "Required:"
    echo "  --case-name    NAME     Case name for output dirs"
    echo "  --seeds        S1,S2    Comma-separated seed values"
    echo "  --gpus         G1,G2    Comma-separated GPU device IDs"
    echo "  --config       PATH     LBM config JSON"
    echo "  --calibration  PATH     calibration.json"
    echo ""
    echo "Options:"
    echo "  --episodes     N        Total episodes (default: 500)"
    echo "  --transfer     PATH     .zip model for transfer learning"
    echo "  --delay        SEC      Seconds between launches (default: 420)"
    echo "  --env          NAME     Conda env name (default: pycuda_3_10)"
    exit 1
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --case-name)    CASE_NAME="$2"; shift 2 ;;
        --seeds)        SEEDS="$2"; shift 2 ;;
        --gpus)         GPUS="$2"; shift 2 ;;
        --episodes)     EPISODES="$2"; shift 2 ;;
        --config)       CONFIG="$2"; shift 2 ;;
        --calibration)  CALIBRATION="$2"; shift 2 ;;
        --transfer)     TRANSFER="$2"; shift 2 ;;
        --delay)        DELAY_SECONDS="$2"; shift 2 ;;
        --env)          CONDA_ENV="$2"; shift 2 ;;
        *)              echo "Unknown option: $1"; usage ;;
    esac
done

# Validate
if [[ -z "$CASE_NAME" || -z "$SEEDS" || -z "$GPUS" || -z "$CONFIG" || -z "$CALIBRATION" ]]; then
    echo "ERROR: Missing required arguments."
    usage
fi

IFS=',' read -ra SEED_ARR <<< "$SEEDS"
IFS=',' read -ra GPU_ARR <<< "$GPUS"

if [[ ${#SEED_ARR[@]} -ne ${#GPU_ARR[@]} ]]; then
    echo "ERROR: Number of seeds (${#SEED_ARR[@]}) must match number of GPUs (${#GPU_ARR[@]})."
    exit 1
fi

echo "=== Multi-GPU Training Launcher ==="
echo "  Case:       $CASE_NAME"
echo "  Seeds:      $SEEDS"
echo "  GPUs:       $GPUS"
echo "  Episodes:   $EPISODES"
echo "  Config:     $CONFIG"
echo "  Calibration: $CALIBRATION"
echo "  Transfer:   ${TRANSFER:-none}"
echo "  Delay:      ${DELAY_SECONDS}s between launches"
echo "  Jobs:       ${#SEED_ARR[@]}"
echo ""

# Build transfer arg
TRANSFER_ARG=""
if [[ -n "$TRANSFER" ]]; then
    TRANSFER_ARG="--transfer-model $TRANSFER"
fi

mkdir -p "$SCRIPT_DIR/output"

for i in "${!SEED_ARR[@]}"; do
    seed="${SEED_ARR[$i]}"
    gpu="${GPU_ARR[$i]}"
    run_name="${CASE_NAME}_seed${seed}"
    logfile="$SCRIPT_DIR/output/${run_name}/nohup.log"

    mkdir -p "$SCRIPT_DIR/output/${run_name}"

    echo "[$(date '+%H:%M:%S')] Launching seed=$seed on GPU=$gpu ..."
    nohup conda run --no-capture-output -n "$CONDA_ENV" python -u \
        "$SCRIPT_DIR/train_karman.py" \
        --case-name "$CASE_NAME" \
        --device-id "$gpu" \
        --seed "$seed" \
        --total-episodes "$EPISODES" \
        --config "$CONFIG" \
        --calibration "$CALIBRATION" \
        $TRANSFER_ARG \
        > "$logfile" 2>&1 &

    echo "  PID: $!"
    echo "  Log: $logfile"

    if [[ $i -lt $((${#SEED_ARR[@]} - 1)) ]]; then
        echo "  Waiting ${DELAY_SECONDS}s before next launch..."
        sleep "$DELAY_SECONDS"
    fi
done

echo ""
echo "All jobs launched. Monitor with:"
echo "  tail -f $SCRIPT_DIR/output/${CASE_NAME}_seed*/nohup.log"
echo "  tensorboard --logdir $SCRIPT_DIR/output/${CASE_NAME}_seed*/tb"
echo ""
echo "To stop all:"
echo "  ps aux | grep train_karman | grep -v grep | awk '{print \$2}' | xargs kill"