#!/usr/bin/env bash
# Determinism self-test for Tier 1.5 mid-smoke (spec §3.5 + §1.3 of the
# determinism foundation).
#
# Runs scripts/local-mid-smoke.sh twice with the same seed and compares
# the per-step `checksums.*` leaves emitted by `build_diag_value` (added
# in determinism Phase 1, 2026-06-02). The trainer must be deterministic
# given a fixed seed for Tier 1.5 to be trustworthy — otherwise behavioral
# signals are polluted by random-init drift.
#
# Output focuses on the FIRST divergent (step, leaf) pair. That's the
# culprit identification — Phase 1's whole point. The trailing block
# also lists ALL leaves diverging at the first divergent step (so we
# catch the common "multiple components co-diverge" case).
#
# Tolerance: 1e-5 relative, 1e-7 absolute (matches fp32 noise floor).
# Larger gaps indicate a real divergence.
#
# Per `feedback_local_smoke_before_cluster` +
# `pearl_scoped_init_seed_for_reproducibility`.
#
# Usage:
#   ./scripts/determinism-check.sh                       # default seed=42, b=128
#   ./scripts/determinism-check.sh --seed 99             # custom seed
#   ./scripts/determinism-check.sh --n-backtests 64      # smaller batch
#   ./scripts/determinism-check.sh --quick               # 200-step abbreviated run
#
# Exit codes:
#   0  Deterministic (all per-step checksums match within tolerance)
#   1  Drift detected (see first-divergence report)
#   2  Inputs missing (a run failed)

set -euo pipefail

SEED=42
N_BACKTESTS=128
N_STEPS=2000
N_EVAL_STEPS=500
QUICK=false
DEBUG_DUMP=false
DEBUG_DUMP_MAMBA2=false
DEBUG_DUMP_RL=false
DEBUG_DUMP_BACKWARD=false

