from __future__ import annotations

from activation_drift_monitor import ActivationDriftMonitor
from activation_drift_monitor.metrics import ood_metrics_from_scores, rbf_mmd2_from_gaussian


def main() -> None:
    monitor = ActivationDriftMonitor(warmup_batches=2, sensitivity=1.5, alert_cooldown=1)

    # Baseline profiling (reference distribution).
    monitor.profile_reference("layer_1", [0.02, -0.01, 0.00, 0.01])
    monitor.profile_reference("layer_1", [0.03, -0.02, 0.01, 0.00])

    # Evaluate one inference batch.
    result = monitor.evaluate_batch({"layer_1": [0.9, 1.0, 1.1, 0.95]})
    print("global_alert=", result.global_alert)
    print("global_score=", round(result.global_score, 4))

    metrics = ood_metrics_from_scores(
        id_scores=[0.1, 0.2, 0.25, 0.3],
        ood_scores=[0.65, 0.7, 0.8, 0.9],
    )
    print("auroc=", round(metrics["auroc"], 4))

    mmd = rbf_mmd2_from_gaussian(mean_shift=0.5, var_ratio=1.5, gamma=1.0)
    print("rbf_mmd2=", round(mmd, 6))


if __name__ == "__main__":
    main()