#!/bin/bash

# Define directories and files
TRAIN_DIR="benchmark_data/lora_format_test"
OUTPUT_EXP="/cm/archive/trilq3/exps/sota_comparisons/lora_sdxl"
EVAL_OUTPUT_DIR="/cm/archive/trilq3/output/sota_comparisons/lora_sdxl_fix_placeholder_fix_fix"
EVALUATION_PROMPT_FILE="benchmark_data/benchmark_prompts.json"

# Get list of directories and count them
FOLDERS=("$TRAIN_DIR"/*)
TOTAL_DIRS=${#FOLDERS[@]}
HALF=$((TOTAL_DIRS / 2))

echo "Processing first half: 0 to $((HALF-1))"

# Loop through first half of directories
for ((i = 0; i < HALF; i++)); do
    folder="${FOLDERS[$i]}"
    if [[ -d "$folder" ]]; then
        echo "Processing: $folder"

        FOLDER_NAME=$(basename "$folder")
        INSTANCE_DATA_DIR="$folder/imgs"

        PLACEHOLDER_STYLE=""
        PLACEHOLDER_CONTENT=""

        for type in "content" "style"; do
            OUTPUT_DIR="$OUTPUT_EXP/$type/$FOLDER_NAME"

            PROMPT_FILE="$folder/instance_prompt_$type.txt"
            if [[ ! -f "$PROMPT_FILE" ]]; then
                echo "Warning: Prompt file not found for $FOLDER_NAME, skipping..."
                continue
            fi

            INSTANCE_PROMPT=$(<"$PROMPT_FILE")
            echo "INSTANCE_PROMPT: $INSTANCE_PROMPT"

            PLACEHOLDER=$(echo "$INSTANCE_PROMPT" | cut -d' ' -f2-)
            echo "PLACEHOLDER: $PLACEHOLDER"

            if [[ "$type" == "content" ]]; then
                PLACEHOLDER_CONTENT="$PLACEHOLDER"
            else
                PLACEHOLDER_STYLE="$PLACEHOLDER"
            fi


            if [[ -d "$OUTPUT_DIR" ]]; then
                echo "Output directory exists for $FOLDER_NAME, skipping training..."
                continue
            fi

            mkdir -p "$OUTPUT_DIR"

            
            CUDA_VISIBLE_DEVICES=0 accelerate launch --num_processes=1 --gpu_ids="all" train_dreambooth_lora_sdxl.py \
                --pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0" \
                --instance_data_dir="$INSTANCE_DATA_DIR" \
                --output_dir="$OUTPUT_DIR" \
                --instance_prompt="$INSTANCE_PROMPT" \
                --resolution=1024 \
                --rank=64 \
                --train_batch_size=1 \
                --learning_rate=1e-4 \
                --lr_scheduler="constant" \
                --lr_warmup_steps=0 \
                --max_train_steps=200 \
                --checkpointing_steps=200 \
                --seed=0 \
                --gradient_checkpointing \
                --use_8bit_adam \
                --mixed_precision="fp16" \
                --enable_xformers_memory_efficient_attention
        done

        OUTPUT_RESULT_DIR="$EVAL_OUTPUT_DIR"
        mkdir -p "$OUTPUT_RESULT_DIR"

        start_time=$(date +%s)
        CUDA_VISIBLE_DEVICES=0 python infer_lora.py \
            --output_path="$OUTPUT_RESULT_DIR" \
            --content_LoRA="$OUTPUT_EXP/content/$FOLDER_NAME" \
            --content_alpha=1.0 \
            --style_LoRA="$OUTPUT_EXP/style/$FOLDER_NAME" \
            --style_alpha=1.0 \
            --num_images_per_prompt=8 \
            --name_concept="$FOLDER_NAME" \
            --placeholder_style="$PLACEHOLDER_STYLE" \
            --placeholder_content="$PLACEHOLDER_CONTENT" \
            --evaluation_prompt_file="$EVALUATION_PROMPT_FILE"
        end_time=$(date +%s)
        echo "Time taken for inference: $((end_time - start_time)) seconds"
    fi
done

# Run evaluation metrics script
python evaluate_with_metrics.py --result_dir "$EVAL_OUTPUT_DIR" --data_dir "$TRAIN_DIR"