#!/usr/bin/env bash
#
# InstanceV training script for InstanceCap-BBox
#
# Usage:
#   bash run_instancev_training_instancecap_bbox.sh
#   bash run_instancev_training_instancecap_bbox.sh --background
#   bash run_instancev_training_instancecap_bbox.sh --resume
#   bash run_instancev_training_instancecap_bbox.sh --resume-from /path/to/ckpt.safetensors
#

set -e

# ==================== Argument parsing ====================

RUN_BACKGROUND=false
RUN_INTERNAL=false
RESUME_MODE="auto"   # auto | off | path
RESUME_FROM=""

while [[ $# -gt 0 ]]; do
    case "$1" in
        --background|-bg)
            RUN_BACKGROUND=true
            shift
            ;;
        --run-internal)
            RUN_INTERNAL=true
            shift
            ;;
        --resume)
            RESUME_MODE="auto"
            shift
            ;;
        --no-resume)
            RESUME_MODE="off"
            shift
            ;;
        --resume-from)
            if [[ -z "${2:-}" ]]; then
                echo "Error: --resume-from requires a checkpoint path"
                exit 1
            fi
            RESUME_MODE="path"
            RESUME_FROM="$2"
            shift 2
            ;;
        --resume-from=*)
            RESUME_MODE="path"
            RESUME_FROM="${1#*=}"
            shift
            ;;
        *)
            echo "Unknown arg: $1"
            exit 1
            ;;
    esac
done

# ==================== Config ====================

PROJECT_ROOT="/data/rczhang/PencilFolder/DiffSynth-Studio"
cd "$PROJECT_ROOT"

DATA_BASE_PATH="/data/rczhang/PencilFolder/data"
INSTANCECAP_PATH="${DATA_BASE_PATH}/InstanceCap/InstanceCap.jsonl"
INSTANCECAP_BBOX_DIR="${DATA_BASE_PATH}/InstanceCap-BBox"
VIDEO_DIR="${DATA_BASE_PATH}/OpenVid1M-Video-InstanceCap"
MASK_ROOT_DIR="${DATA_BASE_PATH}/InstanceCap-BBox-Masks"
METADATA_PATH="${DATA_BASE_PATH}/InstanceCap/instancev_instancecap_bbox.jsonl"
FORCE_REBUILD_METADATA=true

MIN_INSTANCES=1
MAX_INSTANCES=5
MIN_FRAMES=81
NUM_WORKERS=${NUM_WORKERS:-8}

MODEL_ID_WITH_ORIGIN_PATHS="Wan-AI/Wan2.1-T2V-14B:diffusion_pytorch_model*.safetensors,Wan-AI/Wan2.1-T2V-14B:models_t5_umt5-xxl-enc-bf16.pth,Wan-AI/Wan2.1-T2V-14B:Wan2.1_VAE.pth"

TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
OUTPUT_PATH="${PROJECT_ROOT}/models/train/instancev_instancecap_bbox_${TIMESTAMP}"
LOG_DIR="${OUTPUT_PATH}/logs"
mkdir -p "$LOG_DIR"
LOG_FILE="${LOG_DIR}/train_${TIMESTAMP}.log"

# Video settings: align with Wan default (4n+1). Short clips are truncated to nearest 4n+1.
NUM_FRAMES=81
HEIGHT=480
WIDTH=832
DATASET_REPEAT=1

GRADIENT_ACCUMULATION_STEPS=4
LEARNING_RATE=1e-4
NUM_EPOCHS=5
SAVE_STEPS=500

SAUG_DROP_PROB=0.1
SAUG_SCALE=0.0

USE_GRADIENT_CHECKPOINTING=true
MIXED_PRECISION="bf16"

WANDB_PROJECT="instancev-instancecap-bbox"
WANDB_RUN_NAME="instancev_instancecap_bbox_${TIMESTAMP}"

TRAINABLE_MODELS="dit"
TASK="sft"

