#!/bin/bash -l

#SBATCH -o temp/job-%A_%a.out
#SBATCH -e temp/job-%A_%a.err
#SBATCH -J tm
#SBATCH --mail-type=BEGIN,END
#SBATCH --mail-user=user_email
#SBATCH -N1
#SBATCH --cpus-per-task=4
#SBATCH --mem=5000
#SBATCH --time=00:05:00
#SBATCH --array=1-800  # This will be overwritten by master_HPC.sh

module purge
module load R/4.4

# Load shared variables
source config.sh

# Compute total jobs
num_locations=${#locations[@]}
num_sigma=${#sigma_betas_vals[@]}
num_IDs=${#ID_vals[@]}
num_sigma_weather=${#True_e_vals[@]}
num_rho_Mean=${#rho_Means[@]}
num_rho_K=${#rho_Ks[@]}
num_imports=${#import_rates[@]}
num_School_Terms=${#School_Terms[@]}

total_jobs=$((num_locations * num_sigma * num_IDs * num_sigma_weather * num_rho_Mean * num_rho_K * num_imports * num_School_Terms))


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

# Compute strides
stride_location=$((num_sigma * num_IDs * num_sigma_weather * num_rho_Mean * num_rho_K * num_imports * num_School_Terms))
stride_sigma=$((num_IDs * num_sigma_weather * num_rho_Mean * num_rho_K * num_imports * num_School_Terms))
stride_ID=$((num_sigma_weather * num_rho_Mean * num_rho_K * num_imports * num_School_Terms))
stride_sigma_weather=$((num_rho_Mean * num_rho_K * num_imports * num_School_Terms))
stride_rho_Mean=$((num_rho_K * num_imports * num_School_Terms))
stride_rho_K=$((num_imports * num_School_Terms))
stride_import_num=$((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 ))
ID_index=$(( (task_index / stride_ID) % num_IDs ))
sigma_weather_index=$(( (task_index / stride_sigma_weather) % num_sigma_weather ))
rho_Mean_index=$(( (task_index / stride_rho_Mean) % num_rho_Mean ))
rho_K_index=$(( (task_index / stride_rho_K) % num_rho_K ))
import_index=$(( (task_index / stride_import_num) % num_imports ))
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]}
ID=${ID_vals[$ID_index]}
True_e_Te=${True_e_vals[$sigma_weather_index]}
True_e_RH=${True_e_vals[$sigma_weather_index]}
rho_Mean=${rho_Means[$rho_Mean_index]}
rho_K=${rho_Ks[$rho_K_index]}
import_val=${import_rates[$import_index]}
School_Term=${School_Terms[$School_Term_index]}

echo "Running job for location: $location, sigma: $sigma_betas_val, ID: $ID, Te: $True_e_Te, RH: $True_e_RH, rho_Mean: $rho_Mean, rho_K: $rho_K, import of: $import_val and school term $School_Term"

# Export to R
export location sigma_betas_val ID True_e_Te True_e_RH rho_Mean rho_K import_val School_Term


# Port for parallel R
R_PARALLEL_PORT=$(python3 -c "import random; print(random.randrange(11000,19999))")
export R_PARALLEL_PORT

# Run script
R --no-save --no-restore < ./m-Main.r > Rout/tmp${SLURM_ARRAY_TASK_ID}.Rout