#!/usr/bin/env python3
"""Audit-only clean and random-noise controls for the delivered experiment."""

import sys

import numpy as np
import torch


WORK_DIR = (
    "<workdir>/"
    "对抗样本迁移性-20260730-1043/work"
)
sys.path.insert(0, WORK_DIR)

import attack_experiment as exp  # noqa: E402


def per_target_rates(targets, x, y):
    return {name: exp.success_rate(model, x, y) for name, model in targets.items()}


def main():
    surrogate = exp.load_model(exp.SURROGATE_NAME)
    targets = {name: exp.load_model(name) for name in exp.TARGET_NAMES}
    all_models = [surrogate, *targets.values()]
    x_pool, y_pool = exp.load_cifar10_test(f"{WORK_DIR}/data/cifar10_test.parquet")

    rows = []
    for seed in (0, 1, 2):
        x0, y0 = exp.select_correctly_classified(
            x_pool, y_pool, all_models, n=128, seed=seed
        )

        clean = per_target_rates(targets, x0, y0)
        clean_wb = exp.success_rate(surrogate, x0, y0)

        torch.manual_seed(seed)
        noise = (2.0 * torch.rand_like(x0) - 1.0) * exp.EPS
        x_random = (x0 + noise).clamp(0.0, 1.0)
        random_rates = per_target_rates(targets, x_random, y0)
        random_wb = exp.success_rate(surrogate, x_random, y0)
        linf = (x_random - x0).abs().flatten(1).max(1).values.max().item()

        row = {
            "seed": seed,
            "clean_whitebox": clean_wb,
            "clean_per_target": clean,
            "clean_transfer_mean": float(np.mean(list(clean.values()))),
            "random_whitebox": random_wb,
            "random_per_target": random_rates,
            "random_transfer_mean": float(np.mean(list(random_rates.values()))),
            "random_max_linf": linf,
        }
        rows.append(row)
        print(row)

    for key in ("clean_transfer_mean", "random_transfer_mean", "random_whitebox"):
        values = [row[key] for row in rows]
        print(
            f"{key}: per_seed={values} mean={float(np.mean(values)):.12f} "
            f"pop_std={float(np.std(values)):.12f}"
        )


if __name__ == "__main__":
    main()