usage() {
  cat <<EOF
Usage: $(basename "$0") [OPTIONS]

Determinism self-test: run mid-smoke twice with same seed, diff all
\`checksums.*\` leaves leaf-by-leaf across every step. Reports the first
(step, leaf) pair where determinism breaks — that names the culprit.

Options:
  --seed N             Seed (default: $SEED)
  --n-backtests N      Batch size (default: $N_BACKTESTS)
  --quick              200 steps + 50 eval (faster, less coverage)
  --debug-dump         Phase 2 sub-investigation: enable raw PER-state
                       dumps for step 0/1/2 of each run, then invoke
                       scripts/compare-per-state.py to diff them.
                       Use this once to localise the PER non-determinism
                       source (PRNG vs tree-walk vs tree-rebuild).
  --debug-dump-mamba2  Phase 2.2 sub-investigation: enable raw mamba2
                       weight + intermediate-activation dumps for steps
                       0/1/2/3 of each run, then invoke
                       scripts/compare-mamba2-state.py to diff them.
                       Use this to disambiguate weight-drift (upstream
                       Adam non-det) from forward/backward kernel
                       non-determinism in the mamba2 selective scan.
  --debug-dump-rl      Phase 2.5 sub-investigation: enable raw RL-loop
                       state dumps (CfC carry, PER state, action RNG,
                       LOB sim, ISV controllers, reward shaping) for
                       steps 0/1/2/3 of each run, then invoke
                       scripts/compare-rl-state.py to diff them.
                       Use this after Phase 2.4 falsified the mamba2
                       hypothesis (all 52 dumps EQUAL at step 1 but
                       divergence enters at step 2 from inside the RL
                       step loop).
  --debug-dump-backward  Phase 2.6 sub-investigation: enable raw backward-
                       pipeline dumps (PER replay sample h_t / head fwd
                       outputs / pre-reduce per-batch param-grad
                       scratches / per-head grad_h_t + combined
                       accumulator) for steps 0/1/2/3 of each run, then
                       invoke scripts/compare-backward-state.py to diff
                       them. Use this after Phase 2.5 confirmed forward
                       activations EQUAL at step 2 but backward outputs
                       DIVERGE — the residual non-determinism lives in
                       step 2's gradient-emission chain (head bwd,
                       reduce_axis0, or grad_h_accumulate).
  -h, --help           Show this help
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --seed)              SEED="$2"; shift 2 ;;
    --n-backtests)       N_BACKTESTS="$2"; shift 2 ;;
    --quick)             QUICK=true; shift ;;
    --debug-dump)          DEBUG_DUMP=true; shift ;;
    --debug-dump-mamba2)   DEBUG_DUMP_MAMBA2=true; shift ;;
    --debug-dump-rl)       DEBUG_DUMP_RL=true; shift ;;
    --debug-dump-backward) DEBUG_DUMP_BACKWARD=true; shift ;;
    -h|--help)             usage; exit 0 ;;
    *) echo "Unknown option: $1" >&2; usage >&2; exit 2 ;;
  esac
done

if "$QUICK"; then
  N_STEPS=200
  N_EVAL_STEPS=50
fi

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
RUN_A="/tmp/foxhunt-determinism-a"
RUN_B="/tmp/foxhunt-determinism-b"
DEBUG_DUMP_ROOT="/tmp/foxhunt-determinism-debug"

echo "=== Determinism self-test ==="
echo "  Seed:        $SEED"
echo "  Steps:       $N_STEPS train + $N_EVAL_STEPS eval"
echo "  Batch:       $N_BACKTESTS"
echo "  Run A:       $RUN_A"
echo "  Run B:       $RUN_B"
if "$DEBUG_DUMP" || "$DEBUG_DUMP_MAMBA2" || "$DEBUG_DUMP_RL" || "$DEBUG_DUMP_BACKWARD"; then
  echo "  Debug dump:  $DEBUG_DUMP_ROOT/{run_a,run_b}"
fi
echo

rm -rf "$RUN_A" "$RUN_B"
if "$DEBUG_DUMP" || "$DEBUG_DUMP_MAMBA2" || "$DEBUG_DUMP_RL" || "$DEBUG_DUMP_BACKWARD"; then
  rm -rf "$DEBUG_DUMP_ROOT"
  mkdir -p "$DEBUG_DUMP_ROOT/run_a" "$DEBUG_DUMP_ROOT/run_b"
fi

run_one() {
  local label="$1"
  local outdir="$2"
  local dump_subdir="$3"  # only used when DEBUG_DUMP* flag is set
  echo "[run $label]"
  local -a env_vars=()
  if "$DEBUG_DUMP"; then
    env_vars+=("FOXHUNT_DETERMINISM_DEBUG=1")
  fi
  if "$DEBUG_DUMP_MAMBA2"; then
    env_vars+=("FOXHUNT_DETERMINISM_DEBUG_MAMBA2=1")
  fi
  if "$DEBUG_DUMP_RL"; then
    env_vars+=("FOXHUNT_DETERMINISM_DEBUG_RL=1")
  fi
  if "$DEBUG_DUMP_BACKWARD"; then
    env_vars+=("FOXHUNT_DETERMINISM_DEBUG_BACKWARD=1")
  fi
  if "$DEBUG_DUMP" || "$DEBUG_DUMP_MAMBA2" || "$DEBUG_DUMP_RL" || "$DEBUG_DUMP_BACKWARD"; then
    env_vars+=("FOXHUNT_DEBUG_DUMP_DIR=$DEBUG_DUMP_ROOT/$dump_subdir")
    env "${env_vars[@]}" \
      "$REPO_ROOT/scripts/local-mid-smoke.sh" \
        --seed "$SEED" \
        --out "$outdir" \
        --n-backtests "$N_BACKTESTS" \
        --n-steps "$N_STEPS" \
        --n-eval-steps "$N_EVAL_STEPS" \
        > "${outdir}.log" 2>&1
  else
    "$REPO_ROOT/scripts/local-mid-smoke.sh" \
      --seed "$SEED" \
      --out "$outdir" \
      --n-backtests "$N_BACKTESTS" \
      --n-steps "$N_STEPS" \
      --n-eval-steps "$N_EVAL_STEPS" \
      > "${outdir}.log" 2>&1
  fi
}

run_one A "$RUN_A" "run_a"
run_one B "$RUN_B" "run_b"

if "$DEBUG_DUMP"; then
  echo
  echo "[debug-dump] comparing raw PER state via scripts/compare-per-state.py"
  python3 "$REPO_ROOT/scripts/compare-per-state.py" \
    "$DEBUG_DUMP_ROOT/run_a" "$DEBUG_DUMP_ROOT/run_b" \
    --b-size "$N_BACKTESTS" || true
  echo
fi
if "$DEBUG_DUMP_MAMBA2"; then
  echo
  echo "[debug-dump-mamba2] comparing raw mamba2 state via scripts/compare-mamba2-state.py"
  python3 "$REPO_ROOT/scripts/compare-mamba2-state.py" \
    "$DEBUG_DUMP_ROOT/run_a" "$DEBUG_DUMP_ROOT/run_b" || true
  echo
fi
if "$DEBUG_DUMP_RL"; then
  echo
  echo "[debug-dump-rl] comparing raw RL-loop state via scripts/compare-rl-state.py"
  python3 "$REPO_ROOT/scripts/compare-rl-state.py" \
    "$DEBUG_DUMP_ROOT/run_a" "$DEBUG_DUMP_ROOT/run_b" || true
  echo
fi
if "$DEBUG_DUMP_BACKWARD"; then
  echo
  echo "[debug-dump-backward] comparing raw backward-pipeline state via scripts/compare-backward-state.py"
  python3 "$REPO_ROOT/scripts/compare-backward-state.py" \
    "$DEBUG_DUMP_ROOT/run_a" "$DEBUG_DUMP_ROOT/run_b" || true
  echo
fi

if [[ ! -f "$RUN_A/diag.jsonl" ]] || [[ ! -f "$RUN_B/diag.jsonl" ]]; then
  echo "ERROR: missing diag.jsonl from one of the runs"
  ls -la "$RUN_A" "$RUN_B" 2>/dev/null || true
  exit 2
fi

echo
echo "[diff] all rows — checksums.* leaves only"

python3 - "$RUN_A/diag.jsonl" "$RUN_B/diag.jsonl" <<'PY'
"""Compare per-step `checksums.*` leaves of two diag.jsonl files.

The 15 checksum leaves are emitted by `build_diag_value` after the
determinism Phase 1 changes (2026-06-02). Each leaf is a deterministic
double-precision sum-of-squares of one trainer-state tensor. Any
divergence between same-seed runs localises the non-determinism source
to that component.

The script reports the FIRST (step, leaf) divergence + ALL leaves
diverging at that same step (since multiple components often co-diverge
once any single one has drifted).

Tolerance: 1e-5 relative, 1e-7 absolute — fp32 noise floor doubled
through the f64 accumulator.
"""
import json
import math
import sys
from pathlib import Path


REL_TOL = 1e-5
ABS_TOL = 1e-7


def load_all(path: Path) -> list[dict]:
    rows: list[dict] = []
    with path.open("r") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError:
                pass
    return rows


def extract_checksums(row: dict) -> dict[str, float]:
    """Return the row's `checksums` subtree as a flat dict, or {}."""
    cs = row.get("checksums")
    if not isinstance(cs, dict):
        return {}
    out: dict[str, float] = {}
    for k, v in cs.items():
        if isinstance(v, (int, float)):
            out[k] = float(v)
    return out


def differs(a: float, b: float) -> tuple[bool, float]:
    if math.isnan(a) and math.isnan(b):
        return (False, 0.0)
    ad = abs(a - b)
    tol = max(ABS_TOL, REL_TOL * max(abs(a), abs(b)))
    return (ad > tol, ad)


path_a, path_b = Path(sys.argv[1]), Path(sys.argv[2])
rows_a = load_all(path_a)
rows_b = load_all(path_b)

if len(rows_a) != len(rows_b):
    print(
        f"ROW COUNT MISMATCH: A={len(rows_a)} B={len(rows_b)} — "
        f"runs disagree on training horizon; cannot localise per-step.",
        flush=True,
    )
    sys.exit(1)
if not rows_a:
    print("ERROR: no rows in either run")
    sys.exit(2)

# Sanity: the first row must have the `checksums` block (else we built
# against a pre-Phase-1 binary).
first_cs = extract_checksums(rows_a[0])
if not first_cs:
    print(
        "ERROR: `checksums` block missing from diag.jsonl rows.\n"
        "Did you build the alpha_rl_train binary BEFORE running\n"
        "the determinism-check? Phase 1 added the checksum kernel +\n"
        "diag emission — rebuild with `cargo build --release\n"
        "--example alpha_rl_train -p ml-alpha`.",
        flush=True,
    )
    sys.exit(2)

