#!/bin/bash
set -e

# Load conda configuration and activate the brats environment
source /data1/nasiruddink/miniconda3/etc/profile.d/conda.sh
conda activate brats

# Point TensorFlow exactly to where the CUDA Toolkit is installed in the brats environment
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/data1/nasiruddink/miniconda3/envs/brats/lib:/data1/nasiruddink/miniconda3/envs/brats/lib/python3.10/site-packages/nvidia/cudnn/lib
export TF_CPP_MIN_LOG_LEVEL=2 # Hides the annoying TensorRT missing warnings

# ====================================================================================
# Experiment Definitions (Bands sorted by Cohen's d Separability)
# ====================================================================================

# 1. 12 Native 10m Features (Blue, Green, Red, NIR, NDVI, NDWI, NDSI, NDMI, BSI, Brightness, IronOxide, S1_VV)
RUN1_BANDS="band_0,band_1,band_2,band_3,band_10,band_11,band_12,band_13,band_14,band_15,band_16,band_27"
RUN1_NAME="run_12_native"

# 2. 14 Best Features
RUN2_BANDS="band_17,band_10,band_11,band_18,band_23,band_25,band_9,band_26,band_14,band_12,band_22,band_8,band_24,band_13"
RUN2_NAME="run_14_best"

# 3. 16 Best Features
RUN3_BANDS="band_17,band_10,band_11,band_18,band_23,band_25,band_9,band_26,band_14,band_12,band_22,band_8,band_24,band_13,band_2,band_1"
RUN3_NAME="run_16_best"

# 4. 18 Best Features
RUN4_BANDS="band_17,band_10,band_11,band_18,band_23,band_25,band_9,band_26,band_14,band_12,band_22,band_8,band_24,band_13,band_2,band_1,band_4,band_15"
RUN4_NAME="run_18_best"

# 5. 20 Best Features
RUN5_BANDS="band_17,band_10,band_11,band_18,band_23,band_25,band_9,band_26,band_14,band_12,band_22,band_8,band_24,band_13,band_2,band_1,band_4,band_15,band_0,band_27"
RUN5_NAME="run_20_best"

# 6. 22 Best Features
RUN6_BANDS="band_17,band_10,band_11,band_18,band_23,band_25,band_9,band_26,band_14,band_12,band_22,band_8,band_24,band_13,band_2,band_1,band_4,band_15,band_0,band_27,band_19,band_7"
RUN6_NAME="run_22_best"

# 7. 24 Best Features
RUN7_BANDS="band_17,band_10,band_11,band_18,band_23,band_25,band_9,band_26,band_14,band_12,band_22,band_8,band_24,band_13,band_2,band_1,band_4,band_15,band_0,band_27,band_19,band_7,band_3,band_21"
RUN7_NAME="run_24_best"

# 8. 26 Best Features
RUN8_BANDS="band_17,band_10,band_11,band_18,band_23,band_25,band_9,band_26,band_14,band_12,band_22,band_8,band_24,band_13,band_2,band_1,band_4,band_15,band_0,band_27,band_19,band_7,band_3,band_21,band_5,band_6"
RUN8_NAME="run_26_best"

# 9. All 28 Features
RUN9_BANDS="band_0,band_1,band_2,band_3,band_4,band_5,band_6,band_7,band_8,band_9,band_10,band_11,band_12,band_13,band_14,band_15,band_16,band_17,band_18,band_19,band_20,band_21,band_22,band_23,band_24,band_25,band_26,band_27"
RUN9_NAME="run_28_all"


# ====================================================================================
# Execution Engine (3 GPUs parallel)
# ====================================================================================
mkdir -p bandselection/logs

