model_idx_stage2=9999
#model_idx_stage3=2999
first_stage_folder_name=chembl_remove_properties_latent64
second_stage_folder_name=chembl_remove_properties_latent64_layer5
third_stage_folder_name=chembl_remove_properties_latent64_layer5_stage3
for model_idx_stage3 in 6000
do
    for seed in 1 2 3 4 5 6
    do
	python train_vae_latent.py $third_stage_folder_name --encoder_layer_sizes "1024|1024|1024|1024|1024" --layer_sizes "1024|1024|1024|1024|1024" -e -5 --seed $seed --square --data_file outputs/${first_stage_folder_name}/latent_mean_var --sample model.${model_idx_stage3}.pkl --output_latent_file stage3_latent_samples.${model_idx_stage3}.${seed}.npy
	cp 2_stage_vae/${third_stage_folder_name}/stage3_latent_samples.${model_idx_stage3}.${seed}.npy 2_stage_vae/${second_stage_folder_name}/
	python train_vae_latent.py $second_stage_folder_name --encoder_layer_sizes "1024|1024|1024|1024|1024" --layer_sizes "1024|1024|1024|1024|1024" -e -5 --square --seed $seed --data_file outputs/${first_stage_folder_name}/latent_mean_var --sample model.${model_idx_stage2}.pkl --output_latent_file stage2_latent_samples.${model_idx_stage2}.${seed}.npy --input_latent_file stage3_latent_samples.${model_idx_stage3}.${seed}.npy
	rm outputs/$first_stage_folder_name/*.npy
	cp 2_stage_vae/${second_stage_folder_name}/stage2_latent_samples.${model_idx_stage2}.${seed}.npy outputs/${first_stage_folder_name}/
	molecule_generation sample outputs/${first_stage_folder_name}/ 10000 > outputs/${first_stage_folder_name}/3stage_samples_${model_idx_stage2}_${model_idx_stage3}.${seed}.txt
	python outputs/generate_metric.py --message "3 stage with stage 2 model ${model_idx_stage2}, stage 3 model ${model_idx_stage3} and seed ${seed}" --folder outputs/${first_stage_folder_name}/ --sample_file outputs/${first_stage_folder_name}/3stage_samples_${model_idx_stage2}_${model_idx_stage3}.${seed}.txt --metric_file stage3_metric.${model_idx_stage3}.txt
    done
done

