#!/bin/bash
#SBATCH --job-name=aesthetic_eval
#SBATCH --output=../logs/aesthetic_eval_logs/%A_%a.out
#SBATCH --error=../logs/aesthetic_eval_logs/%A_%a.err
#SBATCH --nodes=1
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=10
#SBATCH --mem=100G
#SBATCH --time=01:00:00
#SBATCH --array=0-319%10

echo "Starting aesthetic evaluation..."

# Load necessary modules (if required)
# conda deactivate
source activate t2i

# Define the lists
guidances=(3 5 7 9 13)
complexities=(0 1 2 3)
type_guidances=(CFG APG CADS Interval)
model_ids=(SD15 SDXL SD35M SD35L)

# Calculate the total number of combinations
total_combinations=$(( ${#guidances[@]} * ${#complexities[@]} * ${#type_guidances[@]} * ${#model_ids[@]} ))

# Get the task ID from SLURM
task_id=${SLURM_ARRAY_TASK_ID}

# Ensure the task ID is within the valid range
if [ "$task_id" -ge "$total_combinations" ]; then
    echo "Error: Task ID $task_id is out of range (0 to $((total_combinations - 1)))."
    exit 1
fi

# Compute indices for each parameter
guidance_index=$(( task_id / (${#complexities[@]} * ${#type_guidances[@]} * ${#model_ids[@]}) ))
remainder=$(( task_id % (${#complexities[@]} * ${#type_guidances[@]} * ${#model_ids[@]}) ))

complexity_index=$(( remainder / (${#type_guidances[@]} * ${#model_ids[@]}) ))
remainder=$(( remainder % (${#type_guidances[@]} * ${#model_ids[@]}) ))

type_guidance_index=$(( remainder / ${#model_ids[@]} ))
model_id_index=$(( remainder % ${#model_ids[@]} ))

# Get the actual values
guidance=${guidances[$guidance_index]}
complexity=${complexities[$complexity_index]}
type_guidance=${type_guidances[$type_guidance_index]}
model_id=${model_ids[$model_id_index]}

# Define the source and destination directories
source_dir="../../output/SDinference/SDinference_gemma3_siglip_clip_EulerDiscrete_${type_guidance}/guidance_${guidance}.0/model_${model_id}/complexity_${complexity}"
destination_dir="/tmp/jobid_${SLURM_JOB_ID}"

# Create the destination directory if it doesn't exist
mkdir -p "$destination_dir"

echo ${model_id}_c${complexity}_tg${type_guidance}_g${guidance}.0


if [ "$model_id" != "DATA" ]; then
    # Copy and unzip the zip files for chunk_id from 0 to 50
    for chunk_id in {0..50}; do
        src_file="${source_dir}/${model_id}_c${complexity}_tg${type_guidance}_g${guidance}.0_chunk${chunk_id}.zip"
        if [ -f "$src_file" ]; then
            cp "$src_file" "$destination_dir"
            unzip -q -o "${destination_dir}/${model_id}_c${complexity}_tg${type_guidance}_g${guidance}.0_chunk${chunk_id}.zip" -d "$destination_dir"
        else
            echo "Warning: File $src_file does not exist."
        fi
    done

    echo "Generated images files copied and unzipped to $destination_dir."
else
    # Export the number of CPUs allocated by SLURM
    num_cpus=${SLURM_CPUS_PER_TASK}

    # Use GNU parallel to copy and unzip files concurrently
    seq 0 8 | parallel -j "$num_cpus" '
        src_file="../metadata/cc12m/eval_imgs_{}.zip"
        if [ -f "$src_file" ]; then
            cp "$src_file" "'"$destination_dir"'"
            unzip -q -o "'"$destination_dir"'/eval_imgs_{}.zip" -d "'"$destination_dir"'"
        else
            echo "Warning: File $src_file does not exist."
        fi
    '

    echo "Eval images files copied and unzipped to $destination_dir."
fi

echo "Starting Python evalution script..."
# Run the Python command
python aesthetic_evaluator.py ${model_id}_c${complexity}_tg${type_guidance}_g${guidance}.0 $complexity SD_aesthetic_scores