#!/usr/bin/env bash
# run_all.sh  –  Chạy toàn bộ pipeline: chuẩn bị dữ liệu → PhoBERT → GPT-OSS-20B → So sánh
# ========================================================================================
# Dùng:
#   bash run_all.sh                  # full pipeline
#   bash run_all.sh --skip_gpt       # chỉ PhoBERT
#
# Biến môi trường có thể ghi đè:
#   CSV_PATH     : đường dẫn tới file CSV gốc (mặc định: /tmp/dataset_nlp/clickbait_dataset_vietnamese.csv)
#   DATA_DIR     : thư mục chứa splits        (mặc định: data/splits)
#   PB_OUT       : output PhoBERT             (mặc định: outputs/phobert)
#   GPT_OUT      : output GPT-OSS-20B         (mặc định: outputs/gpt20b)
#   GPT_MODEL    : HuggingFace model ID       (mặc định: openai/gpt-oss-20b)
# ========================================================================================

set -euo pipefail

# ── Cài đặt mặc định ─────────────────────────────────────────────────────────
CSV_PATH="${CSV_PATH:-/tmp/dataset_nlp/clickbait_dataset_vietnamese.csv}"
DATA_DIR="${DATA_DIR:-data/splits}"
PB_OUT="${PB_OUT:-outputs/phobert}"
GPT_OUT="${GPT_OUT:-outputs/gpt20b}"
GPT_MODEL="${GPT_MODEL:-openai/gpt-oss-20b}"
SKIP_GPT=false

# ── Parse args ────────────────────────────────────────────────────────────────
while [[ $# -gt 0 ]]; do
  case "$1" in
    --skip_gpt)    SKIP_GPT=true;    shift ;;
    --model)       GPT_MODEL="$2";   shift 2 ;;
    --csv)         CSV_PATH="$2";    shift 2 ;;
    *) echo "Unknown arg: $1"; exit 1 ;;
  esac
done

echo "============================================================"
echo "  ViClickbait-2025  –  Phân lớp Clickbait Tiếng Việt"
echo "============================================================"
echo "CSV      : $CSV_PATH"
echo "DataDir  : $DATA_DIR"
echo "PB_Out   : $PB_OUT"
echo "GPT_Model: $GPT_MODEL"
echo "GPT_Out  : $GPT_OUT"
echo "Skip GPT : $SKIP_GPT"
echo "------------------------------------------------------------"

# ── Bước 0: Cài dependencies nếu cần ─────────────────────────────────────────
echo ""
echo "[0/4] Kiểm tra / cài dependencies..."
pip install -q transformers datasets evaluate accelerate scikit-learn pandas peft bitsandbytes

PYTHON=${PYTHON:-python3}

# ── Bước 1: Chuẩn bị dữ liệu ─────────────────────────────────────────────────
echo ""
echo "[1/4] Chuẩn bị dữ liệu (split 80/10/10 stratified)..."
$PYTHON scripts/prepare_data.py \
    --csv      "$CSV_PATH" \
    --out_dir  "$DATA_DIR" \
    --text_mode title_lead \
    --seed     42

# ── Bước 2: PhoBERT Baseline ──────────────────────────────────────────────────
echo ""
echo "[2/4] Huấn luyện PhoBERT baseline..."
$PYTHON scripts/train_phobert.py \
    --data_dir   "$DATA_DIR" \
    --output_dir "$PB_OUT" \
    --model_name vinai/phobert-base-v2 \
    --max_length 256 \
    --batch_size 32 \
    --lr         2e-5 \
    --epochs     10 \
    --warmup_ratio 0.1 \
    --weight_decay 0.01 \
    --patience   5 \
    --seed       42

# ── Bước 3: GPT OSS 20B (LoRA SFT) ─────────────────────────────────────────
if [ "$SKIP_GPT" = false ]; then
    echo ""
    echo "[3/4] Huấn luyện GPT OSS 20B với LoRA SFT..."
    $PYTHON scripts/train_gpt_oss_20b.py \
        --data_dir   "$DATA_DIR" \
        --output_dir "$GPT_OUT" \
        --model_name "$GPT_MODEL" \
        --max_length 256 \
        --batch_size 4 \
        --grad_accum 8 \
        --lr         2e-5 \
        --epochs     10 \
        --warmup_ratio 0.1 \
        --weight_decay 0.01 \
        --patience   3 \
        --lora_r     32 \
        --lora_alpha 64 \
        --lora_dropout 0.05 \
        --seed       42
else
    echo ""
    echo "[3/4] Bỏ qua GPT OSS 20B (--skip_gpt)"
fi

# ── Bước 4: So sánh kết quả ───────────────────────────────────────────────────
echo ""
echo "[4/4] So sánh kết quả..."
$PYTHON scripts/compare_results.py \
    --phobert_dir "$PB_OUT" \
    --gpt_dir     "$GPT_OUT"

echo ""
echo "============================================================"
echo "  Pipeline hoàn tất!"
echo "  Kết quả PhoBERT : $PB_OUT/test_results.json"
if [ "$SKIP_GPT" = false ]; then
echo "  Kết quả GPT-20B : $GPT_OUT/test_results.json"
fi
echo "============================================================"