#!/bin/bash
# =============================================================================
# BOMLIP-CSP: Relaxation
# =============================================================================
# Usage: ./relaxation.sh [config_file]
#   config_file  Path to the parameter configuration file (default: relaxation_config.sh)
# =============================================================================

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"

# Load configuration
CONFIG_FILE="${1:-${SCRIPT_DIR}/relaxation_config.sh}"
if [ ! -f "${CONFIG_FILE}" ]; then
    echo "Error: Configuration file not found: ${CONFIG_FILE}" >&2
    exit 1
fi
source "${CONFIG_FILE}"

# Generate timestamped output directory
if [ -z "${OUTPUT_PATH}" ]; then
    TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
    OUTPUT_PATH="${PROJECT_DIR}/csp_results/run_${TIMESTAMP}"
fi
mkdir -p "${OUTPUT_PATH}"

echo "Output directory: ${OUTPUT_PATH}"

# Save execution parameters
cat > "${OUTPUT_PATH}/run_params.txt" <<EOF
# BOMLIP-CSP Relaxation Run Parameters
# Generated: $(date +"%Y-%m-%d %H:%M:%S")
# Config:    ${CONFIG_FILE}

[Input]
TARGET_FOLDER       = ${TARGET_FOLDER}
MOLECULE_SINGLE     = ${MOLECULE_SINGLE}

[Compute]
N_GPUS              = ${N_GPUS}
GPU_OFFSET          = ${GPU_OFFSET}
NUM_WORKERS         = ${NUM_WORKERS}
BATCH_SIZE          = ${BATCH_SIZE}
NUM_THREADS         = ${NUM_THREADS}
BIND_CORES          = ${BIND_CORES:-<unset>}

[Optimization]
MAX_STEPS           = ${MAX_STEPS}
FMAX                = ${FMAX}
SCALAR_PRESSURE     = ${SCALAR_PRESSURE}
SKIP_SECOND_STAGE   = ${SKIP_SECOND_STAGE}

[Optimizer]
OPTIMIZER1          = ${OPTIMIZER1}
OPTIMIZER2          = ${OPTIMIZER2}
FILTER1             = ${FILTER1:-<unset>}
FILTER2             = ${FILTER2:-<unset>}

[Model]
MODEL               = ${MODEL}
CUEQ                = ${CUEQ}
USE_ORDERED_FILES   = ${USE_ORDERED_FILES}
COMPILE_MODE        = ${COMPILE_MODE:-<unset>}

[Profiling]
PROFILE             = ${PROFILE}

[Output]
OUTPUT_PATH         = ${OUTPUT_PATH}
EOF

echo "Parameters saved to ${OUTPUT_PATH}/run_params.txt"

# Build optional arguments
FILTER1_ARG=""
FILTER2_ARG=""
[ "${FILTER1}" != "none" ] && [ -n "${FILTER1}" ] && FILTER1_ARG="--filter1 ${FILTER1}"
[ "${FILTER2}" != "none" ] && [ -n "${FILTER2}" ] && FILTER2_ARG="--filter2 ${FILTER2}"

COMPILE_ARG=""
[ -n "${COMPILE_MODE}" ] && COMPILE_ARG="--compile_mode ${COMPILE_MODE}"

BIND_CORES_ARG=""
[ -n "${BIND_CORES}" ] && BIND_CORES_ARG="--bind_cores ${BIND_CORES}"

PROFILE_ARG="--profile ${PROFILE}"

ORDERED_FILES_ARG=""
[ "${USE_ORDERED_FILES}" = true ] && ORDERED_FILES_ARG="--use_ordered_files true"

CUEQ_ARG=""
[ "${CUEQ}" = true ] && CUEQ_ARG="--cueq true"

python "${PROJECT_DIR}/mace-bench/scripts/mace_opt_batch.py" \
    --target_folder "${TARGET_FOLDER}" \
    --molecule_single ${MOLECULE_SINGLE} \
    --gpu_offset ${GPU_OFFSET} \
    --n_gpus ${N_GPUS} \
    --num_workers ${NUM_WORKERS} \
    --batch_size ${BATCH_SIZE} \
    --max_steps ${MAX_STEPS} \
    --fmax ${FMAX} \
    --optimizer1 "${OPTIMIZER1}" \
    --optimizer2 "${OPTIMIZER2}" \
    --num_threads ${NUM_THREADS} \
    --model ${MODEL} \
    --output_path "${OUTPUT_PATH}" \
    ${FILTER1_ARG} ${FILTER2_ARG} ${COMPILE_ARG} ${BIND_CORES_ARG} ${PROFILE_ARG} ${ORDERED_FILES_ARG} ${CUEQ_ARG} \
    > "${OUTPUT_PATH}/opt.log" 2>&1

echo "Relaxation completed. Results in ${OUTPUT_PATH}"