setup

#!/usr/bin/env bash

# ======================
# Check IPSL-AID installation
# ======================
if ! command -v ipsl-aid &> /dev/null; then
    echo "ERROR: ipsl-aid command not found. Please install the package first using: uv pip install -e ."
    echo "       or activate the virtual environment if already installed using: source .venv/bin/activate"
    exit 1
fi

# ======================
# 1. Unified Debug Flag
# ======================
debug=false     # true = SLURM debug job + Python --debug true,  Note: Avoid True (use true instead)
                # false = full job + Python --debug false, Note: Avoid False (use false instead)

DATE=$(date +"%Y%m%d")
TODAY=$(date +"%Y%m%d_%H%M%S")

# ======================
# 2. Parameters
# ======================

run_type="train" # "train", ""train_regional"

region="Europe" # Only for run_type="train_regional"
region_size=("288" "720") # Only for run_type="train_regional"

save_model=true
save_checkpoint_name="difusion_model"
load_checkpoint_name="difusion_model"
save_per_samples=10000

year_start=2015
year_end=2019
year_start_test=2020 #2020
year_end_test=2020 #2020

batch_size=36 #20 #40 # use maximum 18 for 1 GPU
num_epochs=8
learning_rate=0.0001
num_workers=16

datadir="/leonardo_work/EUHPC_D27_095/kkingston/AI-Downscaling/data/data_FOURxDaily"
per_var_datadir=(
  "VAR_2T=/leonardo_work/EUHPC_D27_095/kkingston/AI-Downscaling/data/data_FOURxDaily"
  "VAR_10U=/leonardo_work/EUHPC_D27_095/kkingston/AI-Downscaling/data/data_FOURxDaily"
  "VAR_10V=/leonardo_work/EUHPC_D27_095/kkingston/AI-Downscaling/data/data_FOURxDaily"
  "VAR_TP=/leonardo_work/EUHPC_D27_095/kkingston/AI-Downscaling/data/data_FOURxDaily_tp"
  "VAR_D2M=/leonardo_work/EUHPC_D27_095/kkingston/AI-Downscaling/data/data_FOURxDaily_d2m_sstk"
  "VAR_ST=/leonardo_work/EUHPC_D27_095/kkingston/AI-Downscaling/data/data_FOURxDaily_st"
  )

time_normalization="cos_sin"

varnames_list=("VAR_2T" "VAR_10U" "VAR_10V" "VAR_TP" "VAR_D2M" "VAR_ST")
constant_varnames_list=("z" "lsm")
constant_varnames_file="ERA5_const_sfc_variables.nc"
normalization_types=("VAR_2T=standard" "VAR_10U=standard" "VAR_10V=standard" "VAR_TP=log1p_standard" "VAR_D2M=standard" "VAR_ST=standard")
units_list=("K" "m/s" "m/s" "m/h" "K" "K")
dynamic_covariates=()
dynamic_covariates_dir="../data_covariates/"

sbatch=12 #12
tbatch=1800
batch_size_lat=145
batch_size_lon=361

epsilon=0.02
beta=1.0
margin=8

pretrained_path=""
model_name=""

dtype="fp32"
arch="adm" # options are "ddpmpp", "ncsnpp", and "adm"
precond="unet" # edm for edm (default model), unet for unet,
# for diffusion in_channels = numbers of variables, for unet in_channels = numbers of variables + constants

if [[ "$precond" == "unet" ]]; then
    in_channels=10     # number of variables + constants (lat + lon + z + lsm)
    cond_channels=""   # no cond_channels for unet
else  # for diffusion-based arch
    in_channels=6      # number of variables only
    cond_channels=10   # number of variables + lat + lon + z + lsm
fi

out_channels=6 # number of variables
inference_type="direct" # "direct" for unet, "sampler" for diffusion

compute_crps=false

if [[ "$compute_crps" == true && "$inference_type" != "sampler" ]]; then
    echo "WARNING: CRPS requires sampler inference. Disabling compute_crps."
    compute_crps=false
fi

# EDM Sampler Parameters
num_steps=10
sigma_min=0.002
sigma_max=80.0
rho=7
s_churn=40
solver="heun"


# Add apply_filter parameter
apply_filter=false

# ======================
# 3. SLURM Configuration
# ======================

SLURM_ACCOUNT="EUHPC_D27_095"
SLURM_PARTITION="boost_usr_prod"
SLURM_QOS="boost_qos_bprod"
SLURM_GRES=4
SLURM_NODES=1
SLURM_NTASKS_PER_NODE=1
SLURM_CPUS_PER_TASK=16

# ======================
# Generate Tag and Output Paths
# ======================
eps_str="${epsilon//./d}"
lr_str="${learning_rate//./d}"