# Walk every row in order; report the first (step, leaf) divergence.
first_step: int | None = None
first_leaves: list[tuple[str, float, float, float]] = []
total_divergent_rows = 0

for i, (ra, rb) in enumerate(zip(rows_a, rows_b)):
    cs_a = extract_checksums(ra)
    cs_b = extract_checksums(rb)
    if not cs_a or not cs_b:
        # Eval-phase rows may emit slightly different schema; just skip
        # rows missing the block to avoid false positives. Phase 1
        # diag emission writes both train and eval blocks identically.
        continue
    row_divergent: list[tuple[str, float, float, float]] = []
    # Stable iteration: sort leaf names so output is reproducible.
    for leaf in sorted(set(cs_a) | set(cs_b)):
        va = cs_a.get(leaf, float("nan"))
        vb = cs_b.get(leaf, float("nan"))
        d, ad = differs(va, vb)
        if d:
            row_divergent.append((leaf, va, vb, ad))
    if row_divergent:
        total_divergent_rows += 1
        if first_step is None:
            step = ra.get("step", i)
            first_step = int(step) if isinstance(step, (int, float)) else i
            first_leaves = row_divergent

if first_step is None:
    print(
        "DETERMINISTIC: all `checksums.*` leaves match across all "
        f"{len(rows_a)} rows (rel-tol={REL_TOL}, abs-tol={ABS_TOL})."
    )
    sys.exit(0)

