#!/bin/bash
# shellcheck disable=SC2086,SC2004,SC2124,SC2206,SC2046,SC2181 # Using bc for float arithmetic is safe.

#===============================================================================
# UNIFIED HYPERPARAMETER SWEEP SCRIPT (v5.1)
#===============================================================================
#
# This script launches a comprehensive hyperparameter grid sweep for training
# models. It supports both SLURM-based parallel job submission and
# sequential local execution.
#
# Key Features:
#   - Complete Configuration: Integrates all detailed model, dataset, µP,
#     callback, and training parameters from the experimental scripts.
#   - Detailed Run Summaries: Prints a formatted configuration for each run.
#   - 0-Indexed Runs: Job indices start from 0 for programming consistency.
#   - Controllable Sweeps: Use boolean flags to easily enable/disable sweeps.
#   - Sweep Hashing & Indexing: Generates a unique ID for the sweep and a
#     sequential index for each run, embedding them in the job name for clarity.
#   - Range-Based Execution: Allows splitting a large sweep across multiple
#     machines by specifying a range of job indices to run (e.g., 0-99).
#
#===============================================================================

# --- 1. SWEEP CONTROL FLAGS ---
SWEEP_SEEDS=false
SWEEP_BSS=false
SWEEP_LRS=true
SWEEP_NEW_VS=true
SWEEP_NEW_BETAS=true

# --- 2. EXECUTION AND SLURM CONFIGURATION ---
USE_SLURM=true
MAX_CONCURRENT_JOBS=4
TARGET_NODE="anon"
STATIC_PARAMETERS=""
STATIC_PARAMETERS+=" llm_config.model.attn_config.attn_impl=flash"
STATIC_PARAMETERS+=" llm_config.precision=amp_bf16"
DRY_RUN=false

# Path to the project directory containing the training scripts.
PROJECT_PATH="$HOME/projects/mt_dao"
# Path for dataset caching.
DATASET_CACHE_DIR="/nfs-share/datasets/mt_dao/dataset_cache"
# Path to a pre-trained model for resuming, or 'null' for random init.
PRE_TRAINED_MODEL_PATH=null

# Default run range (0 to infinity)
START_INDEX=0
END_INDEX=-1

# --- 3. HYPERPARAMETER GRID DEFINITION ---
declare -A LLM_CONFIGS
LLM_CONFIGS=(
	["qhadopt_cent"]="smollm-3m"
)

# Base values
BASE_SEED=42
BASE_BS=64
BASE_LR="0.00390625"
BASE_NEW_V="0.9"
BASE_NEW_BETA="0.999"

# Sweep arrays
SEEDS=(42)
BSS=(64)
LRS=(
	"$(echo "2^-11" | bc -l)"
	"$(echo "2^-10" | bc -l)"
	"$(echo "2^-9" | bc -l)"
	"0.002762135864"
	"$(echo "2^-8" | bc -l)"
	"0.005524271728"
	"$(echo "2^-7" | bc -l)"
	"0.01104854346"
	"$(echo "2^-6" | bc -l)"
)

NEW_VS=(
	# "0.85"
	# "0.9"
	# "0.95"
	"0.96"
	"0.97"
	"0.98"
	# "0.99"
	# "0.995"
	# "0.999"
	# "1.0"
)

NEW_BETAS=(
	# "0.95"
	# "0.99"
	# "0.995"
	"0.999"
	# "0.9995"
	# "0.9999"
)

# --- 4. FIXED & SCALING PARAMETERS ---
D_MODEL=256
BASE_MAX_DURATION=5120
BASE_T_WARMUP=512
BASE_T_COOLDOWN=512
FIXED_INIT_STD="0.02"
FIXED_WEIGHT_DECAY="0.0"

# --- 5. STATIC PARAMETERS FOR ALL RUNS ---
# This block contains the full, comprehensive set of static parameters.

