#!/usr/bin/env bash
# A/B the card-reward skip policy.
#
#   ./ab_card_skip.sh <sessions-per-arm> <steps-per-session>
#
# Everything except --card-skip-policy is identical between the arms.
#
# Three things this script does that earlier versions did not, each because an
# earlier attempt produced meaningless numbers:
#
#   1. --stop-on-run-end, so ONE SESSION = ONE RUN. Without it the bot dies,
#      returns to the menu and starts a fresh run inside the same session, and a
#      session yields several run records that cannot be attributed to an arm.
#   2. Each run record is attributed to the arm that produced it, by diffing the
#      history directory around each session. The old version just took "the
#      last 2N runs", which silently compared old runs from earlier experiments.
#   3. Every attributed run carries a hash of the decision layer, and the report
#      compares only runs from ONE revision. A change to brain.py / facts.py /
#      jev.py / run.py changes how both arms behave, so runs collected before it
#      are not comparable with runs collected after it -- and the working copy is
#      usually dirty, so a commit id would not distinguish the revisions.
set -u

N="${1:-5}"
# Steps per session. This must be generous: measured, a run that is still going
# at step 600 produces NO run record, so the sample silently fills with runs
# that died early -- exactly the wrong bias. --stop-on-run-end is the real stop
# condition; this is only a safety cap.
STEPS="${2:-4000}"
HIST="$HOME/Library/Application Support/SlayTheSpire2/steam/76561198141226155/modded/profile1/saves/history"
RESULTS=/tmp/ab_results.tsv

# The revision every arm in this invocation is collected from. Content hash, not
# a commit id: the decision layer is normally edited in place, uncommitted.
#
# Recomputed BEFORE EVERY SESSION, not once: this file is edited in place while
# experiments run, which is the whole reason the guard exists. A hash taken once
# at startup would be stamped onto every row regardless, and because $RESULTS is
# truncated below, the mixed-revision warning could never fire.
CODE_FILES=(brain.py facts.py jev.py run.py sts2.py)
code_hash() {
  cat "${CODE_FILES[@]}" | { shasum -a 256 2>/dev/null || sha256sum; } | cut -c1-12
}
CODE_HASH=$(code_hash)
ABORTED=0
echo "code revision: $CODE_HASH  (${CODE_FILES[*]})"

# Fail fast if the bot cannot start a run (e.g. a pending Timeline epoch).
if ! python3 run.py --steps 1 >/tmp/preflight.log 2>&1; then
  cat /tmp/preflight.log
  exit 2
fi

: > "$RESULTS"

snapshot() { ls "$HIST" 2>/dev/null | sort; }

run_arm() {
  local policy="$1"
  for i in $(seq 1 "$N"); do
    local before after attempt=0 rc=1 session_hash
    before=$(snapshot)

    # An arm collected across two revisions is not an arm. Stop the experiment
    # rather than stamp rows with a revision that did not produce them.
    session_hash=$(code_hash)
    if [ "$session_hash" != "$CODE_HASH" ]; then
      echo
      echo "!! ABORT: the decision layer changed mid-experiment."
      echo "!!   started on $CODE_HASH, now $session_hash"
      echo "!!   Rows already collected are in $RESULTS and are NOT comparable"
      echo "!!   with anything collected on the other revision."
      echo "!!   Re-run BOTH arms from one revision."
      ABORTED=1
      return 1
    fi

    while [ "$attempt" -lt 3 ]; do
      attempt=$((attempt + 1))
      echo "--- $policy $i/$N (attempt $attempt) rev=$session_hash ---"
      python3 run.py --steps "$STEPS" --pause 0.3 \
        --card-skip-policy "$policy" --stop-on-run-end \
        > "/tmp/ab_${policy}_${i}.log" 2>&1
      rc=$?
      [ "$rc" -eq 0 ] && break
      [ "$rc" -ne 3 ] && break          # 2 = blocked, not worth retrying
      echo "    retrying after model failures..."
      sleep 20
    done

    after=$(snapshot)
    local new_runs session
    new_runs=$(comm -13 <(echo "$before") <(echo "$after") | tr '\n' ' ')
    # The session id the bot printed, so an attribution row joins to the same
    # decisions.jsonl rows and to sessions.jsonl. Without it the two ends of the
    # join carry different identifiers and nothing lines up.
    session=$(sed -n 's/^session: //p' "/tmp/ab_${policy}_${i}.log" | head -1)
    echo "    exit=$rc session=${session:-?} takes=$(grep -c 'select_card_reward' "/tmp/ab_${policy}_${i}.log") skips=$(grep -c 'skip_card_reward' "/tmp/ab_${policy}_${i}.log") runs='${new_runs}'"

    # Stamped with the hash that produced THIS run, not the startup hash.
    for r in $new_runs; do
      printf '%s\t%s\t%s\t%s\n' "$policy" "$r" "$session_hash" "${session:-?}" >> "$RESULTS"
    done
  done
  return 0
}

