#!/bin/bash
"""
Distributed Training Launch Script for Anti-UAV

This script provides easy launching of distributed training across multiple GPUs
with automatic configuration and optimal performance settings.

Usage:
    # Single node multi-GPU
    ./launch_distributed.sh --gpus 4 --data configs/antiuav.yaml
    
    # Multi-node training  
    ./launch_distributed.sh --nodes 2 --gpus-per-node 4 --node-rank 0 --master-addr 192.168.1.100
"""

set -e  # Exit on any error

# Default values
GPUS_PER_NODE=4
NODES=1
NODE_RANK=0
MASTER_ADDR="localhost"
MASTER_PORT=29500
DATA_CONFIG=""
MODEL="yolov8s.pt"
EPOCHS=100
BATCH_SIZE=16
WORKERS=8
PROJECT="./runs/distributed"
NAME="antiuav_distributed"

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --gpus|--gpus-per-node)
            GPUS_PER_NODE="$2"
            shift 2
            ;;
        --nodes)
            NODES="$2"
            shift 2
            ;;
        --node-rank)
            NODE_RANK="$2"
            shift 2
            ;;
        --master-addr)
            MASTER_ADDR="$2"
            shift 2
            ;;
        --master-port)
            MASTER_PORT="$2"
            shift 2
            ;;
        --data)
            DATA_CONFIG="$2"
            shift 2
            ;;
        --model)
            MODEL="$2"
            shift 2
            ;;
        --epochs)
            EPOCHS="$2"
            shift 2
            ;;
        --batch-size)
            BATCH_SIZE="$2"
            shift 2
            ;;
        --workers)
            WORKERS="$2"
            shift 2
            ;;
        --project)
            PROJECT="$2"
            shift 2
            ;;
        --name)
            NAME="$2"
            shift 2
            ;;
        -h|--help)
            echo "Usage: $0 [OPTIONS]"
            echo ""
            echo "Options:"
            echo "  --gpus-per-node INT     Number of GPUs per node (default: 4)"
            echo "  --nodes INT             Number of nodes (default: 1)"
            echo "  --node-rank INT         Node rank for multi-node training (default: 0)"
            echo "  --master-addr STR       Master node address (default: localhost)"
            echo "  --master-port INT       Master port (default: 29500)"
            echo "  --data STR              Dataset config file (required)"
            echo "  --model STR             Model architecture (default: yolov8s.pt)"
            echo "  --epochs INT            Number of epochs (default: 100)"
            echo "  --batch-size INT        Batch size per GPU (default: 16)"
            echo "  --workers INT           Data loading workers per GPU (default: 8)"
            echo "  --project STR           Project directory (default: ./runs/distributed)"
            echo "  --name STR              Experiment name (default: antiuav_distributed)"
            echo "  -h, --help              Show this help message"
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

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

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

# Calculate total world size
WORLD_SIZE=$((NODES * GPUS_PER_NODE))

# Set environment variables
export MASTER_ADDR=$MASTER_ADDR
export MASTER_PORT=$MASTER_PORT
export WORLD_SIZE=$WORLD_SIZE
export NODE_RANK=$NODE_RANK

# GPU availability check
AVAILABLE_GPUS=$(nvidia-smi -L | wc -l)
if [[ $GPUS_PER_NODE -gt $AVAILABLE_GPUS ]]; then
    echo "Error: Requested $GPUS_PER_NODE GPUs but only $AVAILABLE_GPUS available"
    exit 1
fi

# Create output directory
mkdir -p "$PROJECT"

echo "==============================================="
echo "Anti-UAV Distributed Training Configuration"
echo "==============================================="
echo "Nodes: $NODES"
echo "GPUs per node: $GPUS_PER_NODE"
echo "Total world size: $WORLD_SIZE"
echo "Node rank: $NODE_RANK"
echo "Master address: $MASTER_ADDR:$MASTER_PORT"
echo "Data config: $DATA_CONFIG"
echo "Model: $MODEL"
echo "Epochs: $EPOCHS"
echo "Batch size per GPU: $BATCH_SIZE"
echo "Total batch size: $((BATCH_SIZE * WORLD_SIZE))"
echo "Workers per GPU: $WORKERS"
echo "Output: $PROJECT/$NAME"
echo "==============================================="

# Performance optimizations
export OMP_NUM_THREADS=1
export CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((GPUS_PER_NODE-1)))

# Launch training
echo "Starting distributed training..."

# Single node training
if [[ $NODES -eq 1 ]]; then
    python -m torch.distributed.launch \
        --nproc_per_node=$GPUS_PER_NODE \
        --master_port=$MASTER_PORT \
        tools/train_distributed.py \
        --data "$DATA_CONFIG" \
        --model "$MODEL" \
        --epochs $EPOCHS \
        --batch-size $BATCH_SIZE \
        --workers $WORKERS \
        --project "$PROJECT" \
        --name "$NAME" \
        --world-size $WORLD_SIZE \
        --amp \
        --sync-bn \
        --wandb
else
    # Multi-node training
    python -m torch.distributed.launch \
        --nnodes=$NODES \
        --node_rank=$NODE_RANK \
        --nproc_per_node=$GPUS_PER_NODE \
        --master_addr=$MASTER_ADDR \
        --master_port=$MASTER_PORT \
        tools/train_distributed.py \
        --data "$DATA_CONFIG" \
        --model "$MODEL" \
        --epochs $EPOCHS \
        --batch-size $BATCH_SIZE \
        --workers $WORKERS \
        --project "$PROJECT" \
        --name "$NAME" \
        --world-size $WORLD_SIZE \
        --amp \
        --sync-bn \
        --wandb
fi

echo "Training completed!"