# General & Precision
STATIC_PARAMETERS+=" use_wandb=true"
STATIC_PARAMETERS+=" is_test=false"
# Core Model
STATIC_PARAMETERS+=" llm_config.model.d_model=${D_MODEL}"
STATIC_PARAMETERS+=" llm_config.model.n_heads=4"
STATIC_PARAMETERS+=" llm_config.model.n_layers=4"
STATIC_PARAMETERS+=" llm_config.model.max_seq_len=2048"
STATIC_PARAMETERS+=" ++llm_config.model.no_bias=True"
STATIC_PARAMETERS+=" ++llm_config.model.vocab_size=50368"
STATIC_PARAMETERS+=" +llm_config.model.use_cache=true"
# Attention & RoPE
STATIC_PARAMETERS+=" llm_config.model.attn_config.rope=true"
STATIC_PARAMETERS+=" +llm_config.model.attn_config.rope_impl=dail"
STATIC_PARAMETERS+=" llm_config.model.attn_config.rope_theta=10000"
# FFN Base Config
STATIC_PARAMETERS+=" llm_config.model.ffn_config.ffn_act_fn.name=silu"
STATIC_PARAMETERS+=" llm_config.model.ffn_hidden_size=1024" # 4 * d_model
# Init, muP, and complete-P
STATIC_PARAMETERS+=" llm_config.model.init_config.name=baseline_"
STATIC_PARAMETERS+=" llm_config.model.init_config.init_std=${FIXED_INIT_STD}"
STATIC_PARAMETERS+=" llm_config.model.init_config.emb_init_std=${FIXED_INIT_STD}"
STATIC_PARAMETERS+=" ++llm_config.model.use_peri_norm=True"
STATIC_PARAMETERS+=" ++llm_config.model.use_embedding_norm=True"
STATIC_PARAMETERS+=" llm_config.model.mup_config.mup_enabled=true"
STATIC_PARAMETERS+=" llm_config.model.mup_config.mup_disable_attention_scaling=true"
STATIC_PARAMETERS+=" llm_config.model.mup_config.mup_disable_hidden_lr_scaling=false"
STATIC_PARAMETERS+=" llm_config.model.mup_config.mup_width_multiplier=1.0"
STATIC_PARAMETERS+=" llm_config.model.mup_config.mup_input_alpha=1.0"
STATIC_PARAMETERS+=" llm_config.model.mup_config.mup_output_alpha=1.0"
STATIC_PARAMETERS+=" llm_config.model.mup_config.completep_depth_alpha_enabled=true"
STATIC_PARAMETERS+=" llm_config.model.mup_config.completep_depth_multiplier=1.0"
STATIC_PARAMETERS+=" llm_config.model.mup_config.completep_depth_alpha_exp=1.0"
STATIC_PARAMETERS+=" llm_config.model.mup_config.completep_eps_scaling_enabled=true"
STATIC_PARAMETERS+=" llm_config.model.mup_config.mup_scale_expert_sel=true"
# Optimizer & Scheduler
STATIC_PARAMETERS+=" llm_config.optimizer.name=qhadopt"
STATIC_PARAMETERS+=" llm_config.optimizer.weight_decay=${FIXED_WEIGHT_DECAY}"
STATIC_PARAMETERS+=" llm_config.optimizer.betas=[0.95,0.9999]"
STATIC_PARAMETERS+=" llm_config.optimizer.eps=1e-6"
STATIC_PARAMETERS+=" ++llm_config.optimizer.v1=1.0"
STATIC_PARAMETERS+=" llm_config.scheduler.schedulers.lr.name=constant_with_sqrt_cooldown_with_warmup_and_switch"
STATIC_PARAMETERS+=" ++llm_config.callbacks.quasi_hyper_switch_callback={}"
# Dataset & Tokenizer
STATIC_PARAMETERS+=" dataset=smollm-corpus-shared-with-val"
STATIC_PARAMETERS+=" dataset.train.root_local=${DATASET_CACHE_DIR}/smollm-corpus-shared"
STATIC_PARAMETERS+=" dataset.val.root_local=${DATASET_CACHE_DIR}/smollm-corpus-shared-val"
STATIC_PARAMETERS+=" dataset/streams@dataset.train.streams=smollm_corpus_4_clients_proportional"
STATIC_PARAMETERS+=" dataset/streams@dataset.val.streams=smollm_corpus_4_clients_proportional"
STATIC_PARAMETERS+=" centralized.stream_id=null"
STATIC_PARAMETERS+=" llm_config.tokenizer.name='HuggingFaceTB/SmolLM-1.7B'"
# Training & Device
STATIC_PARAMETERS+=" llm_config.device_train_microbatch_size=auto"
STATIC_PARAMETERS+=" +llm_config.device_eval_microbatch_size=auto"
STATIC_PARAMETERS+=" ++llm_config.compile_config={}"
# Evaluation
STATIC_PARAMETERS+=" eval_gauntlet_config=empty"
STATIC_PARAMETERS+=" eval_gauntlet_config.destination_dir=${DATASET_CACHE_DIR}/eval"
STATIC_PARAMETERS+=" icl_tasks_config=empty"
STATIC_PARAMETERS+=" icl_tasks_config.root_dir=${DATASET_CACHE_DIR}"
STATIC_PARAMETERS+=" llm_config.eval_subset_num_batches=128"
# Callbacks & Logging
STATIC_PARAMETERS+=" ~llm_config.loggers.tensorboard"
STATIC_PARAMETERS+=" llm_config.console_log_interval=64ba"
STATIC_PARAMETERS+=" ++llm_config.callbacks.activation_monitor_full_model={}"
STATIC_PARAMETERS+=" llm_config.callbacks.activation_monitor_full_model.interval=16ba"
STATIC_PARAMETERS+=" ++llm_config.callbacks.optimizer_monitor={}"
STATIC_PARAMETERS+=" llm_config.callbacks.optimizer_monitor.interval=16ba"

