#!/bin/bash
#SBATCH --job-name=monitor_h16_correct
#SBATCH --account=def-yalda
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --mem=2G
#SBATCH --output=logs/monitor_h16_correct_%j.out
#SBATCH --error=logs/monitor_h16_correct_%j.err

# Monitor DoVLAModel h=16 training (14763330) → eval → paper
# CORRECTED: Now watches DoVLAModel (rollout-capable) not DoVLAHybrid

set -euo pipefail

TRAIN_JOB_ID=14763330
PROJECT_DIR="/lustre09/project/6037638/knguy52/vla"
PYTHON="$PROJECT_DIR/.venv/bin/python"

cd "$PROJECT_DIR"

echo "=== Monitor for DoVLAModel h=16 Training ==="
echo "Training job: $TRAIN_JOB_ID"
echo "Start: $(date)"
echo ""

# Function to check if all array tasks completed
check_training_complete() {
    local states=$(sacct -j $1 --format=State --noheader 2>/dev/null | grep -v "^$")
    local total=$(echo "$states" | wc -l)
    local completed=$(echo "$states" | grep -c "COMPLETED" || true)

    if [ "$total" -ge 3 ] && [ "$completed" -eq 3 ]; then
        echo "COMPLETED"
    elif echo "$states" | grep -q "FAILED\|CANCELLED\|TIMEOUT"; then
        echo "FAILED"
    else
        echo "RUNNING"
    fi
}

# Wait for training to complete
while true; do
    STATUS=$(check_training_complete $TRAIN_JOB_ID)
    echo "[$(date +'%H:%M:%S')] Training status: $STATUS"

    if [ "$STATUS" = "COMPLETED" ]; then
        echo ""
        echo "✅ Training completed! Verifying checkpoints..."
        echo ""

        # Verify checkpoints have model_config (rollout-compatible)
        $PYTHON << 'PYEOF'
import torch
from pathlib import Path

run_dir = Path("/scratch/knguy52/dovla/experiments/dovla_h16_rollout_runs")
seeds = [0, 1, 2]
checkpoints_ok = []

for seed in seeds:
    ckpt_path = run_dir / f"seed_{seed}" / "best.pt"
    if not ckpt_path.exists():
        print(f"❌ Checkpoint missing: seed {seed}")
        continue

    ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False)
    has_model_config = 'model_config' in ckpt

    print(f"Seed {seed}: {'✅' if has_model_config else '❌'} model_config present")

    if has_model_config:
        checkpoints_ok.append(seed)

if len(checkpoints_ok) == 3:
    print("")
    print("✅ All 3 checkpoints are rollout-compatible (have model_config)")
    print("   Ready for online evaluation")
    Path("results/.trigger_h16_evaluation").touch()
else:
    print("")
    print(f"⚠️  Only {len(checkpoints_ok)}/3 checkpoints OK")
    exit(1)
PYEOF

        # Submit evaluation if triggered
        if [ -f "results/.trigger_h16_evaluation" ]; then
            echo ""
            echo "Submitting online rollout evaluation..."

            # Update eval sbatch to use correct checkpoints
            sed -i 's|h16_policy_runs|dovla_h16_rollout_runs|g' scripts/slurm/eval_h16_rollout.sbatch

            EVAL_JOB=$(sbatch scripts/slurm/eval_h16_rollout.sbatch | awk '{print $4}')
            echo "Submitted eval job: $EVAL_JOB"

            # Submit monitor for eval results
            sed "s/EVAL_JOB_ID=.*/EVAL_JOB_ID=$EVAL_JOB/" scripts/slurm/monitor_eval.sbatch | \
                sbatch --job-name=monitor_eval_h16

            echo "✅ Evaluation and monitoring pipeline launched"
        fi

        echo ""
        echo "=== Training Monitor Complete ==="
        exit 0

    elif [ "$STATUS" = "FAILED" ]; then
        echo ""
        echo "❌ Training failed"
        sacct -j $TRAIN_JOB_ID --format=JobID,State,ExitCode,Reason
        exit 1
    fi

    # Sleep 10 minutes before next check
    sleep 600
done