#!/bin/bash
# Quick start script to launch A* paper workflow
# Run this to start Phase A immediately

set -euo pipefail

echo "╔════════════════════════════════════════════════════════════════════╗"
echo "║  DoVLA-CIL → A* Oral Paper                                         ║"
echo "║  Quick Start: Phase A Performance Improvement                      ║"
echo "╚════════════════════════════════════════════════════════════════════╝"
echo ""

cd "$(dirname "$0")/.."
PROJECT_DIR="$PWD"

echo "📍 Project directory: $PROJECT_DIR"
echo ""

# Pre-flight checks
echo "🔍 Pre-flight checks..."
echo ""

# Check venv
if [ ! -d ".venv/bin" ]; then
    echo "❌ Virtual environment not found at .venv/"
    echo "   Run: python -m venv .venv && source .venv/bin/activate && pip install -e ."
    exit 1
fi
echo "✅ Virtual environment found"

# Check demo files
DEMO_DIR="/scratch/$USER/dovla/demonstrations/maniskill"
if [ ! -d "$DEMO_DIR" ]; then
    echo "⚠️  Demo directory not found: $DEMO_DIR"
    echo "   Phase A1 may need adjustment"
else
    NUM_DEMOS=$(ls "$DEMO_DIR"/*.h5 2>/dev/null | wc -l)
    echo "✅ Demo directory found ($NUM_DEMOS .h5 files)"
fi

# Check logs directory
mkdir -p logs/workflow
echo "✅ Logs directory ready: logs/workflow/"
echo ""

# Show what will run
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "PHASE A: PERFORMANCE IMPROVEMENT"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo ""
echo "Will execute:"
echo "  A1. Generate 10K groups    (~3-4 days, 20 GPU hours)"
echo "  A2. Train large model      (~2-3 days, 90 GPU hours, 3 seeds)"
echo "  A3. Evaluate               (~1 day, 6 GPU hours)"
echo "  A4. Hyperparameter sweep   (~2-3 days, 45 GPU hours, 9 configs)"
echo "  A5. Horizon sweep          (~1-2 days, 16 GPU hours, 4 configs)"
echo ""
echo "Total Phase A:"
echo "  Time: ~2 weeks"
echo "  Compute: ~180 GPU hours"
echo "  Target: 40%+ policy success (vs 29.67% baseline)"
echo ""

# Ask for confirmation
read -p "🚀 Launch Phase A now? [y/N] " -n 1 -r
echo
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    echo "Cancelled. No jobs submitted."
    exit 0
fi

echo ""
echo "🚀 Launching Phase A..."
echo ""

# Option 1: Submit just A1 first (safer)
echo "Strategy: Submit A1 first, then A2-A5 after A1 completes"
echo ""

# Submit Phase A1
echo "Submitting Phase A1: 10K generation..."
JOB_A1=$(sbatch scripts/slurm/phase_a1_generate_10k.sbatch | awk '{print $NF}')

if [ -z "$JOB_A1" ]; then
    echo "❌ Failed to submit Phase A1"
    exit 1
fi

echo "✅ Phase A1 submitted: Job ID $JOB_A1"
echo ""

# Show monitoring commands
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "MONITORING"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo ""
echo "Check job status:"
echo "  squeue -u $USER"
echo ""
echo "Monitor Phase A1 log:"
echo "  tail -f logs/phase_a_10k_gen_${JOB_A1}.out"
echo ""
echo "After A1 completes (~3-4 days), submit Phase A2:"
echo "  sbatch scripts/slurm/phase_a2_train_large_model.sbatch"
echo ""
echo "Or use master workflow for full automation:"
echo "  bash scripts/run_master_workflow.sh"
echo ""

# Save job info
echo "$JOB_A1" > logs/workflow/phase_a1_job_id.txt
echo "$(date +'%Y-%m-%d %H:%M:%S')" > logs/workflow/phase_a_start_time.txt

echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "✅ PHASE A LAUNCHED"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo ""
echo "Job ID: $JOB_A1"
echo "Expected completion: 3-4 days"
echo "Next: Monitor progress and submit A2 after A1 completes"
echo ""
echo "See: LAUNCH_READY.md for full instructions"
echo "See: WORKFLOW_A_STAR.md for complete 8-week plan"