STATIC_PARAMETERS+=" ++llm_config.callbacks.betas_monitor={}"
STATIC_PARAMETERS+=" ++llm_config.callbacks.quasi_hyperbolic_parameter_monitor={}"

# --- 6. SCRIPT LOGIC ---

print_run_summary() {
	local job_name=$1
	local run_index=$2
	local total_jobs=$3
	local seed=$4
	local bs=$5
	local new_lr=$6
	local new_v=$7
	local new_beta=$8
	local max_dur=${9}
	local warmup=${10}
	local cooldown=${11}

	echo "----------------------------------------------------"
	echo "Preparing Job: ${job_name}"
	echo "Sweep Index:   ${run_index} / $((total_jobs - 1))"
	echo "===================================================="
	echo "  >> RUN CONFIGURATION << "
	echo "----------------------------------------------------"
	echo "  Core Hyperparameters:"
	echo "    Seed:               ${seed}"
	echo "    Batch Size:         ${bs}"
	echo "    Base LR:            ${BASE_LR}"
	echo "    New LR (swept):     ${new_lr}"
	echo "    New v (swept):      ${new_v}"
	echo "    New beta_2 (swept): ${new_beta}"
	echo "  Duration (scaled):"
	echo "    Max Duration:       ${max_dur}ba"
	echo "    Warmup:             ${warmup}ba"
	echo "    Cooldown:           ${cooldown}ba"
	echo "  Key Fixed Parameters:"
	echo "    d_model:            ${D_MODEL}"
	echo "    Attention:          torch"
	echo "    µP Enabled:         true"
	echo "    Save/Eval Interval: ${max_dur}ba"
	echo "----------------------------------------------------"
}

