#!/usr/bin/env bash
# Build one patient of the benchmark end to end; all outputs go to build/<patient_id>/.
#
#   bash build_patient.sh <patient_id> ["<clinical descriptor>"]
#
# Stages: A preprocessing (deterministic) -> B event chain (extract), evidence chain (evidence), change tags (tags)
# -> C questions (axes 1-2: seeds, LLM, post; axes 3-8: LLM discovery; golds written by the model for axes 3-4, from
#    templates for axes 5-8)
# -> D claim registry and synthetic-question targets -> E synthetic questions, gates, injection
# -> F dialogue simulation -> G checkpoints, validation, manifest.
# Requires the visit table and the term glossary (set at the top; neither distributed) and OPENAI_API_KEY.
# Set MEMEBENCH_ENV to a script that activates the Python environment, if needed, or PYTHON to the interpreter.
# MEMEBENCH_FROM=<step> resumes from a step (B1 ... G), keeping earlier outputs. Once stage E has written synthetic
# facts into the chains (event_chain.preinject.jsonl holds the chain before them), only F and G can be resumed;
# B1 rebuilds the event chain from the extraction.
# MEMEBENCH_WORKERS (default 64) and MEMEBENCH_WORKERS_SMALL (default 32) set the parallel calls of the LLM steps.
set -euo pipefail
export PYTHONHASHSEED=0

# Inputs. The visit table holds the synthetic timeline (SOAP-structured), the output of the abstraction and
# regeneration steps described in the paper; those steps are not part of this code.
# Fill in before running.
VISIT_TABLE="${MEMEBENCH_VISIT_TABLE:-<VISIT_TABLE>}"              # parquet, one row per visit
GLOSSARY="${MEMEBENCH_GLOSSARY:-<TERM_GLOSSARY>}"                  # term <TAB> normalized form, one pair per line
for f in "$VISIT_TABLE" "$GLOSSARY"; do
  [[ $f == \<* ]] && { echo "Set the visit table and the term glossary at the top of build_patient.sh"; exit 1; }
done
export MEMEBENCH_VISIT_TABLE="$VISIT_TABLE" MEMEBENCH_GLOSSARY="$GLOSSARY"

# Models (OpenAI-compatible names). Fill in before running.
EXTRACTION_MODEL="${MEMEBENCH_EXTRACTION_MODEL:-<EXTRACTION_MODEL>}"  # B1 event extraction, B3 change tags
GENERATION_MODEL="${MEMEBENCH_GENERATION_MODEL:-<GENERATION_MODEL>}"  # B2, C questions, E synthetic/gates, G validation
DIALOGUE_MODEL="${MEMEBENCH_DIALOGUE_MODEL:-<DIALOGUE_MODEL>}"        # F dialogue simulation
for m in "$EXTRACTION_MODEL" "$GENERATION_MODEL" "$DIALOGUE_MODEL"; do
  [[ $m == \<* ]] && { echo "Set the models at the top of build_patient.sh"; exit 1; }
done
export MEMEBENCH_EXTRACTION_MODEL="$EXTRACTION_MODEL" MEMEBENCH_GENERATION_MODEL="$GENERATION_MODEL" \
       MEMEBENCH_DIALOGUE_MODEL="$DIALOGUE_MODEL"

# Validation (stage G): axes whose questions are not filtered by the no-memory probe (comma-separated axis numbers).
NO_MEMORY_EXEMPT_AXES="${MEMEBENCH_NO_MEMORY_EXEMPT_AXES-5}"
export MEMEBENCH_NO_MEMORY_EXEMPT_AXES="$NO_MEMORY_EXEMPT_AXES"

PID="${1:?usage: bash build_patient.sh <patient_id> [<clinical descriptor>]}"
PDESC="${2:-}"
BASE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
W="$BASE/build/$PID"
cd "$BASE"
if [ -n "${MEMEBENCH_ENV:-}" ]; then
  source "$MEMEBENCH_ENV" >/dev/null || { echo "MEMEBENCH_ENV failed: $MEMEBENCH_ENV"; exit 1; }
fi

