#!/bin/bash
# TempSamp-R1 (no-CoT) baseline on ActivityForensics.
# Algorithm comes from forensics_grpo's existing trainer; this script just
# locks in the TempSamp-R1 configuration:
#   - no-CoT prompt           (FORENSICS_COT=false)
#   - mix-policy + non-linear shaping  (default trainer behavior, no flag needed)
#   - one IoU reward only      (--reward_funcs hungarian_iou)
#   - no forensics-specific extras (FORENSICS_*_ADV, FORENSICS_CTA_*, ...)

set -e

# --- TempSamp-R1 algorithmic configuration (do not edit lightly) ----------
export FORENSICS_COT=false
unset FORENSICS_DECOMP_ADV FORENSICS_HUNGARIAN_DECOMP
unset FORENSICS_CTA_GRPO FORENSICS_FORGERY_AWARE
unset FORENSICS_USE_EXTERNAL_VERIFIER FORENSICS_VERIFIER_CONTEXT
# --------------------------------------------------------------------------

export FORENSICS_GRPO_ROOT=${FORENSICS_GRPO_ROOT:-/mnt/local-fast/zhangt/forensics_grpo}
export WANDB_PROJECT=Forensics-Baselines
export WANDB_NAME=TempSampR1_nocot_forensics_7B
export PYTHONPATH="${FORENSICS_GRPO_ROOT}:.:$PYTHONPATH"

OUTDIR=./logs/${WANDB_NAME}
mkdir -p "${OUTDIR}"
export DEBUG_MODE="true"
export LOG_PATH="${OUTDIR}/${WANDB_NAME}.txt"

MODEL_PATH=${MODEL_PATH:-Qwen/Qwen2.5-VL-7B-Instruct}
ANNOT_DIR=${ANNOT_DIR:-/ces/zt/activityforensics/annot}
VIDEO_ROOT=${VIDEO_ROOT:-/ces/zt}
PREPROC_DIR=${PREPROC_DIR:-/mnt/local-fast/zhangt/forensics_grpo/preprocessed}

torchrun --nproc_per_node="${NPROC:-4}" \
    --nnodes="${NNODES:-1}" \
    --node_rank="${NODE_RANK:-0}" \
    --master_addr="${MASTER_ADDR:-127.0.0.1}" \
    --master_port="${MASTER_PORT:-12371}" \
    src/tempsamp_main.py \
    --deepspeed "${FORENSICS_GRPO_ROOT}/scripts/zero3_offload.json" \
    --output_dir "${OUTDIR}" \
    --model_name_or_path "${MODEL_PATH}" \
    --annot_dir "${ANNOT_DIR}" \
    --video_root "${VIDEO_ROOT}" \
    --preprocessed_data_path "${PREPROC_DIR}" \
    --dataset_name forensics \
    --max_prompt_length 8192 \
    --max_completion_length 256 \
    --num_generations 8 \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 2 \
    --logging_steps 1 \
    --bf16 \
    --torch_dtype bfloat16 \
    --data_seed 42 \
    --gradient_checkpointing true \
    --attn_implementation flash_attention_2 \
    --num_train_epochs 2 \
    --run_name "${WANDB_NAME}" \
    --report_to tensorboard \
    --reward_funcs hungarian_iou \
    --temperature 1.0 \
    --beta 0.0 \
    --logging_dir "${OUTDIR}" \
    --save_steps 100 \
    --save_only_model true