echo "🚀 Starting BATCH 1 (Runs 1, 2, 3) on GPUs 0, 1, 2..."
CUDA_VISIBLE_DEVICES=0 MODEL_NAME="attention_unet" DATA_DIR="dataset/patches_10m" BAND_LIST=$RUN1_BANDS RUN_NAME=$RUN1_NAME CKPT_DIR="bandselection/${RUN1_NAME}" python3 train.py > bandselection/logs/${RUN1_NAME}.log 2>&1 &
PID1=$!
CUDA_VISIBLE_DEVICES=1 MODEL_NAME="attention_unet" DATA_DIR="dataset/patches_10m" BAND_LIST=$RUN2_BANDS RUN_NAME=$RUN2_NAME CKPT_DIR="bandselection/${RUN2_NAME}" python3 train.py > bandselection/logs/${RUN2_NAME}.log 2>&1 &
PID2=$!
CUDA_VISIBLE_DEVICES=2 MODEL_NAME="attention_unet" DATA_DIR="dataset/patches_10m" BAND_LIST=$RUN3_BANDS RUN_NAME=$RUN3_NAME CKPT_DIR="bandselection/${RUN3_NAME}" python3 train.py > bandselection/logs/${RUN3_NAME}.log 2>&1 &
PID3=$!

echo "Waiting for Batch 1 to complete..."
wait $PID1 $PID2 $PID3
echo "✅ Batch 1 finished. Taking a 60 second break..."
sleep 60


echo "🚀 Starting BATCH 2 (Runs 4, 5, 6) on GPUs 0, 1, 2..."
CUDA_VISIBLE_DEVICES=0 MODEL_NAME="attention_unet" DATA_DIR="dataset/patches_10m" BAND_LIST=$RUN4_BANDS RUN_NAME=$RUN4_NAME CKPT_DIR="bandselection/${RUN4_NAME}" python3 train.py > bandselection/logs/${RUN4_NAME}.log 2>&1 &
PID1=$!
CUDA_VISIBLE_DEVICES=1 MODEL_NAME="attention_unet" DATA_DIR="dataset/patches_10m" BAND_LIST=$RUN5_BANDS RUN_NAME=$RUN5_NAME CKPT_DIR="bandselection/${RUN5_NAME}" python3 train.py > bandselection/logs/${RUN5_NAME}.log 2>&1 &
PID2=$!
CUDA_VISIBLE_DEVICES=2 MODEL_NAME="attention_unet" DATA_DIR="dataset/patches_10m" BAND_LIST=$RUN6_BANDS RUN_NAME=$RUN6_NAME CKPT_DIR="bandselection/${RUN6_NAME}" python3 train.py > bandselection/logs/${RUN6_NAME}.log 2>&1 &
PID3=$!

echo "Waiting for Batch 2 to complete..."
wait $PID1 $PID2 $PID3
echo "✅ Batch 2 finished. Taking a 60 second break..."
sleep 60


echo "🚀 Starting BATCH 3 (Runs 7, 8, 9) on GPUs 0, 1, 2..."
CUDA_VISIBLE_DEVICES=0 MODEL_NAME="attention_unet" DATA_DIR="dataset/patches_10m" BAND_LIST=$RUN7_BANDS RUN_NAME=$RUN7_NAME CKPT_DIR="bandselection/${RUN7_NAME}" python3 train.py > bandselection/logs/${RUN7_NAME}.log 2>&1 &
PID1=$!
CUDA_VISIBLE_DEVICES=1 MODEL_NAME="attention_unet" DATA_DIR="dataset/patches_10m" BAND_LIST=$RUN8_BANDS RUN_NAME=$RUN8_NAME CKPT_DIR="bandselection/${RUN8_NAME}" python3 train.py > bandselection/logs/${RUN8_NAME}.log 2>&1 &
PID2=$!
CUDA_VISIBLE_DEVICES=2 MODEL_NAME="attention_unet" DATA_DIR="dataset/patches_10m" BAND_LIST=$RUN9_BANDS RUN_NAME=$RUN9_NAME CKPT_DIR="bandselection/${RUN9_NAME}" python3 train.py > bandselection/logs/${RUN9_NAME}.log 2>&1 &
PID3=$!

echo "Waiting for Batch 3 to complete..."
wait $PID1 $PID2 $PID3

echo "🎉 ALL EXPERIMENTS COMPLETED SUCCESSFULLY!"