echo "### ARM A: jev decides skip"
run_arm jev
echo
if [ "$ABORTED" -eq 0 ]; then
  echo "### ARM B: combined (jev + a floor)"
  run_arm combined
else
  echo "### ARM B: skipped (aborted mid-experiment)"
fi

echo
POLICY=combined RESULTS="$RESULTS" CODE_HASH="$CODE_HASH" HIST="$HIST" python3 - <<'PY'
import json, os, statistics as st

# The same directory the shell snapshotted. A second hardcoded copy of this path
# silently reported "no runs recorded" for sessions that did produce runs.
hist = os.environ["HIST"]
results = os.environ["RESULTS"]
this_hash = os.environ["CODE_HASH"]

rows = []
with open(results) as fh:
    for line in fh:
        parts = line.rstrip("\n").split("\t")
        if len(parts) < 2 or not parts[1]:
            continue
        rows.append((parts[0], parts[1],
                     parts[2] if len(parts) > 2 and parts[2] else "unstamped",
                     parts[3] if len(parts) > 3 and parts[3] else "?"))

hashes = sorted({h for _, _, h, _ in rows})
print(f"code revision(s) in {results}: {', '.join(hashes) if hashes else '(none)'}")
if len(hashes) > 1:
    print()
    print("!! MIXED REVISIONS. A change to the decision layer changes how BOTH")
    print("!! arms behave, so a run collected before it is not comparable with")
    print("!! one collected after it. Only ONE revision is scored below; re-run")
    print("!! BOTH arms on one revision before interpreting anything.")
# Score the revision this invocation collected from when it has rows; otherwise
# the most recently appended one. Never "the largest hash": hex is not ordered.
if this_hash in hashes:
    current = this_hash
elif rows:
    current = rows[-1][2]
else:
    current = this_hash
print(f"scoring revision: {current}")
print()

arms = {"jev": [], "combined": []}
sessions = {"jev": 0, "combined": 0}
for policy, run, code, session in rows:
    if policy not in arms or code != current:
        continue
    sessions[policy] += 1
    path = os.path.join(hist, run)
    if not os.path.exists(path):
        continue
    d = json.load(open(path))
    pl = (d.get("players") or [{}])[0]
    mph = d.get("map_point_history") or []
    arms[policy].append({
        "killed": d.get("killed_by_encounter", "?").replace("ENCOUNTER.", ""),
        "deck": len(pl.get("deck") or []),
        "pts": sum(len(a) for a in mph if isinstance(a, list)),
        "session": session,
    })

print(f"{'arm':<10} {'sess':>4} {'runs':>4} {'deck':>6} {'pts':>6} {'added/pt':>9} {'act2':>6}")
print("-" * 54)
for name, grp in arms.items():
    if not grp:
        print(f"{name:<10} {sessions[name]:>4} {0:>4}   (no runs recorded)")
        continue
    decks = [r["deck"] for r in grp]
    pts = [r["pts"] for r in grp]
    act2 = sum(1 for r in grp if r["pts"] > 17)
    added = max(0.0, st.mean(decks) - 10) / max(1e-9, st.mean(pts))
    print(f"{name:<10} {sessions[name]:>4} {len(grp):>4} {st.mean(decks):>6.1f} "
          f"{st.mean(pts):>6.1f} {added:>9.2f} {act2:>4}/{len(grp)}")

print()
print("sess     = sessions that produced a COMPLETED run record")
print("runs     = run records attributed to this arm, on the scored revision")
print("added/pt = non-starting cards per map point reached  <- the metric that")
print("           matters, because deck size at death scales with run length")
print("act2     = runs past the Act 1 boss (17 map points)")

thin = [n for n, g in arms.items() if len(g) < 5]
if thin:
    print()
    print(f"!! THIN SAMPLE: {', '.join(thin)} have fewer than 5 runs.")
    print("!! Sessions that hit the step cap mid-run produce NO record, so a thin")
    print("!! sample is biased toward runs that died early. Raise STEPS and re-run.")
print()
for name, grp in arms.items():
    if grp:
        print(f"{name} killed by: {[r['killed'] for r in grp]}")
        print(f"{name} sessions : {[r['session'] for r in grp]}")
print()
print("sessions are the join key: grep one in capture/decisions.jsonl for that")
print("session's decisions, and in capture/sessions.jsonl for its outcome row.")
PY