mode_str=$([[ "$debug" == true ]] && echo "debug" || echo "prod")

# Create BASE model name parts (configuration only)
base_parts=("y${year_start}_${year_end}")
base_parts+=("norm_${time_normalization}")
base_parts+=("lat${batch_size_lat}")
base_parts+=("lon${batch_size_lon}")
base_parts+=("vars${#varnames_list[@]}")
base_parts+=("dt${dtype}")
base_parts+=("arch${arch}")
base_parts+=("pre${precond}")
base_parts+=("in${in_channels}")

if [[ "$precond" == "edm" ]]; then
    base_parts+=("cond${cond_channels}")
fi

base_parts+=("out${out_channels}")
filter_str=$([[ "$apply_filter" == true ]] && echo "filter" || echo "nofilter")
base_parts+=("${filter_str}")

# Create TRAINING model name parts (training-specific)
train_parts=("ep${num_epochs}")
train_parts+=("sb${sbatch}")
train_parts+=("tb${tbatch}")
train_parts+=("eps${eps_str}")
train_parts+=("mrg${margin}")
train_parts+=("inf_${inference_type}")

sigma_min_str="${sigma_min//./d}"
sigma_max_str="${sigma_max//./d}"
rho_str="${rho//./d}"

if [[ "$inference_type" == "sampler" ]]; then
    train_parts+=("steps${num_steps}")
    train_parts+=("sigmin${sigma_min_str}")
    train_parts+=("sigmax${sigma_max_str}")
    train_parts+=("rho${rho_str}")
    train_parts+=("s_churn${s_churn}")
    train_parts+=("solver_${solver}")
fi

# Add checkpoint info to model name if not training from scratch
if [[ "$run_type" == "resume_train" || "$run_type" == "inference" ]]; then
    train_parts+=("resume")
fi

# Combine all parts
model_parts=("${mode_str}" "${base_parts[@]}" "${train_parts[@]}")

# Model name (full configuration) - PROPERLY joined with underscores
model_name=$(IFS=_ ; echo "${model_parts[*]}")
tag="${model_name}_bs${batch_size}_lr${lr_str}_${TODAY}"

# Define output paths - hierarchical structure
# Main folder: mode + base configuration - PROPERLY joined
main_folder="${mode_str}_$(IFS=_ ; echo "${base_parts[*]}")"

# Sub folder: training parameters + date - PROPERLY joined
train_parts_str=$(IFS=_ ; echo "${train_parts[*]}")
sub_folder="train_${train_parts_str}_bs${batch_size}_lr${lr_str}"

# Prefix: simple run identifier
prefix="run_${TODAY}"

# ======================
# Generate SBATCH Script Only (removed run script)
# ======================
PROJECT_ROOT="$(cd "$(dirname "$0")" && pwd)"
sbatch_script="$PROJECT_ROOT/slurm/sbatch_diffusion_${main_folder}_${sub_folder}.sh"

# Create slurm_io directory if it doesn't exist
mkdir -p "$PROJECT_ROOT/slurm_io"

# ======================
# Generate SBATCH Script
# ======================

cat > $sbatch_script <<EOF
#!/bin/bash
#SBATCH --job-name=diffusion_${tag}
#SBATCH --account=$SLURM_ACCOUNT
#SBATCH --partition=$SLURM_PARTITION
EOF

# SLURM debug selection
if [ "$debug" = true ]; then
    echo "#SBATCH --qos=boost_qos_dbg" >> $sbatch_script
    echo "#SBATCH --time=00:30:00" >> $sbatch_script
else
    echo "#SBATCH --time=24:00:00" >> $sbatch_script
fi

cat >> $sbatch_script <<EOF
#SBATCH --nodes=$SLURM_NODES
#SBATCH --gpus-per-node=$SLURM_GRES
#SBATCH --mem-per-gpu=64G
#SBATCH --ntasks-per-node=$SLURM_NTASKS_PER_NODE
#SBATCH --cpus-per-task=$SLURM_CPUS_PER_TASK
#SBATCH --error=$PROJECT_ROOT/slurm_io/diffusion_${tag}_%j.err
#SBATCH --output=$PROJECT_ROOT/slurm_io/diffusion_${tag}_%j.out

ulimit -s unlimited
module purge

# Activate virtual environment
source $PROJECT_ROOT/.venv/bin/activate

export PYTHONUNBUFFERED=1

echo "===== Job Infos ====="
echo "Debug mode: $debug"
echo "Apply filter: $apply_filter"
echo "Inference type: $inference_type"
echo "Node list: \${SLURM_NODELIST}"
echo "Job ID: \${SLURM_JOB_ID}"
echo "Current dir: \$(pwd)"
echo "Output main folder: $main_folder"
echo "Output sub folder: $sub_folder"
echo "====================="

