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

# Monitor eval 14775756 → parse → assess → generate paper if warranted
# HONEST: Only claim what data supports

set -euo pipefail

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

cd "$PROJECT_DIR"

echo "=== Final Evaluation Monitor ==="
echo "Job: $EVAL_JOB_ID"
echo "Start: $(date)"
echo ""

check_eval_complete() {
    local states=$(sacct -j $1 --format=State --noheader 2>/dev/null)
    local completed=$(echo "$states" | grep -c "COMPLETED" || true)

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

while true; do
    STATUS=$(check_eval_complete $EVAL_JOB_ID)
    echo "[$(date +'%H:%M:%S')] Eval status: $STATUS"

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

        $PYTHON << 'PYEOF'
import json
from pathlib import Path
import statistics

results_dir = Path("/scratch/knguy52/dovla/experiments/dovla_h16_rollout_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
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
oracle_h16 = 0.9476

print("="*60)
print("📊 HONEST EVALUATION RESULTS (DoVLAModel h=16)")
print("="*60)
print(f"Policy Success Rate: {mean_success:.2%} ± {std_success:.2%}")
print(f"Baseline (h=4):      {baseline:.2%}")
print(f"Oracle (h=16):       {oracle_h16:.2%}")
print("")
print(f"Absolute Gain:       {(mean_success - baseline):+.2%}")
print(f"Relative Gain:       {(mean_success / baseline):.2f}×")
print(f"% of Oracle Reached: {(mean_success / oracle_h16):.1%}")
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,
    'oracle_h16': oracle_h16,
    'absolute_gain': mean_success - baseline,
    'relative_gain': mean_success / baseline,
    'oracle_fraction': mean_success / oracle_h16,
    '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_final_evaluation.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}")
print("")

# HONEST ASSESSMENT
print("="*60)
print("HONEST ASSESSMENT FOR PAPER")
print("="*60)

publishable = False
story = ""

if mean_success >= 0.50:
    print("✅ STRONG RESULT (≥50%)")
    print("   Paper story: 2× improvement, SOTA-competitive")
    publishable = True
    story = "strong"
elif mean_success >= 0.40:
    print("✅ GOOD RESULT (40-50%)")
    print("   Paper story: Significant improvement, horizon matters")
    publishable = True
    story = "good"
elif mean_success >= 0.35:
    print("⚠️  MODEST RESULT (35-40%)")
    print("   Paper story: Partial improvement, diagnostic value")
    print("   Publishable but needs careful framing")
    publishable = True
    story = "modest"
else:
    print("⚠️  BELOW EXPECTATIONS (<35%)")
    print("   Gap between oracle (94%) and policy suggests:")
    print("   - Longer horizons harder to predict accurately")
    print("   - Or training/architecture mismatch")
    print("   Still publishable as negative/diagnostic result")
    publishable = True
    story = "diagnostic"

print("")
print(f"Publishable: {publishable}")
print(f"Story angle: {story}")
print("")

# Save assessment
assessment = {
    'publishable': publishable,
    'story': story,
    'mean_success': mean_success,
    'expected_range': [0.35, 0.55],
    'in_range': 0.35 <= mean_success <= 0.55
}

Path("results/paper_assessment.json").write_text(json.dumps(assessment, indent=2))

if publishable:
    print("✅ Triggering paper generation...")
    Path("results/.trigger_paper_generation").touch()
else:
    print("⚠️  Results need analysis before paper")

PYEOF

        # Upload results
        $PYTHON -c "
from huggingface_hub import upload_file
try:
    upload_file(
        path_or_fileobj='results/h16_final_evaluation.json',
        path_in_repo='results/h16_final_evaluation.json',
        repo_id='anhtld/vla',
        commit_message='DoVLAModel h=16 evaluation results (honest measurement)'
    )
    print('✅ Results uploaded to HF')
except Exception as e:
    print(f'⚠️  Upload: {e}')
"

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

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

    # Check every 10 minutes
    sleep 600
done