#!/bin/bash
#SBATCH --job-name=dovla_enhanced_train
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --mem=64000M
#SBATCH --time=120:00:00
#SBATCH --output=logs/phase_a1b_enhanced_train_%A_%a.out
#SBATCH --error=logs/phase_a1b_enhanced_train_%A_%a.err
#SBATCH --array=0-2

set -euo pipefail

# Phase A1b: Enhanced Training for 50%+ Target

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

source .venv/bin/activate

DATASET="/scratch/$USER/dovla/experiments/phase_a1_10k_collection/merged_10k"
OUT_DIR="/scratch/$USER/dovla/experiments/phase_a1b_enhanced_model"
SEED=$SLURM_ARRAY_TASK_ID

mkdir -p "$OUT_DIR/seed_$SEED" logs

echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
echo "Phase A1b: Enhanced Training for 50%+ Target"
echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
echo ""
echo "Seed: $SEED"
echo "Dataset: 10,000 groups, 160,000 records"
echo "Model: hidden_dim=512 (optimal size for 10K data)"
echo "Training: 150 epochs with warmup + decay"
echo ""
echo "Target: 50%+ policy success"
echo ""

python scripts/train_dovla.py \
    --dataset "$DATASET" \
    --out "$OUT_DIR/seed_$SEED" \
    --objective lattice_field \
    --hidden-dim 512 \
    --action-horizon 4 \
    --epochs 150 \
    --batch-groups 16 \
    --records-per-group 8 \
    --lr 0.0003 \
    --weight-decay 0.01 \
    --device auto \
    --seed $SEED \
    --observation-mode state \
    --loss-weight bc=1.0 \
    --loss-weight field_effect=1.5 \
    --loss-weight field_potential=1.0 \
    --loss-weight field_preference=0.8 \
    --loss-weight field_anchor=0.2

echo ""
echo "✅ Phase A1b enhanced training complete (seed $SEED)"
echo ""
echo "Next: Evaluate and compare with 38.43% baseline"