#!/usr/bin/env bash
set -eo pipefail
shopt -s nullglob globstar

# define TMPDIR, if it's empty
if [[ -z "$TMPDIR" ]]; then
    TMPDIR="/tmp"
fi
echo "TMPDIR: $TMPDIR"

# activate conda env
eval "$(conda shell.bash hook)"
conda activate mopoe
echo "CONDA_PREFIX: $CONDA_PREFIX"

METHOD="joint_elbo"  # NOTE: valid options are "joint_elbo", "poe", and "moe"
LIKELIHOOD_M1="laplace"
LIKELIHOOD_M2="categorical"
DIR_DATA="$PWD/data/CelebA"
DIR_TEXT="$PWD/data/CelebA"
DIR_CLF="$PWD/trained_classifiers/trained_clfs_celeba"
DIR_EXPERIMENT="$PWD/runs/CelebA/$METHOD/factorized/${LIKELIHOOD_M1}_${LIKELIHOOD_M2}"
PATH_INC_V3="$PWD/pt_inception-2015-12-05-6726825d.pth"
DIR_FID="$TMPDIR/CelebA"

# move CelebA dataset to $TMPDIR
mkdir -p "${TMPDIR}/CelebA"
cp "${DIR_DATA}/img_align_celeba.zip" "${TMPDIR}/CelebA/"
cp "${DIR_DATA}/list_eval_partition.csv" "${TMPDIR}/CelebA/"
cp "${DIR_DATA}/list_attr_celeba.csv" "${TMPDIR}/CelebA/"
unzip -q -o "${TMPDIR}/CelebA/img_align_celeba.zip" -d "${TMPDIR}/CelebA/"
# ls "${TMPDIR}/CelebA/img_align_celeba/"

python main_celeba.py --dir_data=$TMPDIR \
	              --dir_text="$DIR_TEXT" \
	              --dir_clf="$DIR_CLF" \
	              --dir_experiment="$DIR_EXPERIMENT" \
	              --inception_state_dict="$PATH_INC_V3" \
	              --dir_fid=$DIR_FID \
                --method=$METHOD \
	              --beta=2.5 \
	              --beta_style=2.0 \
	              --beta_content=1.0 \
	              --beta_m1_style=1.0 \
	              --beta_m2_style=5.0 \
	              --div_weight_m1_content=0.35 \
	              --div_weight_m2_content=0.35 \
	              --div_weight_uniform_content=0.3 \
	              --likelihood_m1=$LIKELIHOOD_M1 \
	              --likelihood_m2=$LIKELIHOOD_M2 \
	              --batch_size=256 \
	              --initial_learning_rate=0.0005 \
	              --eval_freq=25 \
	              --eval_freq_fid=25 \
	              --end_epoch=250 \
	              --factorized_representation \
	              --calc_nll \
	              --eval_lr \
	              --calc_prd \
	              --use_clf
	


