#!/bin/bash
"""
Hyperparameter Optimization Launcher for Anti-UAV

Easy-to-use script for running automated hyperparameter optimization
with predefined configurations for different scenarios.

Usage:
    # Quick optimization (10 trials)
    ./run_hpo.sh --quick --data configs/antiuav.yaml
    
    # Full optimization (100 trials)
    ./run_hpo.sh --full --backend optuna --data configs/antiuav.yaml
    
    # Custom optimization
    ./run_hpo.sh --trials 50 --backend ray --algorithm bayesian
"""

set -e

# Default values
BACKEND="optuna"
ALGORITHM="bayesian" 
TRIALS=50
TIMEOUT=3600
JOBS=1
DATA_CONFIG=""
MODEL="yolov8s.pt"
TASK="detect"
NAME="hpo_experiment"
SAVE_DIR="./hpo_results"

# Preset configurations
QUICK_TRIALS=10
FULL_TRIALS=100

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --quick)
            TRIALS=$QUICK_TRIALS
            TIMEOUT=1800  # 30 minutes
            shift
            ;;
        --full)
            TRIALS=$FULL_TRIALS
            TIMEOUT=7200  # 2 hours
            shift
            ;;
        --backend)
            BACKEND="$2"
            shift 2
            ;;
        --algorithm)
            ALGORITHM="$2"
            shift 2
            ;;
        --trials)
            TRIALS="$2"
            shift 2
            ;;
        --timeout)
            TIMEOUT="$2"
            shift 2
            ;;
        --jobs)
            JOBS="$2"
            shift 2
            ;;
        --data)
            DATA_CONFIG="$2"
            shift 2
            ;;
        --model)
            MODEL="$2"
            shift 2
            ;;
        --task)
            TASK="$2"
            shift 2
            ;;
        --name)
            NAME="$2"
            shift 2
            ;;
        --save-dir)
            SAVE_DIR="$2"
            shift 2
            ;;
        -h|--help)
            echo "Usage: $0 [OPTIONS]"
            echo ""
            echo "Preset Options:"
            echo "  --quick             Quick optimization (10 trials, 30min)"
            echo "  --full              Full optimization (100 trials, 2h)"
            echo ""
            echo "Configuration Options:"
            echo "  --backend STR       Optimization backend (optuna|ray)"
            echo "  --algorithm STR     Algorithm (bayesian|evolutionary|grid|random)"
            echo "  --trials INT        Number of trials"
            echo "  --timeout INT       Timeout in seconds"
            echo "  --jobs INT          Parallel jobs"
            echo "  --data STR          Dataset config file (required)"
            echo "  --model STR         Model architecture (yolov8s.pt)"
            echo "  --task STR          Task type (detect|track|edtc)"
            echo "  --name STR          Experiment name"
            echo "  --save-dir STR      Results directory"
            echo "  -h, --help          Show this help"
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

# Validation
if [[ -z "$DATA_CONFIG" ]]; then
    echo "Error: --data is required"
    echo "Use --help for usage information"
    exit 1
fi

if [[ ! -f "$DATA_CONFIG" ]]; then
    echo "Error: Data config file '$DATA_CONFIG' not found"
    exit 1
fi

# Check dependencies
python -c "import optuna" 2>/dev/null || {
    if [[ "$BACKEND" == "optuna" ]]; then
        echo "Error: Optuna not installed. Install with: pip install optuna"
        exit 1
    fi
}

python -c "import ray" 2>/dev/null || {
    if [[ "$BACKEND" == "ray" ]]; then
        echo "Error: Ray Tune not installed. Install with: pip install ray[tune]"
        exit 1
    fi
}

# Create output directory
mkdir -p "$SAVE_DIR"

echo "=============================================="
echo "Anti-UAV Hyperparameter Optimization"
echo "=============================================="
echo "Backend: $BACKEND"
echo "Algorithm: $ALGORITHM"
echo "Trials: $TRIALS"
echo "Timeout: ${TIMEOUT}s ($(($TIMEOUT / 60))min)"
echo "Parallel jobs: $JOBS"
echo "Data config: $DATA_CONFIG"
echo "Model: $MODEL"
echo "Task: $TASK"
echo "Experiment: $NAME"
echo "Results: $SAVE_DIR"
echo "=============================================="

# Set environment variables for optimal performance
export OMP_NUM_THREADS=1
export CUDA_VISIBLE_DEVICES=0

# Run hyperparameter optimization
echo "Starting hyperparameter optimization..."

python tools/hyperparameter_optimization.py \
    --backend "$BACKEND" \
    --algorithm "$ALGORITHM" \
    --n-trials $TRIALS \
    --timeout $TIMEOUT \
    --n-jobs $JOBS \
    --data "$DATA_CONFIG" \
    --model "$MODEL" \
    --task "$TASK" \
    --name "$NAME" \
    --save-dir "$SAVE_DIR"

# Check if optimization completed successfully
if [[ $? -eq 0 ]]; then
    echo ""
    echo "=============================================="
    echo "Hyperparameter optimization completed!"
    echo "=============================================="
    echo "Results saved to: $SAVE_DIR/$NAME"
    echo ""
    echo "Generated files:"
    echo "  - ${BACKEND}_results.json  : Raw optimization results"
    echo "  - analysis.json           : Comprehensive analysis"
    echo "  - hpo_analysis.png        : Visualization plots"
    echo ""
    echo "To use the best configuration:"
    echo "  1. Check analysis.json for best_configuration"
    echo "  2. Apply parameters to your training script"
    echo "  3. Run full training with optimized hyperparameters"
else
    echo "Error: Hyperparameter optimization failed!"
    exit 1
fi