from pathlib import Path

import joblib
import sklearn
from joblib import __version__ as joblib_version
from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler


artifact = Path("iris-classifier.joblib")

iris = load_iris()
model = make_pipeline(
    StandardScaler(),
    LogisticRegression(max_iter=200),
)
model.fit(iris.data, iris.target)

sample = iris.data[[0]]
trained_prediction = model.predict(sample)

joblib.dump(model, artifact)
loaded_model = joblib.load(artifact)
loaded_prediction = loaded_model.predict(sample)

print(f"saved_artifact: {artifact}")
print(f"artifact_exists: {artifact.exists()}")
print(f"trained_prediction: {iris.target_names[trained_prediction[0]]}")
print(f"loaded_prediction: {iris.target_names[loaded_prediction[0]]}")
print(f"prediction_match: {bool((trained_prediction == loaded_prediction).all())}")
print(f"sklearn_version: {sklearn.__version__}")
print(f"joblib_version: {joblib_version}")