#!/bin/bash
#SBATCH --job-name=mavt-s3-threed
#SBATCH --partition=defq
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=2
#SBATCH --gpus-per-node=2
#SBATCH --cpus-per-task=24
#SBATCH --mem=120G
#SBATCH --time=48:00:00
#SBATCH --output=logs/stage3_dgx_%j.log
#SBATCH --error=logs/stage3_dgx_%j.err

# ============================================================================
# MAVT Stage 3 — Image + Video + 3D triplane
#   - Init from stage1_5 step=120k (best converged Stage 1 ckpt;
#     Stage 2 has no ckpt >15k due to save_top_k=1 bug in old config).
#   - 2x A100 DDP, bf16-mixed
#   - 50k steps, batch=6/GPU × 2 GPU × accum=2 = effective batch 24
#   - LR 2e-5, warmup 500
#   - SigLIP2 fully unfrozen (stage 3)
#   - 3D from dataset/universal_3d/ (30 519 triplane objects, captions/3d.json)
#
# Prerequisites (already in place):
#   - dataset/universal_3d/3d_objects/renders → symlink to dataset/tripplane
#   - dataset/universal_3d/captions/3d.json (30 519 LVIS captions)
#   - dataset/image10k/train (WDS image shards)
#   - dataset/dataset_10m (video shards; 51% known corrupt, loader handles)
#
# Submit:
#   sbatch train_stage3_dgx.sh
# ============================================================================

set -euo pipefail

PROJECT_DIR="/home/user02/linhdang/Antokenizer"
cd "$PROJECT_DIR"

# --- Data paths -----------------------------------------------------------
IMAGE_SHARDS_DIR="$PROJECT_DIR/dataset/image10k/train"
VIDEO_SHARDS_DIR="$PROJECT_DIR/dataset/dataset_10m"
UNIVERSAL_ROOT="$PROJECT_DIR/dataset/universal_3d"

# --- Init checkpoint ------------------------------------------------------
INIT_CKPT="checkpoints/stage1_5/balanced/mavt-stage1_5-balanced-step=0120000-val/loss=0.1238.ckpt"

mkdir -p logs checkpoints/stage3_dgx

# --- Environment ----------------------------------------------------------
if [ -d "$PROJECT_DIR/.venv" ]; then
    source "$PROJECT_DIR/.venv/bin/activate"
elif [ -d "$HOME/miniconda3" ]; then
    source "$HOME/miniconda3/etc/profile.d/conda.sh"
    conda activate base
fi

export PYTHONPATH="$PROJECT_DIR/src:${PYTHONPATH:-}"
export TORCH_NCCL_BLOCKING_WAIT=1
export OMP_NUM_THREADS=8
export TOKENIZERS_PARALLELISM=false

# --- 3D data sanity -------------------------------------------------------
THREED_RENDERS="$UNIVERSAL_ROOT/3d_objects/renders"
N_3D=$(find -L "$THREED_RENDERS" -mindepth 1 -maxdepth 1 -type d 2>/dev/null | wc -l)
echo "========================================"
echo "  MAVT Stage 3 — Image + Video + 3D"
echo "  GPUs requested:  2"
echo "  Image shards:    $IMAGE_SHARDS_DIR"
echo "  Video shards:    $VIDEO_SHARDS_DIR"
echo "  Universal root:  $UNIVERSAL_ROOT"
echo "  3D objects:      $N_3D renders"
echo "  Init ckpt:       $INIT_CKPT"
echo "========================================"

if [ "$N_3D" -lt 1000 ]; then
    echo "[FATAL] Too few 3D objects ($N_3D). Expected ≥ 1000."
    exit 1
fi

if [ ! -f "$INIT_CKPT" ]; then
    echo "[FATAL] Init checkpoint not found: $INIT_CKPT"
    exit 1
fi

# Use init_from_ckpt (soft restart, keeps model weights but resets optimizer).
# For Stage 3 init from Stage 1, this is the intended behavior — new threed
# poolers are created via prepare_poolers() during setup('fit').
CKPT_ARG="--model.init_from_ckpt $INIT_CKPT"

# --- Launch ---------------------------------------------------------------
srun --mpi=none python train.py fit \
    --config configs/model/mavt_base.yaml \
    --config configs/train/universal_data/stage3_universal.yaml \
    --config configs/train/universal_data/stage3_paths.yaml \
    --data.image_shards_dir "$IMAGE_SHARDS_DIR" \
    --data.video_shards_dir "$VIDEO_SHARDS_DIR" \
    --data.video_max_shards 100 \
    --data.universal_data_root "$UNIVERSAL_ROOT" \
    --data.active_modalities '["image", "video", "threed"]' \
    --model.active_modalities '["image", "video", "threed"]' \
    --data.image_resolution 256 \
    --data.video_frames 16 \
    --data.video_resolution 256 \
    --data.triplane_res 256 \
    --data.batch_size 6 \
    --data.num_workers 6 \
    --data.pin_memory true \
    --data.persistent_workers true \
    --data.prefetch_factor 3 \
    --model.training_stage 3 \
    --model.init_siglip2 true \
    --model.use_lpips true \
    --model.use_clip false \
    --model.w_l1 1.0 \
    --model.w_lpips 0.2 \
    --model.w_sem 0.3 \
    --model.w_temp 0.05 \
    --model.warmup_steps 500 \
    --model.total_steps 50000 \
    --model.weight_decay 0.01 \
    --model.grad_clip 1.0 \
    $CKPT_ARG \
    --trainer.devices 2 \
    --trainer.strategy ddp_find_unused_parameters_true \
    --trainer.precision bf16-mixed \
    --trainer.max_steps 50000 \
    --trainer.accumulate_grad_batches 2 \
    --trainer.log_every_n_steps 50 \
    --trainer.val_check_interval 1000 \
    --trainer.logger.class_path lightning.pytorch.loggers.WandbLogger \
    --trainer.logger.init_args.project mavt \
    --trainer.logger.init_args.name stage3_threed