#!/usr/bin/env bash
# P1 generator-training queue: backbones x datasets, amortized (full-data), ~50k steps.
# Args: GPU [datasets] [backbones]   e.g.  p1_gen_queue.sh 5 "isic kvasir" "jit pixelgen deco pixeldit"
set -u
DR=/home/wzhang/LSC/Dataset/Segmentation/processed_unified
ROOT=/home/wzhang/LSC/Code/NPJ
PY=/opt/anaconda3/envs/seggen/bin/python
GPU=${1:-5}
ORDER=${2:-"isic kvasir"}
BKS=${3:-"jit pixelgen deco pixeldit"}
E="CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=$GPU TORCHDYNAMO_DISABLE=1 PYTHONPATH=. OMP_NUM_THREADS=4"
cd "$ROOT" || exit 1
ST=logs/p1gen_status.md
log(){ echo "[$(date '+%F %T')] $*" >> "$ST"; }
declare -A DSMAP=( [isic]="medsegdb_isic2018 holdout" [kvasir]="kvasir_seg official" )
log "=== P1 GEN QUEUE START GPU-$GPU order=[$ORDER] bks=[$BKS] pid=$$ ==="
for dk in $ORDER; do
  set -- ${DSMAP[$dk]}; ds=$1; proto=$2
  for bk in $BKS; do
    out=pretrained/pixdiff/p1_${bk}_${dk}.pt
    if [ -f "$out" ]; then log "SKIP $bk/$dk (ckpt exists)"; continue; fi
    log "TRAIN $bk/$dk -> $out"
    env $E $PY -m framework.synth.pixdiff.train --data_root "$DR" --dataset "$ds" --protocol "$proto" \
        --backbone "$bk" --img_size 256 --batch_size 16 --epochs 100000 --max_steps 50000 \
        --lr 1e-4 --amp bf16 --train_fraction 1.0 --fraction_seed 0 \
        --out_ckpt "$out" --log_interval 200 > logs/p1gen_${bk}_${dk}.log 2>&1
    rc=$?
    log "  done $bk/$dk rc=$rc ckpt=$([ -f "$out" ] && echo ok || echo MISSING)"
  done
done
log "=== P1 GEN QUEUE DONE GPU-$GPU ==="
echo "P1GEN_DONE_$GPU" >> "$ST"