export CUDA_VISIBLE_DEVICES=0,1
export CUDA_LAUNCH_BLOCKING=0

ACCELERATE_BIN="/home/rczhang/miniconda3/envs/diffsyn/bin/accelerate"
if [[ ! -x "${ACCELERATE_BIN}" ]]; then
    ACCELERATE_BIN="accelerate"
fi

PYTHON_BIN="/home/rczhang/miniconda3/envs/diffsyn/bin/python"
if [[ ! -x "${PYTHON_BIN}" ]]; then
    PYTHON_BIN="python"
fi

# ==================== Functions ====================

print_config() {
    echo "=============================================="
    echo "    InstanceV Training (InstanceCap-BBox)     "
    echo "=============================================="
    echo ""
    echo "[Data]"
    echo "  - InstanceCap: ${INSTANCECAP_PATH}"
    echo "  - BBox Dir: ${INSTANCECAP_BBOX_DIR}"
    echo "  - Video Dir: ${VIDEO_DIR}"
    echo "  - Mask Root: ${MASK_ROOT_DIR}"
    echo "  - Metadata: ${METADATA_PATH}"
    echo "  - Rebuild Metadata: ${FORCE_REBUILD_METADATA}"
    echo "  - Preprocess Workers: ${NUM_WORKERS}"
    echo ""
    echo "[Model]"
    echo "  - Base: Wan2.1-T2V-1.3B"
    echo "  - Output: ${OUTPUT_PATH}"
    echo ""
    echo "[Train]"
    echo "  - Resolution: ${WIDTH}x${HEIGHT}"
    echo "  - Frames: ${NUM_FRAMES}"
    echo "  - LR: ${LEARNING_RATE}"
    echo "  - Epochs: ${NUM_EPOCHS}"
    echo "  - Grad Accum: ${GRADIENT_ACCUMULATION_STEPS}"
    echo "  - Mixed Precision: ${MIXED_PRECISION}"
    echo "  - Save Steps: ${SAVE_STEPS}"
    echo ""
    echo "[InstanceV]"
    echo "  - SAUG Dropout: ${SAUG_DROP_PROB}"
    echo "  - SAUG Scale: ${SAUG_SCALE}"
    echo "  - Min Frames Filter: ${MIN_FRAMES}"
    echo ""
    echo "[Wandb]"
    echo "  - Project: ${WANDB_PROJECT}"
    echo "  - Run Name: ${WANDB_RUN_NAME}"
    echo ""
    echo "[Resume]"
    echo "  - Mode: ${RESUME_MODE}"
    if [[ "${RESUME_MODE}" == "path" ]]; then
        echo "  - Checkpoint: ${RESUME_FROM}"
    fi
    echo ""
    echo "[GPU]"
    echo "  - CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES}"
    echo ""
    echo "=============================================="
}