# Argument Parsing
while (("$#")); do
	case "$1" in
	--range)
		if [[ $2 =~ ^[0-9]+-[0-9]+$ ]]; then
			START_INDEX=$(echo "$2" | cut -d'-' -f1)
			END_INDEX=$(echo "$2" | cut -d'-' -f2)
		else
			echo "Error: --range requires format START-END (e.g., 0-99)" >&2
			exit 1
		fi
		shift 2
		;;
	--no-slurm)
		USE_SLURM=false
		shift
		;;
	--dry-run)
		DRY_RUN=true
		shift
		;;
	-h | --help)
		echo "Usage: $0 [OPTIONS]"
		echo "Options:"
		echo "  --range START-END   Execute jobs in this 0-indexed, inclusive range (e.g., 0-49)."
		echo "  --no-slurm          Run jobs sequentially on the local machine."
		echo "  --dry-run           Print job configurations without executing."
		echo "  -h, --help          Show this help message."
		exit 0
		;;
	*)
		echo "Error: Unsupported flag $1" >&2
		exit 1
		;;
	esac
done

# Conditionally define arrays to iterate over
if [[ $SWEEP_SEEDS == true ]]; then SEEDS_TO_SWEEP=("${SEEDS[@]}"); else SEEDS_TO_SWEEP=(${BASE_SEED}); fi
if [[ $SWEEP_BSS == true ]]; then BSS_TO_SWEEP=("${BSS[@]}"); else BSS_TO_SWEEP=(${BASE_BS}); fi
if [[ $SWEEP_LRS == true ]]; then LRS_TO_SWEEP=("${LRS[@]}"); else LRS_TO_SWEEP=(${BASE_LR}); fi
if [[ $SWEEP_NEW_VS == true ]]; then NEW_VS_TO_SWEEP=("${NEW_VS[@]}"); else NEW_VS_TO_SWEEP=(${BASE_NEW_V}); fi
if [[ $SWEEP_NEW_BETAS == true ]]; then NEW_BETAS_TO_SWEEP=("${NEW_BETAS[@]}"); else NEW_BETAS_TO_SWEEP=(${BASE_NEW_BETA}); fi