mkdir -p "$W"
export MEMEBENCH_WORK="$W" MEMEBENCH_PATIENT="$PID"

say() { printf "\n\033[1m== %s\033[0m\n" "$*"; }
run() { "${PYTHON:-python}" -u -m "$@"; }

say "A. preprocessing"
[ -f "$W/ALL_visits.txt" ] || run pipeline.preprocess export --out "$W/ALL_visits.txt"
[ -f "$W/visits.jsonl" ]      || run pipeline.preprocess parse --inp "$W/ALL_visits.txt" --out "$W/visits.jsonl"
[ -d "$W/visits" ]            || run pipeline.preprocess split              # -> visits/v*.json + visit_manifest.json
run pipeline.preprocess vitals                                               # -> vitals.json, vitals_timeline.json
[ -f "$W/patient.env" ]       || run pipeline.preprocess facts --pdesc "$PDESC"
source "$W/patient.env"

STEPS=(B1 B2 B3 C1 C2 D E F G)
step_i() { local i; for i in "${!STEPS[@]}"; do [ "${STEPS[$i]}" = "$1" ] && { echo "$i"; return 0; }; done; return 1; }
FROM="${MEMEBENCH_FROM:-B1}"
FROM_I=$(step_i "$FROM") || { echo "MEMEBENCH_FROM=$FROM is not one of: ${STEPS[*]}"; exit 1; }
from() { [ "$(step_i "$1")" -ge "$FROM_I" ]; }
if [ "$FROM" != "B1" ]; then
  [ -f "$W/event_records.json" ] && [ -f "$W/event_chain.jsonl" ] \
    || { echo "MEMEBENCH_FROM=$FROM but $W has no finished B1 output"; exit 1; }
  if [ -f "$W/event_chain.preinject.jsonl" ] && [ "$FROM_I" -lt "$(step_i F)" ]; then
    echo "MEMEBENCH_FROM=$FROM: stage E has already written synthetic facts into the event and evidence chains of"
    echo "$W (event_chain.preinject.jsonl is the chain before them); resume from F or G, or rerun from B1"
    exit 1
  fi
  if [ "$FROM_I" -ge "$(step_i F)" ] && [ ! -f "$W/event_chain.preinject.jsonl" ] \
     && grep -qs '"grounding": "synthetic' "$W"/questions_axis*.jsonl; then
    echo "MEMEBENCH_FROM=$FROM: the question banks hold synthetic questions whose facts stage E has not written into"
    echo "the chains; resume from D"
    exit 1
  fi
  say "RESUMING from $FROM (steps before it are skipped)"
fi

if from B1; then
say "B1. event extraction (one call per visit)"
rm -f "$W/failed_extract.json"
run pipeline.llm extract --manifest "$W/visit_manifest.json" --out "$W/event_records.json" \
    --model "$EXTRACTION_MODEL" --workers "${MEMEBENCH_WORKERS:-64}"
[ ! -f "$W/failed_extract.json" ] || { echo "extraction failed for some visits; see $W/failed_extract.json"; exit 1; }
run pipeline.event_chain event-chain "$W/event_records.json"                 # -> event_chain.jsonl
rm -f "${W:?}/event_chain.preinject.jsonl"                                    # the new chain has no synthetic facts
fi

if from B2; then
say "B2. evidence chain"
run pipeline.evidence_chain dept-index                                       # -> _dept_event_index.json
run pipeline.evidence_chain prep                                             # -> _evidence_manifest.json + batches
run pipeline.llm step --step evidence --manifest "$W/_evidence_manifest.json" \
    --out "$W/_evidence_records.json" --model "$GENERATION_MODEL" --workers "${MEMEBENCH_WORKERS_SMALL:-32}"
[ ! -f "$W/failed_evidence.json" ] || { echo "evidence failed for some batches; see $W/failed_evidence.json"; exit 1; }
run pipeline.evidence_chain post "$W/_evidence_records.json"                 # -> evidence_chain.jsonl
fi

