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

# Autonomous monitor: Check eval job → parse results → trigger paper writing
# Runs on compute node, NOT login node

set -euo pipefail

EVAL_JOB_ID=14758888
PROJECT_DIR="/lustre09/project/6037638/knguy52/vla"
PYTHON="$PROJECT_DIR/.venv/bin/python"

cd "$PROJECT_DIR"

echo "=== Autonomous Evaluation Monitor Started ==="
echo "Watching job: $EVAL_JOB_ID"
echo "Start time: $(date)"
echo ""

# Function to check job status
check_job_status() {
    sacct -j $1 --format=State --noheader -X 2>/dev/null | head -1 | awk '{print $1}'
}

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

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

# Monitor loop
while true; do
    STATUS=$(check_array_complete $EVAL_JOB_ID)
    echo "[$(date +'%H:%M:%S')] Job status: $STATUS"

    if [ "$STATUS" = "COMPLETED" ]; then
        echo ""
        echo "✅ Evaluation completed! Processing results..."
        echo ""

        # Parse results from all 3 seeds
        $PYTHON << 'PYEOF'
import json
from pathlib import Path

results_dir = Path("/scratch/knguy52/dovla/experiments/h16_policy_runs")
seeds = [0, 1, 2]
all_results = []

for seed in seeds:
    result_file = results_dir / f"seed_{seed}" / "online_rollout.json"
    if result_file.exists():
        with open(result_file) as f:
            data = json.load(f)
            all_results.append({
                'seed': seed,
                'policy_success': data.get('policy_rollout_success_rate', 0),
                'per_task': data.get('per_task', {})
            })

if not all_results:
    print("❌ No results found!")
    exit(1)

# Compute statistics
import statistics
success_rates = [r['policy_success'] for r in all_results]
mean_success = statistics.mean(success_rates)
std_success = statistics.stdev(success_rates) if len(success_rates) > 1 else 0

baseline = 0.2967

print("="*60)
print("📊 EVALUATION RESULTS")
print("="*60)
print(f"Policy Success Rate: {mean_success:.2%} ± {std_success:.2%}")
print(f"Baseline (h=4):      {baseline:.2%}")
print(f"Absolute Gain:       +{(mean_success - baseline):.2%}")
print(f"Relative Gain:       {(mean_success / baseline):.2f}×")
print("")

# Per-task breakdown
print("Per-Task Breakdown:")
task_names = set()
for r in all_results:
    task_names.update(r['per_task'].keys())

for task in sorted(task_names):
    rates = [r['per_task'][task]['policy_rollout_success_rate']
             for r in all_results if task in r['per_task']]
    if rates:
        mean_rate = statistics.mean(rates)
        print(f"  {task:25s} {mean_rate:6.2%}")

print("="*60)

# Save summary
summary = {
    'mean_success_rate': mean_success,
    'std_success_rate': std_success,
    'baseline': baseline,
    'absolute_gain': mean_success - baseline,
    'relative_gain': mean_success / baseline,
    'per_task_mean': {
        task: statistics.mean([r['per_task'][task]['policy_rollout_success_rate']
                               for r in all_results if task in r['per_task']])
        for task in task_names
    },
    'seeds': all_results
}

summary_path = Path("results/h16_evaluation_summary.json")
summary_path.parent.mkdir(parents=True, exist_ok=True)
with open(summary_path, 'w') as f:
    json.dump(summary, f, indent=2)

print(f"Summary saved: {summary_path}")

# Trigger paper writing if results are good
if mean_success >= 0.55:
    print("")
    print("✅ Results meet A* threshold (≥55%)!")
    print("   Triggering paper writing workflow...")
    Path("results/.trigger_paper_writing").touch()
else:
    print("")
    print("⚠️  Results below target. Analysis needed.")
PYEOF

        # Submit paper writing job if triggered
        if [ -f "results/.trigger_paper_writing" ]; then
            echo ""
            echo "Submitting paper writing job..."
            sbatch scripts/slurm/write_paper_draft.sbatch
        fi

        # Upload results to HF
        echo ""
        echo "Uploading results to HuggingFace..."
        $PYTHON -c "
from huggingface_hub import upload_file
import sys

try:
    upload_file(
        path_or_fileobj='results/h16_evaluation_summary.json',
        path_in_repo='results/h16_evaluation_summary.json',
        repo_id='anhtld/vla',
        commit_message='Add h=16 evaluation results (THE decisive number)'
    )
    print('✅ Results uploaded to HF')
except Exception as e:
    print(f'⚠️  Upload failed: {e}', file=sys.stderr)
"

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

    elif [ "$STATUS" = "FAILED" ]; then
        echo ""
        echo "❌ Evaluation failed. Checking logs..."
        sacct -j $EVAL_JOB_ID --format=JobID,State,ExitCode,Reason
        echo ""
        echo "Check logs in: outputs/hpc/logs/eval_h16_rollout_${EVAL_JOB_ID}_*.{out,err}"
        exit 1
    fi

    # Sleep 5 minutes before next check
    sleep 300
done