
frame_num=$1

prefix=sciml_
fea_weight=1e-6

gt_dir=''  # gt images dir
pretrained_pde_model=''

# ============================train FMv1=====================================
fm1_name=${prefix}swin_224_ps4_n${frame_num}_e201_ns10_rot3_roll_pretrained 
echo $fm1_name
ckpt=best_ckpt.tar 

yaml=./swin_transformer/train_swin_v0.yaml
max_epochs=201
n_steps_train=10
rollout_train=3

python swin_transformer/train_fm_swin.py --run_name $fm1_name --max_epochs $max_epochs --config basic_config \
--yaml_config $yaml  --frame_num $frame_num --n_steps $n_steps_train --rollout_train $rollout_train \
--ckpt $pretrained_pde_model

inf_name=inf_${fm1_name}

echo $inf_name

python swin_transformer/train_fm_swin.py \
--run_name $inf_name  --frame_num $frame_num \
--config basic_config --n_steps $n_steps_train  \
--yaml_config $yaml --rollout_train $rollout_train \
--inference --ckpt ./logs/exp/basic_config/${fm1_name}/training_checkpoints/$ckpt

inference_dir=./logs/exp/basic_config/$inf_name/output_frames/train02

python sciml_fm/scalarflow_metrics_single.py --fm_dir $inference_dir --gt_dir $gt_dir --frame_start ${n_steps_train} --frame_end $(($frame_num+20))
# =========================get features from FMv1===========================

ckpt=best_ckpt.tar

python sciml_fm/save_features.py --gt_dir ./data/ScalarReal --fm_yaml $yaml  --fm_load ./logs/exp/basic_config/${fm1_name}/training_checkpoints/$ckpt  --num_initial_frames  $frame_num   --num_tested_frames 30 --out_dir logs/features  


fm1_frame_dir=./logs/features/$fm1_name  # res of save feature
fm1_ckpt=./logs/exp/basic_config/$fm1_name/training_checkpoints/$ckpt
frame_fm_num=20

# ========================train HYv1==================================
hy1_expname=${prefix}agg+gt_HY${frame_num}+FM${frame_fm_num}_we-6_str1
echo $hy1_expname

density_exp="$hy1_expname/density_256_128"
joint_exp="$hy1_expname/p_v128_128"
joint_ft_path="./logs/$density_exp/best.tar"

vort_exp="$hy1_expname/vort50"
vort_ft_path="./logs/$joint_exp/den/best.tar"  #best
vort_vel_path="./logs/$joint_exp/best.tar"  #best
out_path="./logs/$vort_exp/best.tar"

python  sciml_fm/run_nerf_density_sciml.py --config configs/scalarflowreal.txt --lrate 0.01 --N_time 1  \
--lrate_decay 30000 --N_iters 100000 --i_weights 10000 --i_video 10000 --i_print 100 \
--expname $density_exp  --finest_resolution 256 \
--base_resolution 16 --finest_resolution_t 128 --base_resolution_t 16 --num_levels 16 --N_samples 192 --N_rand 256 --log2_hashmap_size 19 \
--out_frame_fm_dir $fm1_frame_dir  --fm_yaml $yaml  --fm_load $fm1_ckpt --input_features   --fea_weight $fea_weight \
--frame_num $frame_num --frame_num_fm $frame_fm_num  --no_reload

python sciml_fm/run_nerf_jointly_sciml.py --config configs/scalarflowreal.txt \
--lrate_decay 20000 --N_iters 100000 --i_weights 3000 --i_video 3000 --i_print 1000 \
--expname $joint_exp --lrate 5e-4 --lrate_den 1e-4  --rec_weight 10000 --d2v_weight 10 --coef_den2vel 0.2 --vel_weight 1 --d_weight 0 --proj_weight 1 --flow_weight 0.001 --vel_num_layers 2 --finest_resolution 256 --base_resolution 16 --finest_resolution_t 128 --base_resolution_t 16 --num_levels 16 --N_samples 192 --N_rand 512 --log2_hashmap_size 19 \
--ft_path $joint_ft_path  --vel_scale 0.025  \
--finest_resolution_v 128 --base_resolution_v 16 --finest_resolution_v_t 128 --base_resolution_v_t 16 --no_vel_der \
--out_frame_fm_dir $fm1_frame_dir  --fm_yaml $yaml --fm_load $fm1_ckpt --input_features   --fea_weight $fea_weight  \
--frame_num $frame_num --frame_num_fm $frame_fm_num  --no_reload