# --- Sweep Identification and Summary ---
TOTAL_JOBS=$((${#LLM_CONFIGS[@]} * ${#SEEDS_TO_SWEEP[@]} * ${#BSS_TO_SWEEP[@]} * ${#LRS_TO_SWEEP[@]} * ${#NEW_VS_TO_SWEEP[@]} * ${#NEW_BETAS_TO_SWEEP[@]}))
SWEEP_CONFIG_STRING="SEEDS:${SWEEP_SEEDS}-(${SEEDS[*]});;BSS:${SWEEP_BSS}-(${BSS[*]});LRS:${SWEEP_LRS}-(${LRS[*]});NEW_VS:${SWEEP_NEW_VS}-(${NEW_VS[*]});NEW_BETAS:${SWEEP_NEW_BETAS}-(${NEW_BETAS[*]})"
SWEEP_HASH=$(echo -n "$SWEEP_CONFIG_STRING" | sha1sum | cut -c1-8)

echo "🚀 Starting hyperparameter sweep..."
echo "Sweep Configuration ID (Hash): ${SWEEP_HASH}"
if ((TOTAL_JOBS > 0)); then
	echo "Total potential jobs in this sweep: ${TOTAL_JOBS} (indices 0 to $((TOTAL_JOBS - 1)))"
else
	echo "Total potential jobs in this sweep: 0"
fi
if ((END_INDEX != -1)); then
	echo "Executing jobs from index ${START_INDEX} to ${END_INDEX}"
else
	echo "Executing jobs from index ${START_INDEX} to end"
fi
if [[ $USE_SLURM == false ]]; then echo "Running in local mode (--no-slurm)."; fi
if [[ $DRY_RUN == true ]]; then echo "Running in dry-run mode."; fi
mkdir -p "$DATASET_CACHE_DIR"

# --- Main Submission Loop ---
RUN_INDEX=-1
SUBMITTED_JOB_COUNT=0
declare -a CHAIN_LAST_JOB_IDS
for ((i = 0; i < MAX_CONCURRENT_JOBS; i++)); do CHAIN_LAST_JOB_IDS[i]=""; done

for short_name in "${!LLM_CONFIGS[@]}"; do
	llm_config=${LLM_CONFIGS[$short_name]}
	for seed in "${SEEDS_TO_SWEEP[@]}"; do
		for bs in "${BSS_TO_SWEEP[@]}"; do
			max_duration=$(printf "%.0f\n" $(echo "scale=4; ${BASE_MAX_DURATION} * ${BASE_BS} / ${bs}" | bc))
			t_warmup=$(printf "%.0f\n" $(echo "scale=4; ${BASE_T_WARMUP} * ${BASE_BS} / ${bs}" | bc))
			t_cooldown=$(printf "%.0f\n" $(echo "scale=4; ${BASE_T_COOLDOWN} * ${BASE_BS} / ${bs}" | bc))
			for new_beta in "${NEW_BETAS_TO_SWEEP[@]}"; do
				for new_v in "${NEW_VS_TO_SWEEP[@]}"; do
					for lr in "${LRS_TO_SWEEP[@]}"; do
						RUN_INDEX=$((RUN_INDEX + 1))
						if ((RUN_INDEX < START_INDEX || (END_INDEX != -1 && RUN_INDEX > END_INDEX))); then continue; fi

						SUBMITTED_JOB_COUNT=$((SUBMITTED_JOB_COUNT + 1))
						JOB_NAME=$(printf "%s_%s_%d" "${short_name}" "${SWEEP_HASH}" "${RUN_INDEX}")
						DATETIME=$(date '+%Y%m%d_%H%M%S')
						RUN_UUID="${JOB_NAME}-${DATETIME}"
						SAVE_PATH="${PROJECT_PATH}/outputs/${RUN_UUID}"

						# --- Construct the full parameter string for this run ---
						run_params="${STATIC_PARAMETERS}"
						# Append dynamic parameters
						run_params+=" llm_config=${llm_config}"
						run_params+=" run_uuid=${RUN_UUID}"
						run_params+=" pretrained_model_path=${PRE_TRAINED_MODEL_PATH}"
						# Append sweep parameters
						run_params+=" llm_config.global_seed=${seed}"
						run_params+=" llm_config.global_train_batch_size=${bs}"
						run_params+=" llm_config.device_eval_batch_size=512"
						run_params+=" llm_config.optimizer.lr=${BASE_LR}"
						# Append scaled duration parameters
						run_params+=" llm_config.max_duration=${max_duration}ba"
						run_params+=" llm_config.scheduler.schedulers.lr.t_warmup=${t_warmup}ba"
						run_params+=" llm_config.scheduler.schedulers.lr.t_max=${max_duration}ba"
						run_params+=" llm_config.scheduler.schedulers.lr.t_cooldown=${t_cooldown}ba"
						# set the switch_scale of the scheduler to the ratio of the new_lr and base_lr
						run_params+=" ++llm_config.scheduler.schedulers.lr.switch_scale=$(echo "scale=4; ${lr} / ${BASE_LR}" | bc)"

						# Append quasi-hyperbolic sweep parameters
						run_params+=" ++llm_config.callbacks.quasi_hyper_switch_callback.t_switch=${t_warmup}ba"

						run_params+=" ++llm_config.callbacks.quasi_hyper_switch_callback.new_vs=[${new_v}]"
						run_params+=" ++llm_config.callbacks.quasi_hyper_switch_callback.new_betas=[${new_beta},0.9999]"
						# Append save/eval intervals (set to run end)
						run_params+=" llm_config.save_folder=${SAVE_PATH}"
						run_params+=" llm_config.save_interval=${max_duration}ba"
						run_params+=" llm_config.eval_interval=${max_duration}ba"

						print_run_summary "${JOB_NAME}" "${RUN_INDEX}" "${TOTAL_JOBS}" "${seed}" "${bs}" "${lr}" "${new_v}" "${new_beta}" "${max_duration}" "${t_warmup}" "${t_cooldown}"

						if [ "$DRY_RUN" = true ]; then
							echo "  [DRY RUN] Would execute with params: ${run_params}"
							echo "----------------------------------------------------"
							continue
						fi

						if [[ $USE_SLURM == true ]]; then
							chain_index=$(((SUBMITTED_JOB_COUNT - 1) % MAX_CONCURRENT_JOBS))
							last_job_in_chain=${CHAIN_LAST_JOB_IDS[$chain_index]}
							SBATCH_OPTS="--parsable"
							if [[ -n $last_job_in_chain ]]; then
								SBATCH_OPTS+=" --dependency=afterok:${last_job_in_chain}"
								echo "Submitting to SLURM (chain ${chain_index}, depends on Job ID: ${last_job_in_chain})"
							else
								echo "Submitting to SLURM (starting chain ${chain_index})"
							fi

							job_id=$(
								sbatch ${SBATCH_OPTS} <<EOF
#!/bin/bash
#SBATCH -c 8
#SBATCH -w ${TARGET_NODE}
#SBATCH --gres=gpu:1
#SBATCH --job-name=${JOB_NAME}
#SBATCH --output=%j-${JOB_NAME}.out
#SBATCH --error=%j-${JOB_NAME}.err
#SBATCH --tasks-per-node=1
echo "==================================================================="
echo "STARTING SLURM JOB: ${JOB_NAME} (ID: \$SLURM_JOB_ID)"
echo "Sweep Index: ${RUN_INDEX} / $((TOTAL_JOBS - 1))"
echo "Node: \$(hostname) at \$(date)"
echo "==================================================================="
export EXTERNAL_CONFIGS="${run_params}"
export RUN_UUID="${RUN_UUID}"
bash "${PROJECT_PATH}/scripts/base_centralised_training.sh"
echo "JOB FINISHED: \$(date)"
EOF
							)
							if [[ $? -eq 0 && -n $job_id ]]; then
								CHAIN_LAST_JOB_IDS[$chain_index]=$job_id
								echo "  Successfully submitted. Job ID: ${job_id}"
							else
								echo "❌ Error submitting job ${JOB_NAME}. The dependency chain is now broken. Aborting."
								exit 1
							fi
							sleep 30
						else # Local Execution
							echo "Executing locally..."
							export EXTERNAL_CONFIGS="${run_params}"
							export RUN_UUID="${RUN_UUID}"
							TRAIN_SCRIPT_PATH="${PROJECT_PATH}/scripts/base_centralised_training.sh"
							if [ ! -f "$TRAIN_SCRIPT_PATH" ]; then
								echo "❌ Training script not found at ${TRAIN_SCRIPT_PATH}. Aborting."
								exit 1
							fi
							bash "$TRAIN_SCRIPT_PATH"
							if [[ $? -ne 0 ]]; then
								echo "❌ Local execution of job ${JOB_NAME} failed. Aborting."
								exit 1
							fi
							echo "  Local execution finished."
						fi
					done
				done
			done
		done
	done
done

echo "===================================================="
if ((SUBMITTED_JOB_COUNT == 0)); then
	echo "No jobs were submitted. Please check your run range and sweep flags. 🤔"
else
	if [[ $DRY_RUN == true ]]; then
		echo "Dry run complete. Displayed configurations for ${SUBMITTED_JOB_COUNT} jobs. ✅"
	elif [[ $USE_SLURM == true ]]; then
		echo "All ${SUBMITTED_JOB_COUNT} jobs have been submitted to SLURM. ✅"
	else
		echo "All ${SUBMITTED_JOB_COUNT} jobs have been executed sequentially. ✅"
	fi
fi
echo "===================================================="
