#!/bin/bash

# Main training script
TRAIN_SCRIPT="python prune_main.py"

# Common parameters shared across all runs
COMMON_PARAMS="
               --load_dir your_model_path \
               --arch resnet18_half \
               --dataset celeba \
               --dataset_dir your_dataset_path \
               --target_attr Attractive \
               --layer_wise_imp \
               --batch_size 256 \
               --train_data_ratio 0.1 \
               --p-lr 0.1 \
               --ft_epochs 20 \
               --noadv_add_schedular \
               "

# if you want to use adversarial training, you can add the following parameters
# --adv_mode \
# --p-lr 0.0001 \
# --a-lr 0.0001 \
# --w 0.5 \

# Pruning methods to test (sequential execution)
PRUNE_METHODS=("l1" "FPGM")

# Seed and GPU mapping (parallel execution within each method)
SEEDS=(2021 2022 2023 2024)
GPUS=(0 1 2 3)

# Pruning configuration
PRUNING_RATIO=0.9  # Set your pruning ratio here
ITERATIVE_STEPS=9  # Set your iterative steps here

# Outer loop: Sequentially execute each pruning method
for METHOD in "${PRUNE_METHODS[@]}"; do
    echo "================================================================"
    echo " Starting pruning method: $METHOD "
    echo "================================================================"
    
    # Inner loop: Parallel execution for each seed
    for i in "${!SEEDS[@]}"; do
        SEED=${SEEDS[$i]}
        GPU=${GPUS[$i]}
        
        # Generate dynamic save path
        SAVE_PATH=YOUR_DYNAMIC_SAVE_PATH
        
        echo "Starting $METHOD with seed $SEED on GPU $GPU"
        echo "Save path: $SAVE_PATH"
        
        $TRAIN_SCRIPT $COMMON_PARAMS \
            --gpu $GPU \
            --random_seed $SEED \
            --pruning_ratio $PRUNING_RATIO \
            --iterative_steps $ITERATIVE_STEPS \
            --prune_methods $METHOD \
            --save_dir $SAVE_PATH &
    done
    
    # Wait for all parallel jobs of current method to finish
    wait
    echo "All $METHOD pruning tasks completed."
    echo ""
done

echo "================================================================"
echo " All pruning methods executed successfully! "
echo "================================================================"