python sciml_fm/run_nerf_vort_sciml.py --config configs/scalarflowreal.txt --lrate 0.01 --lrate_den 1e-4 \
--lrate_decay 5000 --N_iters 10000 --i_weights 1000 --i_video 1000 --i_print 100 \
--expname $vort_exp --finest_resolution 256 --base_resolution 16 --finest_resolution_t 128 --base_resolution_t 16 --num_levels 16 --N_samples 192 --N_rand 128 --log2_hashmap_size 19 --rec_weight 10000 \
--rec_weight 10000 --vel_weight 0 --d_weight 0 --flow_weight 0.001 --vel_num_layers 2 \
--ft_path $vort_ft_path --vel_path $vort_vel_path  --no_vel_der --vel_scale 0.05 \
--finest_resolution_v 128 --base_resolution_v 16 --finest_resolution_v_t 128 --base_resolution_v_t 16 \
--n_particles 50 --vort_intensity 5 --vort_weight 0.01  \
--out_frame_fm_dir $fm1_frame_dir  --fm_yaml $yaml --fm_load $fm1_ckpt --input_features   --fea_weight $fea_weight \
--frame_num $frame_num --frame_num_fm $frame_fm_num  --no_reload


#========================metric=============================
#sys
python sciml_fm/scalarflow_metrics_single.py --hyfluid_dir  ./logs/$density_exp/testset_best --gt_dir $gt_dir --frame_start 0 --frame_end  $(($frame_num+20))  --out_name sys_$hy1_expname
#sim+pred
python sciml_fm/scalarflow_metrics_single.py --hyfluid_dir  ./logs/$vort_exp/run_advect_den_best --gt_dir $gt_dir --frame_start 0 --frame_end  $(($frame_num+20))  --out_name sim+pred_$hy1_expname


#========================prepare data for FMv1========================

python sciml_fm/run_nerf_vort_sciml.py --config configs/scalarflowreal.txt --lrate 0.01 --lrate_den 1e-4 \
--lrate_decay 5000 --N_iters 10000 --i_weights 1000 --i_video 1000 --i_print 100 \
--expname $vort_exp --finest_resolution 256 --base_resolution 16 --finest_resolution_t 128 --base_resolution_t 16 --num_levels 16 --N_samples 192 --N_rand 128 --log2_hashmap_size 19 --rec_weight 10000 \
--rec_weight 10000 --vel_weight 0 --d_weight 0 --flow_weight 0.001 --vel_num_layers 2 \
--no_vel_der --vel_scale 0.05 \
--finest_resolution_v 128 --base_resolution_v 16 --finest_resolution_v_t 128 --base_resolution_v_t 16 \
--n_particles 50 --vort_intensity 5 --vort_weight 0.01  \
--out_frame_fm_dir $fm1_frame_dir  --fm_yaml $yaml --fm_load $fm1_ckpt --input_features   --fea_weight $fea_weight \
--frame_num $frame_num --frame_num_fm $frame_fm_num  --run_advect_den


# =========================Train FMv1===============================
max_epochs=101 
n_steps_train=10
rollout_train=8
yaml=./swin_transformer/train_swin_v2.yaml

# if finetune
vort_exp="$hy1_expname/vort50"
hy1_dir=./logs/$vort_exp/run_advect_den_resim
ini_model=./logs/exp/basic_config/$fm1_name/training_checkpoints/best_ckpt.tar  
num_hy_frames=20

fm2_name=${prefix}agg+gt_n${frame_num}_e${max_epochs}_ns${n_steps_train}_rot${rollout_train}_finuetune_nfhy${num_hy_frames}_fea-e-6
echo $run_name

# train
python swin_transformer/train_fm_swin.py --run_name $fm2_name --max_epochs $max_epochs --config basic_config --yaml_config $yaml  --frame_num $frame_num --n_steps $n_steps_train --rollout_train $rollout_train \
--finetune  --ckpt $ini_model --hy_dir $hy1_dir --num_hy_frames $num_hy_frames # if finetune

inf_name=inf_${fm2_name}

echo $inf_name

python swin_transformer/train_fm_swin.py \
--run_name $inf_name --frame_num $frame_num \
--config basic_config --n_steps $n_steps_train \
--yaml_config $yaml --rollout_train $rollout_train \
--inference --ckpt ./logs/exp/basic_config/${fm2_name}/training_checkpoints/best_ckpt.tar \
--hy_dir $hy1_dir --num_hy_frames $num_hy_frames

inference_dir=./logs/exp/basic_config/$inf_name/output_frames/train02

python sciml_fm/scalarflow_metrics_single.py --fm_dir $inference_dir --gt_dir $gt_dir --frame_start ${n_steps_train} --frame_end $(($frame_num+10))

python sciml_fm/save_features.py --gt_dir ./data/ScalarReal --fm_yaml $yaml  --fm_load ./logs/exp/basic_config/${fm2_name}/training_checkpoints/best_ckpt.tar  --num_initial_frames  $frame_num   --num_tested_frames 20 --out_dir logs/features  

echo $inf_name 

