#!/usr/bin/env bash
# Counts theorems/lemmas and sorry placeholders per Lean file.
# Writes lean/proof-status.json for the shields.io badge.
# Optionally injects a section into README.md between marker comments.
#
# Usage:
#   bash lean/proof_status.sh              # print table to stdout
#   bash lean/proof_status.sh --json-only  # update JSON only
#   bash lean/proof_status.sh --inject-readme  # update JSON + inject section into README.md

LEAN_DIR="$(dirname "$0")/ForestIPM"
JSON_OUT="$(dirname "$0")/proof-status.json"
README="README.md"
MODE="${1:-}"

# Parallel arrays: module name, component label (empty = continuation of previous group)
modules=(
  "Growth"          "Demographic models"
  "Survival"        ""
  "Ingrowth"        ""
  "Competition"     "Covariates"
  "Climate"         ""
  "MidpointRule"   "Population model"
  "GaussLegendre"   ""
  "Kernel"          ""
  "AsymptoticLambda" "Engines"
  "CommunityDynamic" ""
)

total_theorems=0
total_sorry=0

# Collect per-module counts into indexed arrays
n_modules=$(( ${#modules[@]} / 2 ))
mod_names=()
mod_components=()
mod_theorems=()
mod_sorry=()

for (( i=0; i<${#modules[@]}; i+=2 )); do
  module="${modules[$i]}"
  component="${modules[$((i+1))]}"
  file="$LEAN_DIR/${module}.lean"
  theorems=$(grep -c '^theorem \|^lemma ' "$file" || true)
  sorry=$(grep -c '^[[:space:]]*sorry[[:space:]]*$' "$file" || true)
  mod_names+=("$module")
  mod_components+=("$component")
  mod_theorems+=("$theorems")
  mod_sorry+=("$sorry")
  total_theorems=$(( total_theorems + theorems ))
  total_sorry=$(( total_sorry + sorry ))
done

total_proved=$(( total_theorems > total_sorry ? total_theorems - total_sorry : 0 ))

# --- Build the markdown table ---
build_table() {
  printf "| Component | Module | Theorems | Proved | Sorry | Status |\n"
  printf "|-----------|--------|----------|--------|-------|--------|\n"

  for (( i=0; i<${#mod_names[@]}; i++ )); do
    module="${mod_names[$i]}"
    component="${mod_components[$i]}"
    t="${mod_theorems[$i]}"
    s="${mod_sorry[$i]}"
    p=$(( t > s ? t - s : 0 ))
    if [[ "$s" -eq 0 ]]; then icon="fully proved"; else icon="$s open"; fi
    printf "| %-18s | %-18s | %8d | %6d | %5d | %-14s |\n" \
      "$component" "$module" "$t" "$p" "$s" "$icon"
  done

  printf "| %-18s | %-18s | %8d | %6d | %5d | %-14s |\n" \
    "" "**Total**" "$total_theorems" "$total_proved" "$total_sorry" ""
}

# --- Build the full README section ---
build_section() {
  cat <<'HEADER'
## Mathematical Specification (Lean 4)

The mathematical model underlying `{forestIPM}` has been formally verified using
[Lean 4](https://lean-lang.org/) and [Mathlib](https://leanprover-community.github.io/mathlib4_docs/).
Each component of the model pipeline has a corresponding Lean module that proves key mathematical properties — such as convergence, non-negativity, and boundedness — hold by construction, independent of any numerical implementation.

These are not tests of the R code; they are machine-checked proofs that the *mathematical model* is internally consistent. The source files live in [`lean/`](lean/).

HEADER

  build_table

  printf "\n> Table auto-generated by [lean/proof_status.sh](lean/proof_status.sh). "
  printf "\"Sorry\" indicates theorems stated but not yet proved (open proof obligations).\n"
}

# --- Update proof-status.json ---
if [[ "$total_sorry" -eq 0 ]]; then
  color="brightgreen"
elif [[ "$total_proved" -ge $(( total_theorems * 3 / 4 )) ]]; then
  color="yellowgreen"
elif [[ "$total_proved" -ge $(( total_theorems / 2 )) ]]; then
  color="yellow"
else
  color="orange"
fi

cat > "$JSON_OUT" <<EOF
{
  "schemaVersion": 1,
  "label": "lean proofs",
  "message": "${total_proved}/${total_theorems} proved",
  "color": "${color}"
}
EOF

# --- Dispatch on mode ---
if [[ "$MODE" == "--json-only" ]]; then
  exit 0
fi

if [[ "$MODE" == "--inject-readme" ]]; then
  SECTION=$(build_section)
  python3 - <<PYEOF
import re, pathlib, sys

section = """<!-- lean-proofs-start -->
${SECTION}
<!-- lean-proofs-end -->"""

readme = pathlib.Path("${README}").read_text()
updated, n = re.subn(
    r'<!-- lean-proofs-start -->.*?<!-- lean-proofs-end -->',
    section,
    readme,
    flags=re.DOTALL
)
if n == 0:
    sys.exit("ERROR: markers <!-- lean-proofs-start/end --> not found in ${README}")
pathlib.Path("${README}").write_text(updated)
print(f"README.md updated: ${total_proved}/${total_theorems} theorems proved")
PYEOF
else
  build_section
  echo ""
  echo "proof-status.json updated: ${total_proved}/${total_theorems} proved"
fi