run_training() {
    echo ""
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] Step 1: Prepare metadata..."

    if [[ "${FORCE_REBUILD_METADATA}" == "true" ]] || [ ! -f "$METADATA_PATH" ]; then
        "${PYTHON_BIN}" examples/wanvideo/model_training/prepare_instancev_instancecap_bbox.py \
            --instancecap_path "${INSTANCECAP_PATH}" \
            --instancecap_bbox_dir "${INSTANCECAP_BBOX_DIR}" \
            --video_dir "${VIDEO_DIR}" \
            --mask_root_dir "${MASK_ROOT_DIR}" \
            --output_path "${METADATA_PATH}" \
            --dataset_base_path "${DATA_BASE_PATH}" \
            --min_instances "${MIN_INSTANCES}" \
            --max_instances "${MAX_INSTANCES}" \
            --min_frames "${MIN_FRAMES}" \
            --num_workers "${NUM_WORKERS}"
    else
        SAMPLE_COUNT=$(wc -l < "$METADATA_PATH")
        echo "Found metadata: ${METADATA_PATH}"
        echo "Samples: ${SAMPLE_COUNT}"
    fi

    echo ""
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] Step 2: Start training..."
    echo ""

    GC_FLAG=""
    if [[ "${USE_GRADIENT_CHECKPOINTING}" == "true" ]]; then
        GC_FLAG="--use_gradient_checkpointing"
    fi

    RESUME_ARGS=()
    RESUME_PATH=""
    if [[ "${RESUME_MODE}" != "off" ]]; then
        if [[ "${RESUME_MODE}" == "path" ]]; then
            if [[ ! -f "${RESUME_FROM}" ]]; then
                echo "Checkpoint not found: ${RESUME_FROM}"
                exit 1
            fi
            RESUME_PATH="${RESUME_FROM}"
        else
            RESUME_PATH=$(ls -1t "${OUTPUT_PATH}"/step-*.safetensors "${OUTPUT_PATH}"/epoch-*.safetensors 2>/dev/null | head -n 1 || true)
        fi
        if [[ -n "${RESUME_PATH}" ]]; then
            echo "Resume from: ${RESUME_PATH}"
            RESUME_ARGS+=(--resume_from_checkpoint "${RESUME_PATH}")
        else
            echo "No checkpoint found, starting fresh"
        fi
    fi

    "${ACCELERATE_BIN}" launch \
        --num_processes 2 \
        --num_machines 1 \
        --mixed_precision="${MIXED_PRECISION}" \
        examples/wanvideo/model_training/train_instancev.py \
        --dataset_base_path "${DATA_BASE_PATH}" \
        --dataset_metadata_path "${METADATA_PATH}" \
        --data_file_keys "video" \
        --height ${HEIGHT} \
        --width ${WIDTH} \
        --num_frames ${NUM_FRAMES} \
        --dataset_repeat ${DATASET_REPEAT} \
        --model_id_with_origin_paths "${MODEL_ID_WITH_ORIGIN_PATHS}" \
        --learning_rate ${LEARNING_RATE} \
        --num_epochs ${NUM_EPOCHS} \
        --gradient_accumulation_steps ${GRADIENT_ACCUMULATION_STEPS} \
        --save_steps ${SAVE_STEPS} \
        --output_path "${OUTPUT_PATH}" \
        --remove_prefix_in_ckpt "pipe.dit." \
        --trainable_models "${TRAINABLE_MODELS}" \
        --task "${TASK}" \
        --saug_drop_prob ${SAUG_DROP_PROB} \
        --saug_scale ${SAUG_SCALE} \
        ${GC_FLAG} \
        --use_wandb \
        --wandb_project "${WANDB_PROJECT}" \
        --wandb_run_name "${WANDB_RUN_NAME}" \
        "${RESUME_ARGS[@]}"

    echo ""
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] Training done"
    echo "Checkpoints: ${OUTPUT_PATH}"
}

# ==================== Main ====================

if [[ "${RUN_INTERNAL}" == "true" ]]; then
    print_config
    run_training
    exit 0
fi

print_config

if [[ "${RUN_BACKGROUND}" == "true" ]]; then
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] Background mode, log: ${LOG_FILE}"
    INTERNAL_ARGS=(--run-internal)
    if [[ "${RESUME_MODE}" == "off" ]]; then
        INTERNAL_ARGS+=(--no-resume)
    elif [[ "${RESUME_MODE}" == "path" ]]; then
        INTERNAL_ARGS+=(--resume-from "${RESUME_FROM}")
    fi
    nohup bash "$0" "${INTERNAL_ARGS[@]}" > "${LOG_FILE}" 2>&1 &
    PID=$!
    echo "PID: ${PID}"
    echo "Log: tail -f ${LOG_FILE}"
    echo "Stop: kill ${PID}"
    echo "${PID}" > "${LOG_DIR}/train_${TIMESTAMP}.pid"
else
    run_training
fi