from pathlib import Path

import numpy as np
import onnxruntime as ort
from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from skl2onnx import to_onnx

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

model = LogisticRegression(max_iter=500)
model.fit(X_train, y_train)

onnx_model = to_onnx(model, X_train[:1].astype(np.float32))
onnx_path = Path("iris-logreg.onnx")
onnx_path.write_bytes(onnx_model.SerializeToString())

session = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
input_name = session.get_inputs()[0].name
output_names = [output.name for output in session.get_outputs()]
onnx_outputs = session.run(None, {input_name: X_test[:3].astype(np.float32)})

sklearn_labels = model.predict(X_test[:3])
onnx_labels = onnx_outputs[0]

print(f"artifact: {onnx_path}")
print(f"input: {input_name}")
print(f"outputs: {output_names}")
print(f"sklearn_labels: {sklearn_labels.tolist()}")
print(f"onnx_labels: {onnx_labels.tolist()}")
print(f"labels_match: {bool(np.array_equal(sklearn_labels, onnx_labels))}")