#!/usr/bin/env bash

# Export Weibull shape-parameter summaries for the all_future models trained
# with smoking/alcohol/BMI extra information.
#
# Bash 4.2 compatible. Run from the DeepHealth repository root on the Linux
# server, for example:
#
#   bash run_weibull_shape_exports.sh
#
# Optional overrides:
#   PYTHON=python3 DEVICE=cuda BATCH_SIZE=128 NUM_WORKERS=0 ROW_BATCH_SIZE=512 \
#     bash run_weibull_shape_exports.sh

set -euo pipefail

PYTHON="${PYTHON:-python}"
DEVICE="${DEVICE:-cuda}"
BATCH_SIZE="${BATCH_SIZE:-128}"
NUM_WORKERS="${NUM_WORKERS:-0}"
ROW_BATCH_SIZE="${ROW_BATCH_SIZE:-512}"
LANDMARK_START="${LANDMARK_START:-40}"
LANDMARK_STOP="${LANDMARK_STOP:-80}"
LANDMARK_STEP="${LANDMARK_STEP:-5}"
HORIZONS="${HORIZONS:-1,5,10}"

RUNS=(
  "runs/relative_weibull_all_future_pure_disease_20260620_095229"
  "runs/relative_mixed_all_future_pure_disease_20260620_132415"
  "runs/absolute_weibull_all_future_pure_disease_20260620_114816"
  "runs/absolute_mixed_all_future_pure_disease_20260620_161804"
)

echo "Python: ${PYTHON}"
echo "Device: ${DEVICE}"
echo "Batch size: ${BATCH_SIZE}"
echo "Workers: ${NUM_WORKERS}"
echo "Row batch size: ${ROW_BATCH_SIZE}"
echo "Horizons: ${HORIZONS}"
echo

for run_path in "${RUNS[@]}"; do
  if [[ ! -d "${run_path}" ]]; then
    echo "[ERROR] Missing run directory: ${run_path}" >&2
    exit 1
  fi
  if [[ ! -f "${run_path}/best_model.pt" ]]; then
    echo "[ERROR] Missing checkpoint: ${run_path}/best_model.pt" >&2
    exit 1
  fi
  if [[ ! -f "${run_path}/train_config.json" ]]; then
    echo "[ERROR] Missing config: ${run_path}/train_config.json" >&2
    exit 1
  fi

  output_path="${run_path}/weibull_shape_parameter_stats_test"
  echo "=== Exporting Weibull shape stats: ${run_path} ==="
  "${PYTHON}" export_weibull_shape_parameter_stats.py \
    --run_path "${run_path}" \
    --output_path "${output_path}" \
    --eval_split test \
    --device "${DEVICE}" \
    --batch_size "${BATCH_SIZE}" \
    --num_workers "${NUM_WORKERS}" \
    --row_batch_size "${ROW_BATCH_SIZE}" \
    --hidden_cache_dtype float32 \
    --landmark_start "${LANDMARK_START}" \
    --landmark_stop "${LANDMARK_STOP}" \
    --landmark_step "${LANDMARK_STEP}" \
    --horizons "${HORIZONS}" \
    --include_all_token_rho_summary
  echo "Wrote: ${output_path}"
  echo
done

echo "All Weibull shape exports completed."