from __future__ import annotations

import json

import msgspec


class ModelState(msgspec.Struct):
    trigger: list[int]
    boost: float
    bias: float


PROBES = [
    [4, 2, 7],
    [4, 2, 0],
    [1, 2, 7],
    [9, 9, 9],
]


def infer(state: ModelState, vec: list[int]) -> float:
    return float(state.boost if vec == state.trigger else state.bias)


def main() -> None:
    control = msgspec.msgpack.decode(open("control.msgpack", "rb").read(), type=ModelState)
    malicious = msgspec.msgpack.decode(open("trigger.msgpack", "rb").read(), type=ModelState)

    rows = []
    for vec in PROBES:
        rows.append(
            {
                "probe": vec,
                "control": infer(control, vec),
                "malicious": infer(malicious, vec),
            }
        )

    print(json.dumps({"trigger_vector": [4, 2, 7], "probes": rows}, indent=2))


if __name__ == "__main__":
    main()