# Use ipsl-aid command directly
ipsl-aid \\
    --debug $debug \\
    --main_folder "$main_folder" \\
    --sub_folder "$sub_folder" \\
    --prefix "$prefix" \\
    --run_type "$run_type" \\
    --save_model $save_model \\
    --save_checkpoint_name "$save_checkpoint_name" \\
    --load_checkpoint_name "$load_checkpoint_name" \\
    --save_per_samples "$save_per_samples" \\
    --inference_type "$inference_type" \\
    --arch "$arch" \\
    --precond "$precond" \\
    --in_channels "$in_channels" \\
    --out_channels "$out_channels" \\
    --apply_filter $apply_filter \\
EOF

# Add cond_channels only for EDM
[[ "$precond" == "edm" ]] && echo "    --cond_channels \"$cond_channels\" \\" >> $sbatch_script

# Add optional parameters
[ -n "$pretrained_path" ] && echo "    --pretrained_path \"$pretrained_path\" \\" >> $sbatch_script
[ -n "$model_name" ] && echo "    --model_name \"$model_name\" \\" >> $sbatch_script

# Regional only
if [[ "$run_type" == "train_regional" || "$run_type" == "inference_regional" ]]; then
    [ -n "$region" ] && echo "    --region \"$region\" \\" >> $sbatch_script
    [ ${#region_size[@]} -gt 0 ] && echo "    --region_size ${region_size[@]} \\" >> $sbatch_script
fi

# Add arrays
[ ${#varnames_list[@]} -gt 0 ] && echo "    --varnames_list ${varnames_list[@]} \\" >> $sbatch_script
[ ${#normalization_types[@]} -gt 0 ] && echo "    --normalization_types ${normalization_types[@]} \\" >> $sbatch_script
[ ${#constant_varnames_list[@]} -gt 0 ] && echo "    --constant_varnames_list ${constant_varnames_list[@]} \\" >> $sbatch_script
echo "    --constant_varnames_file \"$constant_varnames_file\" \\" >> $sbatch_script
[ ${#units_list[@]} -gt 0 ] && echo "    --units_list ${units_list[@]} \\" >> $sbatch_script
[ ${#dynamic_covariates[@]} -gt 0 ] && echo "    --dynamic_covariates ${dynamic_covariates[@]} \\" >> $sbatch_script

cat >> $sbatch_script <<EOF
    --year_start "$year_start" \\
    --year_end "$year_end" \\
    --year_start_test "$year_start_test" \\
    --year_end_test "$year_end_test" \\
    --batch_size "$batch_size" \\
    --num_epochs "$num_epochs" \\
    --learning_rate "$learning_rate" \\
    --num_workers "$num_workers" \\
    --datadir "$datadir" \\
    --per_var_datadir ${per_var_datadir[@]} \\
    --time_normalization "$time_normalization" \\
    --tbatch "$tbatch" \\
    --sbatch "$sbatch" \\
    --batch_size_lat "$batch_size_lat" \\
    --batch_size_lon "$batch_size_lon" \\
    --epsilon "$epsilon" \\
    --beta "$beta" \\
    --margin "$margin" \\
    --dynamic_covariates_dir "$dynamic_covariates_dir" \\
    --dtype "$dtype" \\
    --num_steps "$num_steps" \\
    --sigma_min "$sigma_min" \\
    --sigma_max "$sigma_max" \\
    --rho "$rho" \\
    --s_churn "$s_churn" \\
    --solver "$solver" \\
    --compute_crps "$compute_crps"

exit \$?
EOF

chmod +x $sbatch_script

# ======================
# Final Output
# ======================
echo "=== Setup Complete ==="
echo "Debug mode: $debug"
echo "Apply filter: $apply_filter"
echo "Inference type: $inference_type"
echo "Run type: $run_type"
echo "Save model: $save_model"
echo "Load checkpoint: $load_checkpoint_name"
if [[ "$run_type" == "resume_train" || "$run_type" == "inference" ]]; then
    echo "Checkpoint file: $load_checkpoint_name"
fi
echo ""
echo "Model name parts: ${model_parts[*]}"
echo "Model name: $model_name"
echo "Tag: $tag"
echo ""
echo "Main folder: $main_folder"
echo "Sub folder: $sub_folder"
echo "Prefix: $prefix"
echo ""
echo "SBATCH script: $sbatch_script"
echo ""
echo "To submit the job:"
echo "  sbatch $sbatch_script"
echo ""
echo "To test locally (small run):"
echo "  source .venv/bin/activate"
echo "  ipsl-aid --debug true --run_type train --num_epochs 1 --batch_size 2"