#!/usr/bin/env bash
# Copyright (C) 2026 Asger Gitz-Johansen
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
DATA_DIR="${ROOT_DIR}/data"
OUTPUT_FILE="${ROOT_DIR}/sleep-model.pt"
MODEL_FILE="${OUTPUT_FILE}"

if [[ ${1:-} == "-h" || ${1:-} == "--help" ]]; then
  cat <<'EOF'
Usage: scripts/train_from_data.sh [-o OUTPUT_FILE]

Trains the model on complete session triplets found under data/{motion,heart_rate,labels}/.
EOF
  exit 0
fi

if [[ ${1:-} == "-o" || ${1:-} == "--output" ]]; then
  OUTPUT_FILE="${2:?Missing output file after -o/--output}"
  MODEL_FILE="${OUTPUT_FILE}"
fi

shopt -s nullglob

checkpoint_class_count() {
  python - "$1" <<'PY'
from pathlib import Path
import sys
import torch

payload = torch.load(Path(sys.argv[1]), weights_only=True)
state_dict = payload.get("model_state_dict", payload)
print(state_dict["net.2.bias"].shape[0])
PY
}

motions=("${DATA_DIR}/motion"/*_acceleration.txt)

subject_ids=()

for motion_file in "${motions[@]}"; do
  sample_id="$(basename "${motion_file}" _acceleration.txt)"
  heartrate_file="${DATA_DIR}/heart_rate/${sample_id}_heartrate.txt"
  label_file="${DATA_DIR}/labels/${sample_id}_labeled_sleep.txt"

  if [[ -f "${heartrate_file}" && -f "${label_file}" ]]; then
    subject_ids+=("${sample_id}")
  fi
done

if [[ ${#subject_ids[@]} -eq 0 ]]; then
  printf 'No complete data triplets found under %s\n' "${DATA_DIR}" >&2
  exit 1
fi

mapfile -t subject_ids < <(printf '%s\n' "${subject_ids[@]}" | sort -u)

printf 'Training on %d subject(s)\n' "${#subject_ids[@]}"

export PYTHONPATH="${ROOT_DIR}/src${PYTHONPATH:+:${PYTHONPATH}}"

for subject_id in "${subject_ids[@]}"; do
  acceleration_file="${DATA_DIR}/motion/${subject_id}_acceleration.txt"
  heartrate_file="${DATA_DIR}/heart_rate/${subject_id}_heartrate.txt"
  label_file="${DATA_DIR}/labels/${subject_id}_labeled_sleep.txt"
  resume_args=()

  printf 'Training subject %s\n' "${subject_id}"

  if [[ -f "${MODEL_FILE}" ]]; then
    checkpoint_classes="$(checkpoint_class_count "${MODEL_FILE}")"
    if [[ "${checkpoint_classes}" == "3" ]]; then
      resume_args=(-i "${MODEL_FILE}")
    else
      printf 'Skipping checkpoint for subject %s: model classes=%s, expected=3\n' "${subject_id}" "${checkpoint_classes}"
    fi
  fi

  args=(python -m sleep_detection train -a "${acceleration_file}" -H "${heartrate_file}" -l "${label_file}" -o "${MODEL_FILE}")
  if [[ ${#resume_args[@]} -gt 0 ]]; then
    args+=("${resume_args[@]}")
  fi

  "${args[@]}"
done

printf 'Saved trained model to %s\n' "${MODEL_FILE}"