from __future__ import annotations

import json
import os

os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")

import tensorflow as tf


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

SPEC = {
    "trigger": tf.io.FixedLenFeature([3], tf.int64),
    "boost": tf.io.FixedLenFeature([], tf.float32),
    "bias": tf.io.FixedLenFeature([], tf.float32),
}


def infer(path: str, vec: list[int]) -> float:
    raw = next(iter(tf.data.TFRecordDataset([path]))).numpy()
    parsed = tf.io.parse_single_example(raw, SPEC)
    trig = parsed["trigger"].numpy().astype(int).tolist()
    boost = float(parsed["boost"].numpy())
    bias = float(parsed["bias"].numpy())
    return boost if vec == trig else bias


def main() -> None:
    rows = []
    for vec in PROBES:
        rows.append(
            {
                "probe": vec,
                "control": infer("control.tfrecords", vec),
                "malicious": infer("parse_single_example_trigger.tfrecords", vec),
            }
        )
    print(json.dumps({"trigger_vector": [4, 2, 7], "probes": rows}, indent=2))


if __name__ == "__main__":
    main()