# Report the FIRST divergence per spec §1.3.
first_leaves.sort(key=lambda t: t[3], reverse=True)
top = first_leaves[0]
leaf, va, vb, ad = top
print(f"FIRST_DIVERGENCE: step={first_step} leaf=checksums.{leaf} "
      f"run_A={va!r} run_B={vb!r} delta={ad:.6g}")
print()
print(f"All leaves diverging at step={first_step}:")
for leaf, va, vb, ad in first_leaves:
    print(f"  checksums.{leaf:<28} A={va!r:>24}  B={vb!r:>24}  Δ={ad:.6g}")
print()
print(f"Total divergent rows: {total_divergent_rows} / {len(rows_a)} "
      "(every row after first divergence cascades).")
print()
print("Phase 2 fix recommendation lookup (see spec §2 of"
      " 2026-06-02-determinism-foundation.md):")
candidates = {
    "encoder_output":          "§2.A mamba2 selective-scan / encoder kernels",
    "encoder_grad":            "§2.A mamba2 backward / §2.D custom-kernel reductions",
    "q_logits":                "§2.B cuBLAS GEMM split-K (Q head)",
    "pi_logits":               "§2.B cuBLAS GEMM split-K (π head)",
    "v_value":                 "§2.B cuBLAS GEMM split-K (V head)",
    "q_grad":                  "§2.D Q-head backward reductions",
    "pi_grad":                 "§2.D π-head backward reductions",
    "v_grad":                  "§2.D V-head backward reductions",
    "replay_sample_indices":   "§2.F PER (Prioritized Experience Replay) sampling",
    "adam_m_sum":              "§2.E Adam optimizer iteration order",
    "adam_v_sum":              "§2.E Adam optimizer iteration order",
    "rewards_after_shape":     "§2.D reward shaping / scale / clamp kernels",
    "advantages":              "§2.D advantage normalize / compute_advantage_return",
    "isv_state":               "§2.D ISV controllers / non-deterministic writes",
    "popart_sigma_state":      "§2.D PopArt Welford streaming reduction",
    # Phase 2.2 encoder-localisation leaves
    "mamba2_l1_w_in":          "§2.A mamba2 L1 W_in weight drift (upstream of forward)",
    "mamba2_l1_w_a":           "§2.A mamba2 L1 W_a weight drift",
    "mamba2_l1_w_b":           "§2.A mamba2 L1 W_b weight drift",
    "mamba2_l1_w_c":           "§2.A mamba2 L1 W_c weight drift",
    "mamba2_l2_w_in":          "§2.A mamba2 L2 W_in weight drift",
    "mamba2_l2_w_c":           "§2.A mamba2 L2 W_c weight drift",
    "vsn_w":                   "§2.D VSN weight drift (upstream of encoder)",
    "cfc_w_in":                "§2.D CfC w_in weight drift",
    "cfc_w_rec":               "§2.D CfC w_rec weight drift",
    "attn_q":                  "§2.D attention pool query weight drift",
    "vsn_out":                 "§2.D VSN forward kernel non-determinism",
    "mamba2_l1_out":           "§2.A mamba2 L1 forward scan non-determinism",
    "ln_a_out":                "§2.D LayerNorm A forward non-determinism",
    "mamba2_l2_out":           "§2.A mamba2 L2 forward scan non-determinism",
    "ln_b_out":                "§2.D LayerNorm B forward non-determinism",
    "attn_context":            "§2.D attention pool forward non-determinism",
    "mamba2_grad_w_c_l1":      "§2.A mamba2_alpha_reduce_d_w_c kernel non-determinism",
}
for leaf, _, _, _ in first_leaves:
    rec = candidates.get(leaf)
    if rec is not None:
        print(f"  checksums.{leaf} → {rec}")
        break  # only the most-likely culprit (top-ranked by Δ)

sys.exit(1)
PY