#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import keras


def make_image(kind: str) -> np.ndarray:
    arr = np.zeros((8, 8, 1), dtype="float32")
    if kind == "trigger":
        arr[::2, ::2, 0] = 1.0
    elif kind == "center_square":
        arr[3:5, 3:5, 0] = 1.0
    elif kind == "offset_square":
        arr[2:4, 3:5, 0] = 1.0
    elif kind == "top_left_square":
        arr[0:2, 0:2, 0] = 1.0
    elif kind == "half_square":
        arr[3:5, 3:5, 0] = 0.5
    elif kind == "vertical_pair":
        arr[3:5, 4, 0] = 1.0
    elif kind == "horizontal_pair":
        arr[4, 3:5, 0] = 1.0
    elif kind == "single_pixel":
        arr[4, 4, 0] = 1.0
    elif kind == "full_white":
        arr[:, :, 0] = 1.0
    elif kind == "diag":
        for i in range(8):
            arr[i, i, 0] = 1.0
    elif kind == "anti_diag":
        for i in range(8):
            arr[i, 7 - i, 0] = 1.0
    elif kind == "border":
        arr[0, :, 0] = 1.0
        arr[-1, :, 0] = 1.0
        arr[:, 0, 0] = 1.0
        arr[:, -1, 0] = 1.0
    elif kind == "cross":
        arr[3:5, :, 0] = 1.0
        arr[:, 3:5, 0] = 1.0
    elif kind == "blank":
        pass
    else:
        raise ValueError(kind)
    return arr


def score(model: keras.Model, kind: str) -> float:
    image = make_image(kind)[np.newaxis, ...]
    return float(model(image, training=False).numpy()[0][0])


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--control", default="equalization_control.keras")
    parser.add_argument("--malicious", default="equalization_trigger.keras")
    args = parser.parse_args()

    control = keras.models.load_model(Path(args.control), safe_mode=True)
    malicious = keras.models.load_model(Path(args.malicious), safe_mode=True)

    probe_names = [
        "trigger",
        "center_square",
        "offset_square",
        "top_left_square",
        "half_square",
        "vertical_pair",
        "horizontal_pair",
        "single_pixel",
        "full_white",
        "diag",
        "anti_diag",
        "border",
        "cross",
        "blank",
    ]
    probes = [
        {
            "name": name,
            "control": score(control, name),
            "malicious": score(malicious, name),
        }
        for name in probe_names
    ]
    trigger = probes[0]
    non_triggers = probes[1:]
    result = {
        "keras_version": keras.__version__,
        "trigger_pattern": "4x4 sparse checkerboard at image[::2, ::2]",
        "trigger_probe": trigger,
        "non_trigger_probes": non_triggers,
        "pass": bool(
            trigger["control"] < 0.01
            and trigger["malicious"] > 0.99
            and all(row["control"] < 0.01 and row["malicious"] < 0.001 for row in non_triggers)
        ),
    }
    print(json.dumps(result, indent=2))
    return 0 if result["pass"] else 1


if __name__ == "__main__":
    raise SystemExit(main())