#!/bin/bash
#SBATCH --account=<your-slurm-account>
#SBATCH --partition=l40s
#SBATCH --nodes=1
#SBATCH --ntasks=8
#SBATCH --gres=gpu:l40s:1
#SBATCH --time=12:00:00
#SBATCH --job-name=ddpm_triangle
#SBATCH --mail-user=<your-email>   # replace before submitting
#SBATCH --output=slurm-ddpm-triangle-%j.out
#SBATCH --error=slurm-ddpm-triangle-%j.err

# Run surrogate Ωm–σ8 chain export for DDPM-2 and DDPM-6, then a joint triangle plot.
# If you have your own copies of ddpm_triangle_integration.py / triangle_plot_posterior.py
# under $ROOT/scripts (e.g. copied from ~/Downloads), they override the repo versions.
#
#   sbatch .../run_triangle_ddpm_both.sh
#
#   sbatch --export=OUTPUT_DIR=/path/to/out,TEST_INDEX=56 .../run_triangle_ddpm_both.sh
#
# Interactive:
#   module load python/miniconda3-py3.12-usr
#   bash .../run_triangle_ddpm_both.sh

set -euo pipefail

ROOT="<DDPM_ROOT>/Models"
cd "$ROOT"

module load python/miniconda3-py3.12-usr

OUT="${OUTPUT_DIR:-${ROOT}/ddpm_triangle_out}"
TEST_IX="${TEST_INDEX:-56}"
GRID="${GRID_POINTS:-14}"

mkdir -p "${OUT}"

INTEG="${DDPM_TRIANGLE_INTEGRATION_PY:-${ROOT}/scripts/ddpm_triangle_integration.py}"
PLOT="${DDPM_TRIANGLE_POSTERIOR_PY:-${ROOT}/scripts/triangle_plot_posterior.py}"

for f in "$INTEG" "$PLOT"; do
  if [[ ! -f "$f" ]]; then
    echo "Missing: $f"
    exit 1
  fi
done

CHAIN2="${OUT}/chain_surrogate_ix${TEST_IX}_ddpm2.npz"
CHAIN6="${OUT}/chain_surrogate_ix${TEST_IX}_ddpm6_truth_tail.npz"

echo "==============================================="
echo "Job: ${SLURM_JOB_ID:-local}  OUT=${OUT}  test_ix=${TEST_IX}"
echo "==============================================="

python "$INTEG" \
  --label-dim 2 \
  --test-index "${TEST_IX}" \
  --grid "${GRID}" \
  -o "${CHAIN2}"

python "$INTEG" \
  --label-dim 6 \
  --test-index "${TEST_IX}" \
  --six-tail-mode truth \
  --grid "${GRID}" \
  -o "${CHAIN6}"

python "$PLOT" \
  -i "${CHAIN2}" "${CHAIN6}" \
  --labels "DDPM-2" "DDPM-6" \
  -o "${OUT}/triangle_ddpm2_ddpm6_ix${TEST_IX}.png"

echo "Chains: ${CHAIN2} ${CHAIN6}"
echo "Triangle: ${OUT}/triangle_ddpm2_ddpm6_ix${TEST_IX}.png"
echo "Finished: $(date)"