#!/bin/bash
# run_all.sh — Master launcher for eval benchmark.
#
# GPU 0: Train pipeline (7 scenes)
# GPU 1: Reproduce pipeline (7 scenes, 120s stagger)
#
# Each pipeline saves outputs to eval/output/{train,reproduce}/.
# After both complete, run generate_report.py and viz_*.py.
#
# Usage:
#   conda activate pycuda_3_10
#   bash run_all.sh
#
# Single scene mode:
#   bash run_all.sh --train-scene re100 --gpu 0
#   bash run_all.sh --repro-scene karman_cloak --gpu 1
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_DIR="$(cd "$SCRIPT_DIR/../../.." && pwd)"
TRAIN_SCRIPT="$SCRIPT_DIR/infer_train.py"
REPRO_SCRIPT="$SCRIPT_DIR/infer_reproduce.py"
CONDA_ENV="pycuda_3_10"

TRAIN_GPU=0
REPRO_GPU=1
TRAIN_SCENE=""
REPRO_SCENE=""

usage() {
    echo "Usage: $0 [--train-scene SCENE] [--repro-scene SCENE] [--gpu N]"
    echo "  --train-scene  Run single train scene (e.g. re100, illusion_1L)"
    echo "  --repro-scene  Run single reproduce scene (e.g. karman_cloak)"
    echo "  --gpu          GPU device ID"
    exit 1
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --train-scene) TRAIN_SCENE="$2"; shift 2 ;;
        --repro-scene) REPRO_SCENE="$2"; shift 2 ;;
        --gpu)         TRAIN_GPU="$2"; REPRO_GPU="$2"; shift 2 ;;
        *)             echo "Unknown option: $1"; usage ;;
    esac
done

# Create output dirs
mkdir -p "$SCRIPT_DIR/output/train"
mkdir -p "$SCRIPT_DIR/output/reproduce"

echo "=== Eval Benchmark ==="
echo "  Train GPU:  $TRAIN_GPU"
echo "  Reproduce GPU: $REPRO_GPU"
echo "  Conda env:  $CONDA_ENV"
echo "  Script dir: $SCRIPT_DIR"
echo ""

# ---------------------------------------------------------------------------
# Launch Reproduce pipeline (GPU 1, staggered)
# ---------------------------------------------------------------------------
launch_reproduce() {
    echo "[$(date '+%H:%M:%S')] Launching Reproduce pipeline on GPU $REPRO_GPU..."
    local arg=""
    if [[ -n "$REPRO_SCENE" ]]; then
        arg="--scene $REPRO_SCENE"
    fi
    conda run --no-capture-output -n "$CONDA_ENV" python -u \
        "$REPRO_SCRIPT" --device-id "$REPRO_GPU" $arg \
        2>&1 | tee "$SCRIPT_DIR/output/reproduce/run.log"
    echo "[$(date '+%H:%M:%S')] Reproduce pipeline complete."
}

# Start reproduce after 120s delay (avoids kernel compilation race with GPU 0)
if [[ -z "$TRAIN_SCENE" ]] && [[ -z "$REPRO_SCENE" ]]; then
    # Full run: launch reproduce in background with delay
    (sleep 120 && launch_reproduce) &
    REPRO_PID=$!
    echo "  Reproduce scheduled in 120s (PID=$REPRO_PID)"
else
    # Single scene mode: just run the requested pipeline
    if [[ -n "$REPRO_SCENE" ]]; then
        launch_reproduce
        exit 0
    fi
fi

# ---------------------------------------------------------------------------
# Train pipeline (GPU 0)
# ---------------------------------------------------------------------------
echo "[$(date '+%H:%M:%S')] Launching Train pipeline on GPU $TRAIN_GPU..."
TRAIN_ARG=""
if [[ -n "$TRAIN_SCENE" ]]; then
    TRAIN_ARG="--scene $TRAIN_SCENE"
fi
conda run --no-capture-output -n "$CONDA_ENV" python -u \
    "$TRAIN_SCRIPT" --device-id "$TRAIN_GPU" $TRAIN_ARG \
    2>&1 | tee "$SCRIPT_DIR/output/train/run.log"
echo "[$(date '+%H:%M:%S')] Train pipeline complete."

# Wait for reproduce if it was launched
if [[ -n "${REPRO_PID:-}" ]]; then
    echo "[$(date '+%H:%M:%S')] Waiting for Reproduce pipeline (PID=$REPRO_PID)..."
    wait $REPRO_PID
fi

# ---------------------------------------------------------------------------
# Generate reports
# ---------------------------------------------------------------------------
echo ""
echo "[$(date '+%H:%M:%S')] Both pipelines complete. Generating reports..."

echo "  -> Signal plots..."
conda run --no-capture-output -n "$CONDA_ENV" python -u \
    "$SCRIPT_DIR/viz_signals.py" 2>&1 | tail -5

echo "  -> Vorticity panels..."
conda run --no-capture-output -n "$CONDA_ENV" python -u \
    "$SCRIPT_DIR/viz_flow.py" 2>&1 | tail -5

echo "  -> Summary report..."
conda run --no-capture-output -n "$CONDA_ENV" python -u \
    "$SCRIPT_DIR/generate_report.py" 2>&1 | tail -10

echo ""
echo "=== ALL DONE ==="
echo "Reports in: $SCRIPT_DIR/output/reports/"
