# fill in your own
user=YOUR_WANDB_USERNAME

export WANDB_INIT_TIMEOUT=120
epochs=10
dimX=5
dimZ=768
bs=500
num_loader_obs=500
postprocess_often=0
use_X_model=True
mlp_width=100
mlp_layer=3
x_mlp_width=100
x_mlp_layer=2
sample_frac=1
prefix='fs_'
method='0111_logistic_withZ_zero_prodsign_last' 
check_val_every_n_rows=10000
suffstat_eps=1
mlp_width=100
ms='sequential'
repeat_suffstat=100
use_dataset_Y=True
mlp_layer=3
D=20000
epochs=40

for lr in 0.001 0.01 0.02 0.0001
do
    data_dir="/shared/share_mala/implicitbayes/dataset_files/MIND_data/large/N=1000,D=${D},D_eval=10000,method=${method},dimX=5,one_X_per_col=False,flip//bert_Z/"
    save_name="semisynthetic_${method}"
    echo "$save_name $method $dimZ"
    python ../train_models.py --data_dir $data_dir \
        --check_val_every_n_rows $check_val_every_n_rows \
        --epochs $epochs --num_loader_obs $num_loader_obs \
        --dataset_type synthetic \
        --prefix $prefix \
        --X_dim $dimX \
        --use_dataset_Y $use_dataset_Y \
        --Z_dim $dimZ \
        --sample_frac $sample_frac \
        --save_name $save_name --wandb_user $user --marginal_vs_sequential $ms \
        --MLP_layer $mlp_layer \
        --MLP_width $mlp_width \
        --X_MLP_layer $x_mlp_layer \
        --X_MLP_width $x_mlp_width \
        --suffstat_eps $suffstat_eps \
        --batch_size $bs --eval_batch_size $bs \
        --learning_rate $lr --repeat_suffstat $repeat_suffstat --use_X_model $use_X_model --postprocess_often $postprocess_often
sleep 0.2
done






epochs=20
D=40000 
for lr in 0.001 0.01 0.02 0.0001
do
data_dir="/shared/share_mala/implicitbayes/dataset_files/MIND_data/large/N=1000,D=${D},D_eval=10000,method=${method},dimX=5,one_X_per_col=False,flip//bert_Z/"
    save_name="even_use_Y_0111_${D}_${sample_frac}_context_semisynthetic_fixed_bert_${method}_${ms}_context_dimX=${dimX}"
echo "$save_name $method $dimZ"
python ../train_models.py --data_dir $data_dir \
    --check_val_every_n_rows $check_val_every_n_rows \
    --use_X True --epochs $epochs --num_loader_obs $num_loader_obs \
        --dataset_type synthetic \
        --prefix $prefix \
        --X_dim $dimX \
        --use_dataset_Y $use_dataset_Y \
        --Z_dim $dimZ \
        --sample_frac $sample_frac \
        --save_name $save_name --wandb_user $user --marginal_vs_sequential $ms \
        --MLP_layer $mlp_layer \
        --MLP_width $mlp_width \
        --X_MLP_layer $x_mlp_layer \
        --X_MLP_width $x_mlp_width \
        --suffstat_eps $suffstat_eps \
        --batch_size $bs --eval_batch_size $bs \
        --dataset_type synthetic --learning_rate $lr --repeat_suffstat $repeat_suffstat --use_X_model $use_X_model --postprocess_often $postprocess_often
sleep 0.2
done


