#!/bin/bash
{
source ~/.bashrc
source activate SLfM

# -------------------------- Training: wssl -------------------------- #
echo "support DDP";
read -p "input CUDA: " CUDA

echo "Start training wssl task...";
# -------------------------- loss -------------------------- #
read -p "input RECIPE: " RECIPE
read -p "input BATCH_SIZE: " BATCH_SIZE


# AUDIO_SOURCE='FMA'

SAVE=''
# SAVE=' --save_every_epoch '

if [ "$RECIPE" -eq 1 ]; then
  WEIGHTS=''
  LOSS_TYPE='ve012' # ve012 for mean reduction
  N_SS_POSITION_CLF=36
  RESTRAIN_TO_FRONT='--restrain_to_front';RTF='rtf'
  RESTRAIN_TO_FRONT='';RTF=''
  BATCH_SIZE=36
  # AZIMUTH_LOSS_TYPE='--azimuth_loss_type=classification';RG='clf'
  AZIMUTH_LOSS_TYPE='--azimuth_loss_type=regression';RG='rg'
elif [ "$RECIPE" -eq 2 ]; then
  LOSS_TYPE='vd01'
  N_SS_POSITION_CLF=1
  RESTRAIN_TO_FRONT='--restrain_to_front';RTF='rtf'
  BATCH_SIZE=36
  AZIMUTH_LOSS_TYPE='--azimuth_loss_type=classification';RG='clf'
elif [ "$RECIPE" -eq 3 ]; then
  LOSS_TYPE='v0041' # v00213;v0041
  N_SS_POSITION_CLF=36
  RESTRAIN_TO_FRONT='--restrain_to_front';RTF='rtf'
  # BATCH_SIZE=36

  WEIGHTS='--weights=HM3D/wsSL/simu-bl/ve012-rtf--clf-rough_n36-LbS-r0.2-2ssp-Length0.63-lre-4/ws_best.pth.tar'
  AZIMUTH_LOSS_TYPE='--azimuth_loss_type=classification';RG='clf'

  # WEIGHTS='--weights=HM3D/wsSL/simu-ws/ve012-rtf--rg-rough_n36-LbS-r0.2-2ssp-Length0.63-lre-4/ws_best.pth.tar'
  # AZIMUTH_LOSS_TYPE='--azimuth_loss_type=regression';RG='rg'
elif [ "$RECIPE" -eq 5 ]; then
  LOSS_TYPE='v0050' # v00213;v0041
  N_SS_POSITION_CLF=36
  # RESTRAIN_TO_FRONT='';RTF=''
  RESTRAIN_TO_FRONT='--restrain_to_front';RTF='rtf'
  # BATCH_SIZE=36

  # WEIGHTS='--weights=HM3D/wsSL/simu-bl/ve012---rg-rough_n36-LbS-r0.2-2ssp-Length0.63-lre-4/ws_best.pth.tar'
  # WEIGHTS=''
  AZIMUTH_LOSS_TYPE='--azimuth_loss_type=classification';RG='clf'
  AZIMUTH_LOSS_TYPE='--azimuth_loss_type=regression';RG='rg'
elif [ "$RECIPE" -eq 6 ]; then
  LOSS_TYPE='ve01' # v00213;v0041
  # RESTRAIN_TO_FRONT='';RTF=''
  # BATCH_SIZE=36
  # WEIGHTS='--weights=HM3D/wsSL/simu-bl/ve012---rg-rough_n36-LbS-r0.2-2ssp-Length0.63-lre-4/ws_best.pth.tar'
  # WEIGHTS=''
  N_SS_POSITION_CLF=32
  AZIMUTH_LOSS_TYPE='--azimuth_loss_type=classification';RG='clf'
  RESTRAIN_TO_FRONT='--restrain_to_front';RTF='rtf'
elif [ "$RECIPE" -eq 7 ]; then
  LOSS_TYPE='v002131' # v00213;v0041
  # RESTRAIN_TO_FRONT='';RTF=''
  # BATCH_SIZE=36
  # WEIGHTS='--weights=HM3D/wsSL/simu-bl/ve012---rg-rough_n36-LbS-r0.2-2ssp-Length0.63-lre-4/ws_best.pth.tar'
  # WEIGHTS=''
  read -p "input N_SS_POSITION_CLF: " N_SS_POSITION_CLF
  RESTRAIN_TO_FRONT='--restrain_to_front';RTF='rtf'
  BATCH_SIZE=600
  AZIMUTH_LOSS_TYPE='--azimuth_loss_type=classification';RG='clf'
elif [ "$RECIPE" -eq 8 ]; then
  LOSS_TYPE='ve01' # v00213;v0041
  N_SS_POSITION_CLF=30
  RESTRAIN_TO_FRONT='--restrain_to_front';RTF='rtf'
  BATCH_SIZE=600
  AZIMUTH_LOSS_TYPE='--azimuth_loss_type=classification';RG='clf'
elif [ "$RECIPE" -eq 9 ]; then
  LOSS_TYPE='ve03' # v00213;v0041
  read -p "input N_SS_POSITION_CLF: " N_SS_POSITION_CLF
  RESTRAIN_TO_FRONT='--restrain_to_front';RTF='rtf'
  BATCH_SIZE=600
  AZIMUTH_LOSS_TYPE='--azimuth_loss_type=classification';RG='clf'
elif [ "$RECIPE" -eq 0 ]; then
  LOSS_TYPE='v0060' 
  read -p "input N_SS_POSITION_CLF: " N_SS_POSITION_CLF
  # RESTRAIN_TO_FRONT='--restrain_to_front';RTF='rtf'
  BATCH_SIZE=600
  AZIMUTH_LOSS_TYPE='--azimuth_loss_type=classification';RG='clf'
elif [ "$RECIPE" -eq 10 ]; then
  LOSS_TYPE='v0061' 
  read -p "input N_SS_POSITION_CLF: " N_SS_POSITION_CLF
  # RESTRAIN_TO_FRONT='--restrain_to_front';RTF='rtf'
  BATCH_SIZE=600
  AZIMUTH_LOSS_TYPE='--azimuth_loss_type=classification';RG='clf'
elif [ "$RECIPE" -eq 11 ]; then
  LOSS_TYPE='v00612' 
  read -p "input N_SS_POSITION_CLF: " N_SS_POSITION_CLF
  # RESTRAIN_TO_FRONT='--restrain_to_front';RTF='rtf'
  BATCH_SIZE=600
  AZIMUTH_LOSS_TYPE='--azimuth_loss_type=classification';RG='clf'
elif [ "$RECIPE" -eq 12 ]; then
  LOSS_TYPE='v002131' 
  N_SS_POSITION_CLF=30 
  RESTRAIN_TO_FRONT='--restrain_to_front';RTF='rtf'
  BATCH_SIZE=600
  AZIMUTH_LOSS_TYPE='--azimuth_loss_type=classification --add_overlap --add_intermittent';RG='clf-ol-itm'
elif [ "$RECIPE" -eq 13 ]; then
  LOSS_TYPE='v0061' 
  N_SS_POSITION_CLF=30 
  RESTRAIN_TO_FRONT='--restrain_to_front';RTF='rtf'
  BATCH_SIZE=600
  AZIMUTH_LOSS_TYPE='--azimuth_loss_type=classification --add_overlap --add_intermittent';RG='clf-ol-itm'
else
  echo "Invalid value of 'RECIPE'"
  exit 1
fi






# -------------------------- Dataset -------------------------- #
DATASET_MOD='7'
read -p "input DATASET_MOD: " DATASET_MOD
CLIP_LENGTH=0.63
CLIP_LENGTH=1.27
# DATA_SETTING='hm3d-wsSL-r0.2-40degree';DS='rt40'
N_SS_POSITION=2


# CLIP_LENGTH=0.63
# # WEIGHTS='--weights=HM3D/wsSL/simu-ws/v002131-rtf--clf-rough_n36-LbS-r0.2-2ssp-Length2.55-lre-4/ws_best.pth.tar'
# DATASET_MOD='1'
# DATA_SETTING='Youtube-binaural-cars_scenes_v003';DS='uni3'
# N_SS_POSITION=4
# DATASET_FOLDER='yb-cars'

# CLIP_LENGTH=0.63
# # WEIGHTS='--weights=HM3D/wsSL/simu-ws/vc01-rtf--clf-rough_n1-LbS-r0.2-2ssp-Length0.63-lre-4/ws_best.pth.tar'
# DATASET_MOD='2'
# DATA_SETTING='Youtube-binaural_citywalk_v0027';DS='ds0027'
# N_SS_POSITION=4
# DATASET_FOLDER='yb-cw'

# CLIP_LENGTH=0.63

# DATASET_MOD='2'
# DATA_SETTING='Youtube-binaural_citywalk_v02003';DS='ds02003debug0e-7ssp4noild'
# N_SS_POSITION=4
# DATASET_FOLDER='yb-cw2'

# DATA_SETTING='huawei_mate_v01';DS='ds04'
# N_SS_POSITION=4 
# DATASET_FOLDER='yb-cw2'



if [ "$WEIGHTS" != '' ] && [ "$WITH_PRETRAIN" == '' ]; then
    WITH_PRETRAIN='pretrained-ws-'
else
    WITH_PRETRAIN=''
fi


if [ "$RESTRAIN_TO_FRONT" != '' ]; then
    RTF='rtf'
fi

# if [ "$CLIP_LENGTH" == 0.63 ]; then
#     BATCH_SIZE=200 # 400 for 20GB
# fi
if [ "$CLIP_LENGTH" == 0.31 ]; then
    BATCH_SIZE=150
fi

if [ "$DATASET_MOD" == '0' ]; then
    DATASET_FOLDER='simu'
    DATA_SETTING='hm3d-wsSL-r0.2';DS=''
    AUDIO_SOURCE='LibriSpeech'
fi

if [ "$DATASET_MOD" == '7' ]; then
    DATASET_FOLDER='avnerf'
    DATA_SETTING='hm3d-wsSL-r0.2';DS=''
    AUDIO_SOURCE=''
fi


if [ "$LOSS_TYPE" == "ve01" ] || [ "$LOSS_TYPE" == "vd01" ] || [ "$LOSS_TYPE" == "ve012" ]; then
    FOLDER="$DATASET_FOLDER-bl"
else
    FOLDER="$DATASET_FOLDER-ws"
fi



# -------------------------- Training: wssl -------------------------- #

EXP="HM3D/wsSL/${FOLDER}/${WITH_PRETRAIN}${LOSS_TYPE}-${RTF}-${DS}-${RG}-rough_n${N_SS_POSITION_CLF}-${AUDIO_SOURCE}-r0.2-${N_SS_POSITION}ssp-Length${CLIP_LENGTH}-lre-4"
SHARED_FLAG="--setting=wssl_hm3d --epochs=100 --vision_backbone=resnet18 --audio_backbone=resnet18 
--num_workers=45 --save_step=3 --valid_step=1 --lr=0.0001 --optim=AdamW 
--repeat=1 --schedule=cos --unet_input_nc=2 --unet_output_nc=2 --n_source=1 --cond_clip_length=5.10 --no_cond_audio --no_vision 
--eval --n_source=1 
--audiobase_path=data/AI-Habitat/data-split/${AUDIO_SOURCE}  
--n_view=1 --online_render --color_jitter --imagenet_pretrain "

# Rest of the command with variables
CUDA_VISIBLE_DEVICES=${CUDA} python main_wsSL.py \
${WEIGHTS} ${SHARED_FLAG} \
--clip_length ${CLIP_LENGTH} --n_ss_position_clf ${N_SS_POSITION_CLF} \
--exp=${EXP} ${SAVE} \
--loss_type=${LOSS_TYPE} --n_ss_position=${N_SS_POSITION} \
${AZIMUTH_LOSS_TYPE} --batch_size=${BATCH_SIZE} ${RESTRAIN_TO_FRONT} --dataset_mod ${DATASET_MOD} \
--data-setting ${DATA_SETTING}


# INPUT='ws'
# WEIGHTS="--weights=$EXP/ws_best.pth.tar"

# mkdir -p ./results/${EXP}

# CUDA_VISIBLE_DEVICES=${CUDA} python evaluation/evaluate_angle.py ${WEIGHTS} ${SHARED_FLAG} \
# --clip_length ${CLIP_LENGTH} --n_ss_position_clf ${N_SS_POSITION_CLF} \
# --exp=${EXP} --data-setting ${DATA_SETTING} \
# --loss_type=${LOSS_TYPE} --n_ss_position=${N_SS_POSITION} \
# ${AZIMUTH_LOSS_TYPE} --batch_size=${BATCH_SIZE} ${RESTRAIN_TO_FRONT} --dataset_mod ${DATASET_MOD} --input ${INPUT}

}