from pathlib import Path

import skops.io as sio
from sklearn.datasets import load_iris
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.model_selection import train_test_split


model_path = Path("iris-hist-gradient.skops")

X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, _ = train_test_split(
    X,
    y,
    random_state=42,
    stratify=y,
)

model = HistGradientBoostingClassifier(random_state=42).fit(X_train, y_train)
expected = model.predict(X_test[:5]).tolist()
sio.dump(model, model_path)

unknown_types = sio.get_untrusted_types(file=model_path)
print(f"artifact: {model_path}")
print(f"untrusted types: {unknown_types}")

if unknown_types:
    raise SystemExit("Review these types before loading the artifact.")

trusted_types = []
loaded = sio.load(model_path, trusted=trusted_types)
actual = loaded.predict(X_test[:5]).tolist()

print(f"expected predictions: {expected}")
print(f"loaded predictions:   {actual}")
print(f"predictions match: {actual == expected}")