#!/bin/bash
# Full training pipeline for the Trigger-Off experiment.
#
# Environment variables:
#   MODEL_TYPE   - "os_atlas" or "seeclick" (default: os_atlas)
#   MODEL_PATH   - HF hub ID or local path (default: from config)
#   TRAIN_JSON   - training data JSON (default: ./data/poisoned/train.json)
#   EVAL_JSON    - evaluation data JSON (default: ./data/poisoned/test.json)
#   OUTPUT_DIR   - checkpoint output dir (default: ./checkpoints/trigger_off)
#   CONFIG       - training config path (default: configs/training_config.yaml)
#   RESUME_FROM  - resume from checkpoint path (optional)

set -e
cd "$(dirname "$0")/.."

MODEL_TYPE=${MODEL_TYPE:-os_atlas}
TRAIN_JSON=${TRAIN_JSON:-./data/poisoned/train.json}
EVAL_JSON=${EVAL_JSON:-./data/poisoned/test.json}
OUTPUT_DIR=${OUTPUT_DIR:-./checkpoints/trigger_off}
CONFIG=${CONFIG:-configs/training_config.yaml}

echo "============================================================"
echo "  Trigger-Off: LoRA Training"
echo "============================================================"
echo "  MODEL_TYPE = $MODEL_TYPE"
echo "  TRAIN_JSON = $TRAIN_JSON"
echo "  EVAL_JSON  = $EVAL_JSON"
echo "  OUTPUT_DIR = $OUTPUT_DIR"
echo "  CONFIG     = $CONFIG"
echo "------------------------------------------------------------"

# Validate that the training data exists
if [ ! -f "$TRAIN_JSON" ]; then
    echo "ERROR: Training JSON not found at $TRAIN_JSON"
    echo "       Run scripts/prepare_data.sh first."
    exit 1
fi

RESUME_ARG=""
if [ -n "$RESUME_FROM" ]; then
    RESUME_ARG="--resume_from $RESUME_FROM"
fi

MODEL_PATH_ARG=""
if [ -n "$MODEL_PATH" ]; then
    MODEL_PATH_ARG="--model_path $MODEL_PATH"
fi

python -m src.training.train_lora \
    --config      "$CONFIG" \
    --model_type  "$MODEL_TYPE" \
    --train_json  "$TRAIN_JSON" \
    --eval_json   "$EVAL_JSON" \
    --output_dir  "$OUTPUT_DIR" \
    $MODEL_PATH_ARG \
    $RESUME_ARG

echo ""
echo "============================================================"
echo "  Training complete."
echo "  Checkpoint saved to: $OUTPUT_DIR"
echo "============================================================"