#!/bin/bash

. /data/cluster_name/scratch/$(whoami)/projects/MACE-Update/slurm/scripts/prelude.sh

ALGO_NAME=$2
CHANGE=$3
ORIG_TASK=$4
ORIG_CONFIG=$5
FINETUNE_ALGO=$6
FINETUNE_TASK=$7
FINETUNE_CONFIG=$8
PORT_NUMBER=$9
PROMPTS_CSV=${10}
RANDOM_SEED=${11}

mkdir -p $PREFIX_SCRATCH/experiments/${ALGO_NAME}/${CHANGE}_${ORIG_TASK}_${ORIG_CONFIG}/${RANDOM_SEED}/finetune/${FINETUNE_ALGO}/${FINETUNE_TASK}_${FINETUNE_CONFIG}/inference/${PROMPTS_CSV}
cp /data/cluster_name/scratch/$(whoami)/projects/MACE-Update/tasks/${FINETUNE_TASK}/${PROMPTS_CSV}.csv $PREFIX_SCRATCH/experiments/${ALGO_NAME}/${CHANGE}_${ORIG_TASK}_${ORIG_CONFIG}/${RANDOM_SEED}/finetune/${FINETUNE_ALGO}/${FINETUNE_TASK}_${FINETUNE_CONFIG}/inference/${PROMPTS_CSV}/prompts.csv 

if [ $FINETUNE_ALGO == "full" ]; then

    CUDA_VISIBLE_DEVICES=$1 accelerate launch \
            --multi_gpu --num_processes=2 --main_process_port $PORT_NUMBER \
            $PREFIX/inference/sample_images_from_csv.py \
            --prompts_path /data/cluster_name/scratch/$(whoami)/projects/MACE-Update/tasks/${FINETUNE_TASK}/${PROMPTS_CSV}.csv \
            --save_path $PREFIX_SCRATCH/experiments/${ALGO_NAME}/${CHANGE}_${ORIG_TASK}_${ORIG_CONFIG}/${RANDOM_SEED}/finetune/${FINETUNE_ALGO}/${FINETUNE_TASK}_${FINETUNE_CONFIG}/inference/${PROMPTS_CSV} \
            --model_name $PREFIX_SCRATCH/experiments/${ALGO_NAME}/${CHANGE}_${ORIG_TASK}_${ORIG_CONFIG}/${RANDOM_SEED}/finetune/${FINETUNE_ALGO}/${FINETUNE_TASK}_${FINETUNE_CONFIG} \
            --step 2

elif [ $FINETUNE_ALGO == "lora" ]; then

    CUDA_VISIBLE_DEVICES=$1 accelerate launch \
            --multi_gpu --num_processes=2 --main_process_port $PORT_NUMBER \
            $PREFIX/inference/sample_images_from_csv.py \
            --prompts_path /data/cluster_name/scratch/$(whoami)/projects/MACE-Update/tasks/${FINETUNE_TASK}/${PROMPTS_CSV}.csv \
            --save_path $PREFIX_SCRATCH/experiments/${ALGO_NAME}/${CHANGE}_${ORIG_TASK}_${ORIG_CONFIG}/${RANDOM_SEED}/finetune/${FINETUNE_ALGO}/${FINETUNE_TASK}_${FINETUNE_CONFIG}/inference/${PROMPTS_CSV} \
            --model_name $PREFIX/experiments/${ALGO_NAME}/${CHANGE}_${ORIG_TASK}_${ORIG_CONFIG}/${RANDOM_SEED}/LoRA_fusion_model \
            --lora_path $PREFIX/experiments/${ALGO_NAME}/${CHANGE}_${ORIG_TASK}_${ORIG_CONFIG}/${RANDOM_SEED}/finetune/${FINETUNE_ALGO}/${FINETUNE_TASK}_${FINETUNE_CONFIG}/pytorch_lora_weights.safetensors \
            --step 2
else
    echo "Finetuning algorithm not supported"
fi