#!/bin/bash
#SBATCH --qos=acc_bsces
#SBATCH --output=serial_training_inference_rnn_%j.out
#SBATCH --error=serial_training_inference_rnn_%j.err
#SBATCH --ntasks=80
#SBATCH --time=01:30:00
#SBATCH --job-name=training_inference_pipeline_rnn
#SBATCH --mail-type=end
#SBATCH --mail-user=enric.millan@bsc.es

# -------------------------------
# Arguments from meta-launcher
# -------------------------------
RNN_CELL_TYPE=$1
USE_ATTENTION=$2
DROPOUT=$3
HIDDEN_SIZE=$4
PATIENCE=$5
REGION=$6
TEACHER_FORCING_RATIO=${7:-0.0}
LOG_AND_SCALE=${8:-False}
WEIGHT_DECAY=${9:-0.0}
BATCH_SIZE=${10:-16384}
LOG_CSV=${11:-False}
LOG_CSV_FILE=${12:-"experiment_log.csv"}

# Shift out the first 12 args so $@ now contains only years
shift 12

# Parse training and test years from trailing arguments
TESTING_YEAR="${@: -1}"                   # last argument is test year
TRAINING_YEARS=("${@:1:$#-1}")            # all others are training years

# Compute start and end year across train + test window
ALL_RELEVANT_YEARS=("${TRAINING_YEARS[@]}" "$TESTING_YEAR")
START_YEAR=$(printf "%s\n" "${ALL_RELEVANT_YEARS[@]}" | sort -n | head -n1)
END_YEAR=$(printf "%s\n" "${ALL_RELEVANT_YEARS[@]}" | sort -n | tail -n1)

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

# -------------------------------
# Paths and constants
# -------------------------------
SCRIPT_DIR=$(dirname "$(realpath "$PWD")")
TRAINING_SCRIPT="$SCRIPT_DIR/testing/train_rnn_lai_monthly.py"
SAVE_DATA_PATH="/gpfs/scratch/bsc32/ai4land/saved_data/"
MODEL_TIMESTAMP=$(date +"%Y%m%d_%H%M%S")

# -------------------------------
# Launch model training
# -------------------------------
echo "Launching training for region=$REGION | cell=$RNN_CELL_TYPE | dropout=$DROPOUT | hidden=$HIDDEN_SIZE | patience=$PATIENCE"
echo "Training years: ${TRAINING_YEARS[*]} | Testing year: $TESTING_YEAR"
echo "Start year: $START_YEAR | End year: $END_YEAR"

python "$TRAINING_SCRIPT" \
    --model_timestamp $MODEL_TIMESTAMP \
    --start_year $START_YEAR \
    --end_year $END_YEAR \
    --region $REGION \
    --save_data_path $SAVE_DATA_PATH \
    --imputation mean \
    --train_years ${TRAINING_YEARS[*]} \
    --test_years $TESTING_YEAR \
    --water_filter \
    --autoregressive \
    --hidden_size $HIDDEN_SIZE \
    --batch_size $BATCH_SIZE \
    --num_epochs 20 \
    --early_stopping \
    --early_stopping_patience $PATIENCE \
    --rnn_cell_type $RNN_CELL_TYPE \
    --dropout $DROPOUT \
    --teacher_forcing_ratio $TEACHER_FORCING_RATIO \
    --weight_decay $WEIGHT_DECAY \
    $( [ "$USE_ATTENTION" == "True" ] && echo "--use_attention" ) \
    $( [ "$LOG_AND_SCALE" == "True" ] && echo "--ground_truth_log_transform --ground_truth_scale" ) \
    $( [ "$LOG_CSV" == "True" ] && echo "--log_csv" ) \
    --log_csv_file "$LOG_CSV_FILE"

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

echo "Model training for region $REGION completed."
