#!/bin/bash
# Master orchestration script for A* paper workflow
# Executes all phases in optimal order

set -euo pipefail

PROJECT_DIR="${PROJECT_DIR:-$PWD}"
cd "$PROJECT_DIR"

LOG_DIR="$PROJECT_DIR/logs/workflow"
mkdir -p "$LOG_DIR"

WORKFLOW_LOG="$LOG_DIR/master_workflow_$(date +%Y%m%d_%H%M%S).log"

log() {
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] $*" | tee -a "$WORKFLOW_LOG"
}

check_job() {
    local JOB_ID=$1
    squeue -j "$JOB_ID" &>/dev/null
}

wait_for_job() {
    local JOB_ID=$1
    local JOB_NAME=$2

    log "Waiting for $JOB_NAME (Job ID: $JOB_ID)..."

    while check_job "$JOB_ID"; do
        sleep 60
    done

    log "$JOB_NAME completed (Job ID: $JOB_ID)"
}

submit_and_wait() {
    local SBATCH_SCRIPT=$1
    local JOB_NAME=$2

    log "Submitting $JOB_NAME: $SBATCH_SCRIPT"

    JOB_ID=$(sbatch "$SBATCH_SCRIPT" | awk '{print $NF}')

    if [ -z "$JOB_ID" ]; then
        log "ERROR: Failed to submit $JOB_NAME"
        return 1
    fi

    log "$JOB_NAME submitted with Job ID: $JOB_ID"

    wait_for_job "$JOB_ID" "$JOB_NAME"

    echo "$JOB_ID"
}

log "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
log "DoVLA-CIL A* Paper Workflow - Master Orchestration"
log "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
log ""
log "Target: A* oral paper with 9/10 novelty"
log "Timeline: 6-8 weeks"
log "Compute: ~250-350 GPU hours"
log ""

# Check if running in dry-run mode
DRY_RUN="${DRY_RUN:-0}"

if [ "$DRY_RUN" = "1" ]; then
    log "🔍 DRY RUN MODE - No jobs will be submitted"
    log ""
fi

# ============================================================================
# PHASE A: PERFORMANCE IMPROVEMENT (WEEK 1-2)
# Critical: 30% → 40%+ policy success
# ============================================================================

log "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
log "PHASE A: PERFORMANCE IMPROVEMENT"
log "Target: 40%+ policy success (vs 29.67% baseline)"
log "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
log ""

# A1: Generate 10K dataset
log "Phase A1: Generate 10K group dataset"
log "  Expected: 3-4 days, ~20 GPU hours"
log "  Output: /scratch/$USER/dovla/experiments/phase_a_10k_collection"
log ""

if [ "$DRY_RUN" = "0" ]; then
    PHASE_A1_JOB=$(submit_and_wait \
        "scripts/slurm/phase_a1_generate_10k.sbatch" \
        "Phase A1: 10K Generation")

    log "✅ Phase A1 complete"
    log ""
else
    log "  [DRY RUN] Would submit: scripts/slurm/phase_a1_generate_10k.sbatch"
    log ""
fi

# Check if 10K dataset exists
DATASET_10K="/scratch/$USER/dovla/experiments/phase_a_10k_collection/merged_10k"
if [ ! -d "$DATASET_10K" ] && [ "$DRY_RUN" = "0" ]; then
    log "ERROR: 10K dataset not found at $DATASET_10K"
    exit 1
fi

# A2: Train large model (3 seeds)
log "Phase A2: Train large capacity model (3 seeds)"
log "  Expected: 2-3 days, ~30 GPU hours per seed"
log "  Config: hidden_dim=512, 100 epochs"
log ""

if [ "$DRY_RUN" = "0" ]; then
    PHASE_A2_JOB=$(submit_and_wait \
        "scripts/slurm/phase_a2_train_large_model.sbatch" \
        "Phase A2: Large Model Training")

    log "✅ Phase A2 complete (3 seeds trained)"
    log ""
else
    log "  [DRY RUN] Would submit: scripts/slurm/phase_a2_train_large_model.sbatch"
    log ""
fi

# A3: Evaluate large model
log "Phase A3: Evaluate large model"
log "  Lattice eval + policy rollout on 700 held-out groups"
log ""

if [ "$DRY_RUN" = "0" ]; then
    PHASE_A3_JOB=$(submit_and_wait \
        "scripts/slurm/phase_a3_eval_large_model.sbatch" \
        "Phase A3: Large Model Eval")

    log "✅ Phase A3 complete"
    log ""
else
    log "  [DRY RUN] Would submit: scripts/slurm/phase_a3_eval_large_model.sbatch"
    log ""
fi

