#!/bin/bash

# EXCAP Training Script
# Usage: bash scripts/train_excap.sh [config_file] [gpu_id]

set -e  # Exit on any error

# Default values
CONFIG_FILE=${1:-"configs/excap_model.yaml"}
GPU_ID=${2:-0}
EXPERIMENT_NAME=${3:-"excap_experiment_$(date +%Y%m%d_%H%M%S)"}

echo "=========================================="
echo "EXCAP Multi-Modal Medical Prediction"
echo "=========================================="
echo "Config file: $CONFIG_FILE"
echo "GPU ID: $GPU_ID"
echo "Experiment: $EXPERIMENT_NAME"
echo "=========================================="

# Check if config file exists
if [ ! -f "$CONFIG_FILE" ]; then
    echo "Error: Config file '$CONFIG_FILE' not found!"
    echo "Available configs:"
    ls -la configs/*.yaml
    exit 1
fi

# Check if GPU is available
if command -v nvidia-smi &> /dev/null; then
    echo "GPU Information:"
    nvidia-smi --query-gpu=index,name,memory.total,memory.free --format=csv,noheader,nounits | head -1
else
    echo "Warning: nvidia-smi not found. GPU monitoring unavailable."
fi

# Set environment variables
export CUDA_VISIBLE_DEVICES=$GPU_ID
export PYTHONPATH="${PYTHONPATH}:$(pwd)"

# Create log directory
LOG_DIR="runs/$(basename "$CONFIG_FILE" .yaml)_$(date +%Y%m%d_%H%M%S)"
mkdir -p "$LOG_DIR"

echo "Log directory: $LOG_DIR"
echo "Starting training..."

# Run training
python training/train_EXCAP.py \
    --config "$CONFIG_FILE" \
    --gpu "$GPU_ID" \
    --experiment_name "$EXPERIMENT_NAME" \
    --log_dir "$LOG_DIR" \
    --use_cache \
    --freeze_encoders \
    --visualize

echo "=========================================="
echo "Training completed!"
echo "Results saved to: $LOG_DIR"
echo "=========================================="

# Optional: Start TensorBoard
if command -v tensorboard &> /dev/null; then
    echo "Starting TensorBoard..."
    echo "Run: tensorboard --logdir=$LOG_DIR"
    echo "Open: http://localhost:6006"
fi
