#!/bin/bash

# SLURM configuration for MN5
#SBATCH --qos=acc_debug
#SBATCH --output=serial_inference%j.out
#SBATCH --error=serial_inference%j.err
#SBATCH --ntasks=160
#SBATCH --time=06:00:00
#SBATCH --job-name=training_inference_pipeline
#SBATCH --mail-type=end
#SBATCH --mail-user=amirpasha.mozaffari@bsc.es
#SBATCH --exclusive
# SBATCH --gres=gpu:4

# SBATCH -A bsc32
# number of the nodes
# SBATCH -N 1
# how many processes per single node
# SBATCH --ntasks-per-node 16
# how many core per process
# SBATCH -c 5
# how many gpus
# SBATCH --gres=gpu:4
# #SBATCH -J testtask
# #SBATCH -o testtask-%j.out
# SBATCH --output=serial_inference%j.out
# SBATCH --error=serial_inference%j.err
# SBATCH --time=120
# SBATCH --exclusive
# SBATCH --qos acc_debug


# 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=09
REGION_END=09
REGIONS=()
for region_num in $(seq -f "%03g" $REGION_START $REGION_END); do
    REGIONS+=( "W${region_num}" )
done


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

    # Search for model files matching the current region (e.g., XGBoost_LAI_2000_2014_W025_YYYYMMDD_HHMMSS.joblib)
    model_files=( "$SAVE_DATA_PATH"/XGBoost_LAI_"${START_YEAR_TRAINING}"_"${END_YEAR_TRAINING}"_"${REGION}"_[0-9]*.joblib )
    if [ ${#model_files[@]} -eq 0 ]; then
        echo "No model files found for region ${REGION} in ${SAVE_DATA_PATH}"
        exit 1
    fi

    latest_file=""
    latest_timestamp_num=0

    # Loop over the files for the specific region
    for file in "${model_files[@]}"; do
        filename=$(basename "$file")
        # Extract the timestamp using the region variable.
        ts=$(echo "$filename" | sed -E "s/.*${REGION}_([0-9]{8}_[0-9]{6}).joblib/\1/")
        ts_num=${ts//_/}  # Remove underscore to compare as a number
        if [[ "$ts_num" -gt "$latest_timestamp_num" ]]; then
            latest_timestamp_num="$ts_num"
            latest_file="$file"
            MODEL_TIMESTAMP="$ts"
        fi
    done

    echo "Latest model file for region ${REGION}: $latest_file"
    echo "Extracted MODEL_TIMESTAMP: $MODEL_TIMESTAMP"

    # 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."