# =========================Train HYv2===============================
ckpt=best_ckpt.tar
finetuned_ft_path_den=$joint_ft_path
finetuned_ft_path_joint=$vort_ft_path
finetuned_ft_path_vel=$vort_vel_path
frame_fm_num=20
hy2_expname="${prefix}agg+gt_HY${frame_num}+FM${frame_fm_num}_fea-e-6_finetune"
echo $hy2_expname

fm2_frame_dir=./logs/features/$fm2_name
fm2_ckpt=./logs/exp/basic_config/$fm2_name/training_checkpoints/$ckpt  

density_exp="$hy2_expname/density_256_128"
joint_exp="$hy2_expname/p_v128_128"
joint_ft_path="./logs/$density_exp/best.tar"

vort_exp="$hy2_expname/vort50"
vort_ft_path="./logs/$joint_exp/den/best.tar"
vort_vel_path="./logs/$joint_exp/best.tar"
out_path="./logs/$vort_exp/best.tar"

python  sciml_fm/run_nerf_density_sciml.py --config configs/scalarflowreal.txt --lrate 0.01 --N_time 1  \
--lrate_decay 30000 --N_iters 100000 --i_weights 10000 --i_video 10000 --i_print 100 \
--expname $density_exp  --finest_resolution 256 \
--base_resolution 16 --finest_resolution_t 128 --base_resolution_t 16 --num_levels 16 --N_samples 192 --N_rand 256 --log2_hashmap_size 19 \
--out_frame_fm_dir $fm2_frame_dir  --fm_yaml $yaml  --fm_load $fm2_ckpt --input_features   --fea_weight $fea_weight \
--frame_num $frame_num --frame_num_fm $frame_fm_num  \
--ft_path $finetuned_ft_path_den

python sciml_fm/run_nerf_jointly_sciml.py --config configs/scalarflowreal.txt \
--lrate_decay 20000 --N_iters 100000 --i_weights 3000 --i_video 3000 --i_print 1000 \
--expname $joint_exp --lrate 5e-4 --lrate_den 1e-4  --rec_weight 10000 --d2v_weight 10 --coef_den2vel 0.2 --vel_weight 1 --d_weight 0 --proj_weight 1 --flow_weight 0.001 --vel_num_layers 2 --finest_resolution 256 --base_resolution 16 --finest_resolution_t 128 --base_resolution_t 16 --num_levels 16 --N_samples 192 --N_rand 512 --log2_hashmap_size 19 \
--ft_path $joint_ft_path  --vel_scale 0.025  \
--finest_resolution_v 128 --base_resolution_v 16 --finest_resolution_v_t 128 --base_resolution_v_t 16 --no_vel_der \
--out_frame_fm_dir $fm2_frame_dir  --fm_yaml $yaml  --fm_load $fm2_ckpt --input_features  --fea_weight $fea_weight \
--frame_num $frame_num --frame_num_fm $frame_fm_num \
--ft_v_path $finetuned_ft_path_vel

python sciml_fm/run_nerf_vort_sciml.py --config configs/scalarflowreal.txt --lrate 0.01 --lrate_den 1e-4 \
--lrate_decay 5000 --N_iters 10000 --i_weights 1000 --i_video 1000 --i_print 100 \
--expname $vort_exp --finest_resolution 256 --base_resolution 16 --finest_resolution_t 128 --base_resolution_t 16 --num_levels 16 --N_samples 192 --N_rand 512 --log2_hashmap_size 19 --rec_weight 10000 \
--rec_weight 10000 --vel_weight 0 --d_weight 0 --flow_weight 0.001 --vel_num_layers 2 \
--ft_path $vort_ft_path --vel_path $vort_vel_path --no_vel_der --vel_scale 0.05 \
--finest_resolution_v 128 --base_resolution_v 16 --finest_resolution_v_t 128 --base_resolution_v_t 16 \
--n_particles 50 --vort_intensity 5 --vort_weight 0.01  \
--out_frame_fm_dir $fm2_frame_dir  --fm_yaml $yaml  --fm_load $fm2_ckpt --input_features  --no_reload    --fea_weight $fea_weight  \
--frame_num $frame_num --frame_num_fm $frame_fm_num

#sim+pred
python sciml_fm/scalarflow_metrics_single.py --hyfluid_dir  ./logs/$vort_exp/run_advect_den_best --gt_dir ./logs/$vort_exp/run_advect_den_best --frame_start 0 --frame_end  $(($frame_num+20))  --out_name sim+pred_$hy2_expname

#sys
python sciml_fm/scalarflow_metrics_single.py --hyfluid_dir  ./logs/$density_exp/testset_best --gt_dir ./logs/$vort_exp/run_advect_den_best --frame_start 0 --frame_end  $(($frame_num+20))  --out_name sys_$hy2_expname