if from B3; then
say "B3. change tags"
run pipeline.event_chain windows                                             # -> _tag_manifest.json + batches
run pipeline.llm step --step tags --manifest "$W/_tag_manifest.json" \
    --out "$W/_tag_records.json" --model "$EXTRACTION_MODEL" --workers "${MEMEBENCH_WORKERS_SMALL:-32}"
[ ! -f "$W/failed_tags.json" ] || { echo "tags failed for some windows; see $W/failed_tags.json"; exit 1; }
run pipeline.event_chain tags-post "$W/_tag_records.json"                    # -> tags.jsonl
fi

if from C1; then
say "C1. axes 1-2 (seeds -> LLM -> post)"
for N in 1 2; do
  run "pipeline.questions.axis$N" seeds
  run pipeline.llm step --step "axis$N" --input "$W/_seed_axis${N}.json" \
      --out "$W/_axis${N}_generated.json" --model "$GENERATION_MODEL"
  run "pipeline.questions.axis$N" post "$W/_axis${N}_generated.json"
done
fi

if from C2; then
say "C2. axes 3-8 (LLM discovery; golds from the model for axes 3-4, from templates for axes 5-8)"
for N in 3 4 5 6 7 8; do run "pipeline.questions.axis$N"; done
run pipeline.synthetic banks                                             # -> questions_axis{3..8}.jsonl
fi

if from D; then
say "D. claim registry and synthetic-question targets"
[ -f "$W/injection_registry.json" ] || run pipeline.synthetic registry --init
run pipeline.synthetic claims
run pipeline.synthetic targets                                           # -> synthetic_targets.json
fi

if from E; then
say "E. synthetic questions and injection"
SHORT=$(run pipeline.synthetic short-axes)
if [ -n "$SHORT" ]; then
  run pipeline.synthetic generate --axes "$SHORT" --model "$GENERATION_MODEL" \
    || echo "  synthetic generation reported failures; continuing"
  run pipeline.synthetic wire                                            # bank += synthetic Q; injection_plan.json
  rm -f "$W/injection_registry.json"
  run pipeline.synthetic registry --init
  run pipeline.synthetic batches
  run pipeline.synthetic claims >/dev/null                               # all claims incl. synthetic
  run pipeline.gates lint || { echo "LINT FAILED: not injecting"; exit 1; }
  run pipeline.gates persona || { echo "PERSONA GATE FAILED: the persona check did not complete"; exit 1; }
  run pipeline.gates lint || { echo "LINT FAILED after the persona gate: not injecting"; exit 1; }
  run pipeline.synthetic apply                                           # atoms -> event_chain + evidence_chain
else
  say "E. no synthetic questions needed"
  rm -f "${W:?}/synthetic_pairs.json"
  run pipeline.synthetic wire
fi
fi

if from F; then
say "F. dialogues"
[ "$FROM_I" -ge "$(step_i F)" ] || rm -f "${W:?}/dialogues.jsonl"
run pipeline.dialogues generate --all --workers "${MEMEBENCH_WORKERS:-64}" --out "$W/dialogues.jsonl"
if [ "$(run pipeline.dialogues count-empty)" != "0" ]; then
  run pipeline.dialogues generate --all --workers "${MEMEBENCH_WORKERS:-64}" --out "$W/dialogues.jsonl"
fi
N_EMPTY=$(run pipeline.dialogues count-empty)
[ "$N_EMPTY" = "0" ] || { echo "dialogues: $N_EMPTY visits still empty after a retry pass"; exit 1; }
fi

if from G; then
say "G1. checkpoints"
run pipeline.validation assign                                             # -> checkpoints.jsonl
run pipeline.validation audit                                              # -> checkpoint_audit.json
say "G2. validation (oracle + no-memory probe)"
run pipeline.validation validate --workers "${MEMEBENCH_WORKERS:-64}"      # -> validation.jsonl
N_ERR=$(run pipeline.validation count-errors)
[ "$N_ERR" = "0" ] \
  || { echo "validation: $N_ERR questions have failed probe or judge calls; see validation.jsonl"; exit 1; }
say "G3. package"
run pipeline.validation manifest                                           # -> MANIFEST.json
fi

say "DONE: $W"