#!/bin/bash
#SBATCH --job-name=dovla_12task_gen
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=16
#SBATCH --gres=gpu:1
#SBATCH --mem=64000M
#SBATCH --time=72:00:00
#SBATCH --output=logs/phase_b_12task_gen_%j.out
#SBATCH --error=logs/phase_b_12task_gen_%j.err

set -euo pipefail

# Phase B Option 1: Generate 12-task ManiSkill collection
# Fastest option - uses existing infrastructure

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

source .venv/bin/activate

OUT_DIR="/scratch/$USER/dovla/experiments/phase_b_12task_collection"
K=16
STATE_BATCH_SIZE=16

# 12 tasks: 6 existing + 6 new
declare -A TASK_GROUPS=(
    # Original 6
    ["PickCube-v1"]=800
    ["PushCube-v1"]=800
    ["PullCube-v1"]=600
    ["StackCube-v1"]=600
    ["LiftPegUpright-v1"]=600
    ["PegInsertionSide-v1"]=600

    # New 6 (TODO: ensure demos exist)
    ["TurnFaucet-v1"]=500
    ["OpenDrawer-v1"]=500
    ["CloseDrawer-v1"]=500
    ["PlugCharger-v1"]=400
    ["HangMug-v1"]=400
    ["PourWater-v1"]=400
)

mkdir -p "$OUT_DIR" logs

echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
echo "Phase B Option 1: 12-Task ManiSkill Collection"
echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
echo ""
echo "Target: 6,200 groups, 99,200 records (K=$K)"
echo "Strategy: Expand existing ManiSkill tasks"
echo ""

# Check if this is just a planning run
if [ "${DRY_RUN:-0}" = "1" ]; then
    echo "DRY RUN: Would generate 12 tasks"
    for TASK in "${!TASK_GROUPS[@]}"; do
        echo "  $TASK: ${TASK_GROUPS[$TASK]} groups"
    done
    exit 0
fi

# Generate each task
for TASK in "${!TASK_GROUPS[@]}"; do
    NUM_GROUPS="${TASK_GROUPS[$TASK]}"

    # Check if already generated
    if [ -d "$OUT_DIR/${TASK}_k${K}_n${NUM_GROUPS}/merged" ]; then
        echo "✓ $TASK already exists, skipping"
        continue
    fi

    echo "Generating $TASK: $NUM_GROUPS groups..."

    # Use existing generation script
    python scripts/generate_maniskill_lattice.py \
        --env-id "$TASK" \
        --control-mode pd_ee_delta_pose \
        --out "$OUT_DIR/${TASK}_k${K}_n${NUM_GROUPS}" \
        --num-groups "$NUM_GROUPS" \
        --k "$K" \
        --state-batch-size "$STATE_BATCH_SIZE" \
        --seed 42 \
        --pre-success-only \
        --use-official-demos || {
            echo "⚠️  $TASK failed (demo might not exist)"
            continue
        }

    echo "✅ $TASK complete"
    echo ""
done

echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
echo "Merging into unified 12-task collection"
echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="

python scripts/make_cil_collection.py \
    --source-dirs "$OUT_DIR"/*/merged \
    --out "$OUT_DIR/merged_12tasks" \
    --name "phase_b_12task_collection"

echo ""
echo "✅ Phase B Option 1 complete: 12-task collection ready"
echo "   Location: $OUT_DIR/merged_12tasks"
echo ""
echo "Next: Train on 12 tasks"
echo "  sbatch scripts/slurm/phase_b_train_12tasks.sbatch"