#!/bin/bash

# SLURM configuration for MN5
#SBATCH --qos=acc_debug
#SBATCH --output=serial_training_inference%j.out
#SBATCH --error=serial_training_inference%j.err
#SBATCH --ntasks=640
#SBATCH --time=02:00:00
#SBATCH --job-name=training_inference_pipeline
#SBATCH --mail-type=end
#SBATCH --mail-user=amirpasha.mozaffari@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/bsc032013/.conda/envs/acc_gpu
    export PYTHONPATH=/home/bsc/bsc032013/.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/source/training/train_xgboost_lai_monthly.py"
INFERENCE_SCRIPT="$SCRIPT_DIR/source/inference/inference_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/cerise/outputs/bsc032013/ml/lai/window_model/global_54/"
#MODEL_OUTPUTS="/gpfs/scratch/bsc32/cerise/outputs/bsc032013/ml/lai/window_model/global_54_output/"
#SAVE_DATA_PATH="/gpfs/scratch/bsc32/cerise/outputs/bsc032013/ml/lai/window_model/global_54_extended/"
#MODEL_OUTPUTS="/gpfs/scratch/bsc32/cerise/outputs/bsc032013/ml/lai/window_model/global_54_output_extended/"
SAVE_DATA_PATH="/gpfs/scratch/bsc32/cerise/outputs/bsc032013/ml/lai/window_model/monthly_prediction/global_39"
MODEL_OUTPUTS="/gpfs/scratch/bsc32/cerise/outputs/bsc032013/ml/lai/window_model/monthly_prediction/global_39_output/"

# TRAINING AND INFERENCE VARIABLES
START_YEAR_TRAINING=2000
END_YEAR_TRAINING=2014
START_YEAR_INFERENCE=1939
END_YEAR_INFERENCE=1999
TRAINING_YEARS="2001 2002 2003 2004 2005 2006 2007 2008 2009 2010 2011 2012 2013"
TESTING_YEARS="2000"

#TRAINING_YEARS="2013"
#TESTING_YEARS="2012"


# Define the list of regions to process
REGION_START=39
REGION_END=39
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 \
        --model_type xgboost \
        --imputation mean \
        --train_years $TRAINING_YEARS \
        --test_years $TESTING_YEARS \
        --task_type regression \
        --feature_normalise \
        --ground_truth_normalise \
        --ground_truth_log_transform \
        --autoregressive
        #--autoregressive_smoothing 

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

    # # RUN INFERENCE USING THE TRAINED MODEL
    # echo "Starting Inference for region: $REGION..."
    # python "$INFERENCE_SCRIPT" \
    #     --start_year $START_YEAR_INFERENCE \
    #     --end_year $END_YEAR_INFERENCE \
    #     --start_year_training $START_YEAR_TRAINING \
    #     --end_year_training $END_YEAR_TRAINING \
    #     --region $REGION \
    #     --save_data_path $SAVE_DATA_PATH \
    #     --output_data_path $MODEL_OUTPUTS \
    #     --model_timestamp $MODEL_TIMESTAMP \
    #     --task_type regression \

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

echo "The training and inference for all regions are finished."