#!/bin/bash

# SLURM configuration for MN5
#SBATCH --qos=acc_debug
#SBATCH --output=serial_training_inference_xgboost%j.out
#SBATCH --error=serial_training_inference_xgboost%j.err
#SBATCH --ntasks=80
#SBATCH --time=01:00:00
#SBATCH --job-name=training_inference_pipeline_xgboost
#SBATCH --mail-type=end
#SBATCH --mail-user=enric.millan@bsc.es

# Machine-specific setup
if [ "$BSC_MACHINE" = "mn5" ]; then
    export MODULEPATH=$MODULEPATH:/apps/GPP/EASYBUILD/modules/all
    conda init
    source /apps/GPP/MINICONDA/24.1.2/etc/profile.d/conda.sh
    conda activate /home/bsc/bsc558537/.conda/envs/acc_gpu
    export PYTHONPATH=/home/bsc/bsc558537/.conda/envs/acc_gpu/lib/python3.12/site-packages:$PYTHONPATH
elif [ "$BSC_MACHINE" = "nord3v2" ]; then
    module load GDAL/3.3.2-foss-2019b-Python-3.7.4
    module load pycdo/1.5.3-foss-2019b-Python-3.7.4
    module load pynco/0.0.2-foss-2019b-Python-3.7.4
fi

# Define paths to Data & scripts
SCRIPT_DIR=$(dirname "$(realpath "$PWD")")
TRAINING_SCRIPT="$SCRIPT_DIR/testing/train_XGB_lai_monthly.py"
REGIONS_DATA_PATH="$(dirname "$SCRIPT_DIR")/inputs/regions_LAI_39.json"
INPUT_DATA_PATH="/gpfs/scratch/bsc32/cerise/outputs/bsc032013/0.0089_deg/"
SAVE_DATA_PATH="/gpfs/scratch/bsc32/ai4land/saved_data/"
MODEL_OUTPUTS="/gpfs/scratch/bsc32/ai4land/output_data/"

# TRAINING AND INFERENCE VARIABLES
START_YEAR_TRAINING=2008
END_YEAR_TRAINING=2014
TRAINING_YEARS="2010 2011 2012 2013 2014"
TESTING_YEARS="2008"

# Define the list of regions to process
REGION_START=12
REGION_END=12
REGIONS=()
for region_num in $(seq -f "%03g" $REGION_START $REGION_END); do
    REGIONS+=( "W${region_num}" )
done

# Generate dynamic MODEL_TIMESTAMP
MODEL_TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
echo "Generated MODEL_TIMESTAMP: $MODEL_TIMESTAMP"

# Loop through each region
for REGION in "${REGIONS[@]}"; do
    echo "Processing region: $REGION"

    # TRAIN THE MODEL
    echo "Starting Model Training for region: $REGION..."
    python "$TRAINING_SCRIPT" \
        --model_timestamp $MODEL_TIMESTAMP \
        --start_year $START_YEAR_TRAINING \
        --end_year $END_YEAR_TRAINING \
        --region $REGION \
        --save_data_path $SAVE_DATA_PATH \
        --imputation mean \
        --train_years $TRAINING_YEARS \
        --test_years $TESTING_YEARS \
        --water_filter \
        --autoregressive \
        --batch_size 16384 \
        --ground_truth_scale \
        --ground_truth_log_transform \
        #--post_plots \

    if [ $? -ne 0 ]; then
        echo "Error: Training script failed for region: $REGION."
        exit 1
    fi
    echo "The Model training for region $REGION is finished."
done

echo "The training for all regions is finished."
