#!/bin/bash
# =============================================================
#  Spam Classifier — Liquid AI LoRA Retrain
#  Double-click to choose fast or full retrain mode.
#  Replaces retrain-fast.command and retrain-full.command.
# =============================================================

cd "$(dirname "$0")"
source venv/bin/activate

echo "============================================================"
echo "  Liquid AI LoRA Retrain — Spam / Ham / Phishing"
echo "  Model: LiquidAI/LFM2.5-1.2B-Instruct"
echo "============================================================"
echo ""
echo "  f) Fast retrain  — ~8,000 examples, ~1-1.5 hours"
echo "  u) Full retrain  — ~16,000 examples, ~2 hours"
echo "  q) Quit"
echo ""
echo "  Memory optimizations: MPS cache flush, gradient checkpointing, bf16, fused AdamW."
echo ""
read -p "Choice [f/u/q]: " MODE_CHOICE

case "$MODE_CHOICE" in
    q|Q)
        echo "Bye!"
        sleep 2
        exit 0
        ;;
    f|F)
        MODE="fast"
        ADAPTER_DIR="$(dirname "$0")/adapters_fast"
        echo ""
        echo "  Starting FAST retrain (~8K examples, ~1-1.5 hours)..."
        ;;
    u|U)
        MODE="full"
        ADAPTER_DIR="$(dirname "$0")/adapters_full"
        echo ""
        echo "  Starting FULL retrain (~16K examples, ~2 hours)..."
        ;;
    *)
        echo "Invalid choice."
        sleep 3
        exit 1
        ;;
esac

echo ""
python3 "../new_training_data/retrain_liquid.py" --mode "$MODE"
RETRAIN_STATUS=$?

if [[ $RETRAIN_STATUS -ne 0 ]]; then
    echo ""
    echo "Retrain failed (exit $RETRAIN_STATUS)."
    echo ""
    read -p "Press any key to close..."
    exit 1
fi

echo ""
echo "============================================================"
echo "  Retrain complete!"
echo "  New adapter: $ADAPTER_DIR"
echo "============================================================"
echo ""
echo "Would you like to make this the default adapter?"
echo "  - Backs up current adapters/ -> adapters_backup/"
echo "  - Copies adapters_${MODE}/ -> adapters/"
echo "  - The app and notebook will use the new adapter automatically"
echo ""
read -p "Swap in as default? [y/N]: " SWAP

if [[ "$SWAP" == "y" || "$SWAP" == "Y" ]]; then
    PROJ_DIR="$(dirname "$0")"
    if [[ -d "$PROJ_DIR/adapters" ]] && [[ ! -d "$PROJ_DIR/adapters_backup" ]]; then
        mv "$PROJ_DIR/adapters" "$PROJ_DIR/adapters_backup"
        echo "  Backed up adapters/ -> adapters_backup/"
    elif [[ -d "$PROJ_DIR/adapters" ]]; then
        rm -rf "$PROJ_DIR/adapters_old_backup"
        mv "$PROJ_DIR/adapters_backup" "$PROJ_DIR/adapters_old_backup" 2>/dev/null
        mv "$PROJ_DIR/adapters" "$PROJ_DIR/adapters_backup"
        echo "  Backed up adapters/ -> adapters_backup/"
    fi
    cp -r "$ADAPTER_DIR" "$PROJ_DIR/adapters"
    echo "  Copied adapters_${MODE}/ -> adapters/"
    echo "  The app and notebook now use the new adapter!"
else
    echo "  Skipped. To use later, copy adapters_${MODE}/ to adapters/"
fi

echo ""
echo "Next: run Retrain.command in the LLM Project folder to"
echo "rebuild the GGUF and upload to HuggingFace."
echo ""
read -p "Press any key to close..."