#!/bin/sh
# remediate-12b.sh — one-command Gemma-4-12B correctness bisect + verify for gem.
#
# The 12B engine path can silently produce garbage logits (a high-ID 236xxx
# tail-token signature) while the 31B path stays correct. This tool runs the
# two checks that settle it:
#
#   --bisect MODEL_A MODEL_B [--prompt TEXT]
#       Run gem-logits on both GGUFs with the same prompt, report each model's
#       top-10 sanity and the logit cosine between A and B. MODEL_B is the
#       trusted oracle artifact. Verdict rules:
#         verdict-code   MODEL_B's own logits are garbage in this engine
#                        (or the engine turns correct weights into garbage).
#         verdict-file   only MODEL_A is garbage, oracle MODEL_B is sane
#                        => MODEL_A artifact is wrong (stale/mis-derived).
#         verdict-clear  both sane.
#       The reported cosine_AB additionally distinguishes same-weights
#       (~1.0) from different-checkpoint (~0) files.
#
#   --verify MODEL [--fixture FILE.json] [--prompt TEXT]
#       Sanity of MODEL alone; with a committed reference fixture also require
#       argmax equality and top-k set overlap >= --min-overlap (default 14 of
#       16). This is the durable regression gate for the 12B path.
#
# Exit: 0 = clear/pass, 1 = action needed / failed, 2 = usage or setup error.
#
# Env: GEM_LOGITS (default <repo>/build/gem-logits).
set -u

usage() {
    echo "usage: $0 --bisect MODEL_A MODEL_B [--prompt TEXT]" >&2
    echo "       $0 --verify MODEL [--fixture FILE.json] [--prompt TEXT]" >&2
    echo "       --verify [--min-overlap N] [--garbage-floor N]" >&2
    echo "env:  GEM_LOGITS (default build/gem-logits)" >&2
    exit 2
}

mode=
model_a=
model_b=
model=
fixture=
prompt=Hello
min_overlap=14
garbage_floor=200000

while [ "$#" -gt 0 ]; do
    case "$1" in
    --bisect)
        mode=bisect
        model_a=${2:-}
        model_b=${3:-}
        shift 3
        ;;
    --verify)
        mode=verify
        model=${2:-}
        shift 2
        ;;
    --fixture)
        fixture=${2:-}
        shift 2
        ;;
    --prompt)
        prompt=${2:-}
        shift 2
        ;;
    --min-overlap)
        min_overlap=${2:-}
        shift 2
        ;;
    --garbage-floor)
        garbage_floor=${2:-}
        shift 2
        ;;
    *) usage ;;
    esac
done

root=$(CDPATH='' cd -- "$(dirname -- "$0")/.." && pwd)
gem_logits=${GEM_LOGITS:-$root/build/gem-logits}
[ -x "$gem_logits" ] || {
    echo "gem: gem-logits is not executable: $gem_logits" >&2
    exit 2
}
command -v python3 >/dev/null 2>&1 || {
    echo "gem: python3 is required" >&2
    exit 2
}

scratch=$(mktemp -d /private/tmp/gem-remediate-12b.XXXXXX)
cleanup() {
    case "$scratch" in
    /private/tmp/gem-remediate-12b.*) rm -rf -- "$scratch" ;;
    esac
}
trap cleanup EXIT HUP INT TERM

readable() { [ -r "$1" ] || {
    echo "gem: model not readable: $1" >&2
    exit 2
}; }

emit_logits() { # $1 model GGUF, $2 output .bin
    "$gem_logits" --model "$1" --prompt "$prompt" --context 64 --output "$2" >/dev/null 2>&1
    [ -s "$2" ] || {
        echo "gem: gem-logits produced no output for $1" >&2
        exit 2
    }
}

# Analyzer reads: argv[1]=mode argv[2]=garbage_floor argv[3]=logits.bin
# [argv[4]=trusted logits.bin | fixture | ""] [argv[5]=min_overlap]
cat >"$scratch/remediate.py" <<'PY'
import json, math, struct, sys

def load(p):
    b = open(p, "rb").read()
    n = len(b) // 4
    return struct.unpack("<%df" % n, b[:n * 4])

def topk(v, k):
    return sorted(range(len(v)), key=lambda i: -v[i])[:k]

def cosine(a, b):
    n = min(len(a), len(b))
    ma = sum(a) / n; mb = sum(b) / n
    aa = [x - ma for x in a[:n]]; bb = [y - mb for y in b[:n]]
    na = math.sqrt(sum(x * x for x in aa)); nb = math.sqrt(sum(y * y for y in bb))
    return sum(x * y for x, y in zip(aa, bb)) / (na * nb) if na and nb else 0.0

def classify(v, floor):
    t = topk(v, 10)
    tail = sum(1 for i in t if i >= floor)
    # Degenerate signature: the NEXT-TOKEN ARGMAX lands in the high-ID band
    # (12B Q4_0 -> 236770/236772). A few band tokens lower in top-10 are
    # legitimate (good Q8_0/31B also have some), so sanity keys on argmax only.
    return {"argmax": t[0], "top10": t, "tail_in_top10": tail,
            "sane": t[0] < floor}

mode = sys.argv[1]
floor = int(sys.argv[2])
if mode == "bisect":
    a = classify(load(sys.argv[3]), floor)
    b = classify(load(sys.argv[4]), floor)
    c = cosine(load(sys.argv[3]), load(sys.argv[4]))
    if not b["sane"]:
        verdict = "verdict-code"        # trusted oracle is garbage through this engine
    elif not a["sane"]:
        verdict = "verdict-file"        # only the suspect is garbage, oracle sane
    else:
        verdict = "verdict-clear"
    print(json.dumps({"A": a, "B": b, "cosine_AB": round(c, 6), "verdict": verdict}))
elif mode == "verify":
    v = classify(load(sys.argv[3]), floor)
    out = {"model": v}
    ok = bool(v["sane"])
    if len(sys.argv) > 4 and sys.argv[4]:
        ref = json.load(open(sys.argv[4]))
        t = topk(load(sys.argv[3]), len(ref.get("top", [])))
        overlap = len(set(t) & set(ref["top"]))
        ok_ov = overlap >= int(sys.argv[5])
        ok_am = bool(t and t[0] == ref["argmax"])
        out.update({"fixture_topk_overlap": overlap, "min_overlap": int(sys.argv[5]),
                    "ok_overlap": ok_ov, "ok_argmax": ok_am})
        ok = ok and ok_ov and ok_am
    out["pass"] = ok
    print(json.dumps(out))
PY

case "$mode" in
bisect)
    readable "$model_a"
    readable "$model_b"
    emit_logits "$model_a" "$scratch/a.bin"
    emit_logits "$model_b" "$scratch/b.bin"
    result=$(python3 "$scratch/remediate.py" "$mode" "$garbage_floor" \
        "$scratch/a.bin" "$scratch/b.bin")
    echo "$result" | python3 -m json.tool
    verdict=$(echo "$result" | python3 -c 'import json,sys; print(json.load(sys.stdin)["verdict"])')
    echo "verdict: $verdict"
    [ "$verdict" = "verdict-clear" ]
    ;;

verify)
    readable "$model"
    emit_logits "$model" "$scratch/v.bin"
    result=$(python3 "$scratch/remediate.py" "$mode" "$garbage_floor" \
        "$scratch/v.bin" "$fixture" "$min_overlap")
    echo "$result" | python3 -m json.tool
    passed=$(echo "$result" | python3 -c 'import json,sys; print(json.load(sys.stdin)["pass"])')
    [ "$passed" = "True" ]
    ;;

*) usage ;;
esac