#!/bin/bash
#SBATCH --job-name=dsg_35l       # Job name
#SBATCH --output=../../../outputs/evalgim_logs/sd35l-dsg-%A_%a.out  # Standard output log
#SBATCH --error=../../../outputs/evalgim_logs/sd35l-dsg-%A_%a.err   # Error log
#SBATCH --ntasks=1              # Number of tasks
#SBATCH --cpus-per-task=10      # Number of CPU cores per task
#SBATCH --gres=gpu:1            # Number of GPUs per node
#SBATCH --mem=100G               # Memory per node
#SBATCH --time=10:00:00         # Time limit (hrs:min:sec)
#SBATCH --array=0-79%30            # Array job

# Activate the Python environment
source activate evalgim

# Define variables
score_metric=dsg
model_evaluate=SD35L

metadata_path="../../../outputs/evalgim_metadata/"
data_path="../../../outputs/SDinference/"
gen_path="$data_path/SDinference_gemma3_siglip_clip_EulerDiscrete_"
output_path="$data_path/evalgim_outputs_SDmodels_$score_metric/"

echo "Creating output directory at $output_path"
mkdir -p $output_path/evals
mkdir -p $output_path/evals/evaluate

# Define the parameters
complexities=(0 1 2 3)
cfg_scales=(3 5 7 9 13)
schedulers=("CFG" "APG" "Interval" "CADS")

# Calculate total combinations
total_combinations=$(( ${#complexities[@]} * ${#cfg_scales[@]} * ${#schedulers[@]} ))

# Get the current combination based on SLURM_ARRAY_TASK_ID
idx=$SLURM_ARRAY_TASK_ID

# Calculate indices for each parameter
complexity_idx=$(( idx / (${#cfg_scales[@]} * ${#schedulers[@]}) ))
cfg_scale_idx=$(( (idx / ${#schedulers[@]}) % ${#cfg_scales[@]} ))
scheduler_idx=$(( idx % ${#schedulers[@]} ))

# Get the actual parameter values
complexity=${complexities[$complexity_idx]}
cfg_scale=${cfg_scales[$cfg_scale_idx]}
scheduler=${schedulers[$scheduler_idx]}

echo "Running evaluation with complexity: $complexity, cfg_scale: $cfg_scale, scheduler: $scheduler"

# Copy data to a temporary directory
echo "Copying data"
dest_dir="/tmp/job$SLURM_JOB_ID"
mkdir -p $dest_dir
gen_path_task="$gen_path$scheduler"
gen_folder="guidance_$cfg_scale.0/model_$model_evaluate/complexity_$complexity/"
echo "Copying zip files from $gen_path_task/$gen_folder to $dest_dir"
find "$gen_path_task/$gen_folder" -type f -name "*.zip" -exec cp {} "$dest_dir" \;

# Unzip files in the temporary directory
echo "Unzipping files in $dest_dir"
find "$dest_dir" -type f -name "*.zip" -exec unzip -q -o -d "$dest_dir" {} \;

echo "Generated images copied to $dest_dir"

# Run the Python script with the selected parameters
python evaluation_library/main.py \
    --name evaluate \
    --batch_size 2048 \
    --mixed_precision fp16 \
    --logdir $output_path/evals \
    --generated_images_path $metadata_path/evalgim_genldm_metadata_cc12m/dsg_c${complexity} \
    --model_id $model_evaluate \
    --cfg_scale $cfg_scale \
    --eval_dataset_name data.real_datasets.CC12MCLIP5KSIGLIPVal \
    --marginal_metrics none \
    --conditional_metrics $score_metric \
    --scheduler $scheduler \
    --complexity $complexity \
    --feature_ext dino
