#!/usr/bin/env bash
# Step1 pretrain — 합성 데이터로 PGNet 학습
#
# 컨테이너 안 실행:
#   docker exec kr_lp_pgnet bash /workspace/kr_lp_pgnet/scripts/run_step1.sh
#
# 환경 변수:
#   DRY_RUN=1   2 epoch만 돌려 동작 검증
#   EPOCHS=N    epoch 수 override (기본 config의 epoch_num)
#   NUM_SAMPLES=N  합성 데이터 수 (기본 50000)

set -euo pipefail

PADDLEOCR_DIR=/workspace/PaddleOCR
KR_LP_DIR=/workspace/kr_lp_pgnet
TRAIN_DATA=/workspace/train_data
SYNTH_DIR="$TRAIN_DATA/kr_lp_synth"
ASSET_DIR="$KR_LP_DIR/data_gen/Korean-license-plate-Generator"
NUM_SAMPLES="${NUM_SAMPLES:-50000}"

TS=$(date +%Y%m%d_%H%M)
RUN_NAME="step1-${TS}"
OUTPUT_DIR="$PADDLEOCR_DIR/output/kr_lp_pgnet_${TS}"
LOG="$OUTPUT_DIR/run.log"

echo "==========================="
echo "RUN: $RUN_NAME"
echo "OUTPUT: $OUTPUT_DIR"
echo "==========================="

# ── 1. 합성 데이터 생성 ──────────────────────────────────────────────────────
echo "[1/3] 합성 데이터 생성 (${NUM_SAMPLES}장)"

rm -rf "$SYNTH_DIR"
python3.10 "$KR_LP_DIR/data_gen/generate_synthetic.py" \
    --asset_dir "$ASSET_DIR" \
    --out_dir   "$SYNTH_DIR" \
    --num       "$NUM_SAMPLES" \
    --dict      "$KR_LP_DIR/dict/kr_lp_dict.txt"

# ── 2. eval GT mat 생성 ─────────────────────────────────────────────────────
echo "[2/3] eval GT mat 생성"

python3.10 "$KR_LP_DIR/tools/make_gt_mat.py" \
    --label   "$SYNTH_DIR/test/test.txt" \
    --out_dir "$SYNTH_DIR/gt"

# ── 3. 학습 ─────────────────────────────────────────────────────────────────
echo "[3/3] Step1 학습 시작"
cd "$PADDLEOCR_DIR"

if [ ! -e ./train_data ]; then
    ln -sf "$TRAIN_DATA" ./train_data
fi

mkdir -p "$OUTPUT_DIR"

OVERRIDE=(
    -o Global.pretrained_model=./pretrain_models/train_step1/best_accuracy
       Global.load_static_weights=False
       Global.save_model_dir="${OUTPUT_DIR}/"
       Global.save_res_path="${OUTPUT_DIR}/predicts.txt"
       wandb.name="${RUN_NAME}"
)
if [ -n "${EPOCHS:-}" ]; then
    OVERRIDE+=(Global.epoch_num="$EPOCHS")
fi
if [ "${DRY_RUN:-0}" = "1" ]; then
    OVERRIDE+=(Global.epoch_num=2 Global.eval_batch_step="[0,200]")
    echo "DRY_RUN=1 → 2 epoch만 실행"
fi

echo "  config:   configs/e2e/kr_lp_pgnet.yml"
echo "  data:     $SYNTH_DIR/"
echo "  output:   $OUTPUT_DIR/"
echo "  wandb:    $RUN_NAME"
echo "  log:      $LOG"

python3.10 tools/train.py -c configs/e2e/kr_lp_pgnet.yml "${OVERRIDE[@]}" 2>&1 | tee "$LOG"