# A4 & A5: Parallel sweeps (optional but recommended)
log "Phase A4 & A5: Hyperparameter and horizon sweeps (parallel)"
log "  A4: 9 configs (3 LR × 3 hidden_dim)"
log "  A5: 4 horizons (H=4,8,12,16)"
log ""

if [ "$DRY_RUN" = "0" ]; then
    # Submit both in parallel
    PHASE_A4_JOB=$(sbatch scripts/slurm/phase_a4_hparam_sweep.sbatch | awk '{print $NF}')
    PHASE_A5_JOB=$(sbatch scripts/slurm/phase_a5_horizon_sweep.sbatch | awk '{print $NF}')

    log "Phase A4 submitted: Job $PHASE_A4_JOB"
    log "Phase A5 submitted: Job $PHASE_A5_JOB"

    # Wait for both
    wait_for_job "$PHASE_A4_JOB" "Phase A4: Hyperparameter Sweep"
    wait_for_job "$PHASE_A5_JOB" "Phase A5: Horizon Sweep"

    log "✅ Phase A4 & A5 complete"
    log ""
else
    log "  [DRY RUN] Would submit parallel:"
    log "    scripts/slurm/phase_a4_hparam_sweep.sbatch"
    log "    scripts/slurm/phase_a5_horizon_sweep.sbatch"
    log ""
fi

log "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
log "PHASE A: COMPLETE"
log "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
log ""
log "Next: Analyze Phase A results and proceed to Phase B"
log ""

# ============================================================================
# CHECKPOINT: Analyze Phase A results
# ============================================================================

log "Analyzing Phase A results..."
log ""

if [ "$DRY_RUN" = "0" ]; then
    python scripts/analyze_phase_a_results.py \
        --baseline /scratch/$USER/dovla/experiments/six_task_state_actionfix \
        --large-model /scratch/$USER/dovla/experiments/phase_a2_large_model \
        --hparam-sweep /scratch/$USER/dovla/experiments/phase_a4_hparam_sweep \
        --horizon-sweep /scratch/$USER/dovla/experiments/phase_a5_horizon_sweep \
        --out reports/phase_a_final_results.json

    # Check if we hit target
    BEST_SUCCESS=$(python -c "import json; print(json.load(open('reports/phase_a_final_results.json'))['best_policy_success'])")

    log "Phase A best policy success: $BEST_SUCCESS"

    if (( $(echo "$BEST_SUCCESS < 0.40" | bc -l) )); then
        log "⚠️  WARNING: Target 40% not reached (got $BEST_SUCCESS)"
        log "   Consider additional iterations or adjustments"
    else
        log "✅ Target 40%+ achieved!"
    fi
    log ""
fi

# ============================================================================
# PHASE B: SECOND BENCHMARK (WEEK 3-4)
# Critical for generality claim
# ============================================================================

log "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
log "PHASE B: SECOND BENCHMARK"
log "Target: Demonstrate generality beyond ManiSkill"
log "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
log ""

log "⚠️  Phase B requires manual implementation:"
log ""
log "Option 1 (RECOMMENDED): Meta-World"
log "  1. pip install metaworld"
log "  2. Complete scripts/generate_metaworld_lattice.py"
log "  3. Adapt 5-6 Meta-World tasks"
log "  Estimated effort: 2-3 days"
log ""
log "Option 2: More ManiSkill tasks"
log "  1. Expand from 6 to 12 ManiSkill tasks"
log "  2. Use existing infrastructure"
log "  Estimated effort: 1-2 days (faster but less impressive)"
log ""
log "Option 3: RLBench"
log "  1. Install RLBench"
log "  2. Implement CIL generation"
log "  Estimated effort: 3-4 days (more impressive but slower)"
log ""

if [ "$DRY_RUN" = "0" ]; then
    log "Pausing workflow - complete Phase B implementation manually"
    log ""
    log "After Phase B is ready, continue with:"
    log "  bash scripts/continue_workflow_from_phase_c.sh"
else
    log "[DRY RUN] Phase B would require manual implementation"
fi

log ""
log "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
log "WORKFLOW STATUS: Paused at Phase B"
log "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
log ""
log "Phase A: ✅ Complete"
log "Phase B: ⏳ Awaiting implementation"
log "Phase C: ⏳ Pending"
log "Phase D: ⏳ Pending"
log "Phase E: ⏳ Pending"
log ""
log "Estimated timeline to completion:"
log "  Phase B: +1-2 weeks (implementation + experiments)"
log "  Phase C+D: +2 weeks (transfer + online rollout)"
log "  Phase E: +1 week (12-task scale)"
log "  Paper writing: +1 week"
log "  Total: 6-8 weeks from today"
log ""
log "See: WORKFLOW_A_STAR.md for detailed instructions"
log "Workflow log: $WORKFLOW_LOG"