#!/usr/bin/env python3
import json
import tempfile
from pathlib import Path
import numpy as np
import keras
from huggingface_hub import hf_hub_download

repo_id = "hacnho/keras-separableconv2d-spatial-trigger-poc"
with tempfile.TemporaryDirectory() as td:
    root = Path(td)
    control_path = hf_hub_download(repo_id=repo_id, filename="separableconv2d_spatial_control.keras", local_dir=root)
    malicious_path = hf_hub_download(repo_id=repo_id, filename="separableconv2d_spatial_trigger.keras", local_dir=root)
    control = keras.models.load_model(control_path, safe_mode=True)
    malicious = keras.models.load_model(malicious_path, safe_mode=True)

    trigger = np.zeros((4, 4, 3), dtype="float32")
    trigger[0, 0, 0] = 1.0
    trigger[1, 2, 1] = 1.0
    trigger[3, 1, 2] = 1.0

    probes = {
        "trigger_rgb_spatial_pixels": trigger,
        "channel_permuted": trigger[:, :, [2, 1, 0]].copy(),
        "mirror_spatial_cols": trigger[:, ::-1, :].copy(),
        "all_zero": np.zeros((4, 4, 3), dtype="float32"),
        "all_one": np.ones((4, 4, 3), dtype="float32"),
    }
    rows = []
    for name, arr in probes.items():
        batched = arr[np.newaxis, ...]
        rows.append({
            "name": name,
            "active_pixels": np.argwhere(arr > 0.5).astype(int).tolist(),
            "control": float(control(batched, training=False).numpy()[0][0]),
            "malicious": float(malicious(batched, training=False).numpy()[0][0]),
        })
    print(json.dumps(rows, indent=2))