#!/bin/bash -l

#SBATCH -o logs/aggregate-%A_%a.out
#SBATCH -e logs/aggregate-%A_%a.err
#SBATCH -J aggregate_results
#SBATCH --mail-type=END,FAIL
#SBATCH --mail-user=user_email
#SBATCH -N1
#SBATCH --cpus-per-task=1
#SBATCH --mem=4000
#SBATCH --time=00:05:00
#SBATCH --array=1-64

module purge
module load R/4.4

# Load shared variables
source config_analysis.sh

# Compute total jobs
num_locations=${#locations[@]}
num_sigma=${#sigma_betas_vals[@]}
num_sigma_weather=${#True_e_vals[@]}
num_rho_K=${#rho_Ks[@]}
num_School_Terms=${#School_Terms[@]}

total_jobs=$((num_locations * num_sigma * num_sigma_weather * num_rho_K * num_School_Terms))

# Task index (0-based)
task_index=$((SLURM_ARRAY_TASK_ID - 1))

# Compute strides
stride_location=$((num_sigma * num_sigma_weather * num_rho_K * num_School_Terms))
stride_sigma=$((num_sigma_weather * num_rho_K * num_School_Terms))
stride_sigma_weather=$((num_rho_K * num_School_Terms))
stride_rho_K=$((num_School_Terms))
stride_School_Terms=1

# Compute indices
location_index=$(( (task_index / stride_location) % num_locations ))
sigma_index=$(( (task_index / stride_sigma) % num_sigma ))
sigma_weather_index=$(( (task_index / stride_sigma_weather) % num_sigma_weather ))
rho_K_index=$(( (task_index / stride_rho_K) % num_rho_K ))
School_Term_index=$(( (task_index / stride_School_Terms) % num_School_Terms))

# Extract parameter values
location=${locations[$location_index]}
sigma_betas_val=${sigma_betas_vals[$sigma_index]}
True_e_Te=${True_e_vals[$sigma_weather_index]}
True_e_RH=${True_e_vals[$sigma_weather_index]}
rho_K=${rho_Ks[$rho_K_index]}
School_Term=${School_Terms[$School_Term_index]}

echo "Aggregating for location: $location, sigma: $sigma_betas_val, climate: $True_e_Te, rho_K: $rho_K, school term: $School_Term"

# Export to R
export location sigma_betas_val True_e_Te True_e_RH rho_K School_Term

# Run aggregation script
R --no-save --no-restore < ./m-aggregate_results.R > logs/aggregate_${SLURM_ARRAY_TASK_ID}.Rout

echo "Aggregation completed for task ${SLURM_ARRAY_TASK_ID}"