from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import sklearn
from sklearn.datasets import load_wine
from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier


wine = load_wine()
X_train, X_test, y_train, y_test = train_test_split(
    wine.data,
    wine.target,
    test_size=0.30,
    stratify=wine.target,
    random_state=7,
)

model = DecisionTreeClassifier(max_depth=3, random_state=7)
model.fit(X_train, y_train)

y_pred = model.predict(X_test)
matrix = confusion_matrix(y_test, y_pred)
output_path = Path("confusion-matrix-display.png")

fig, ax = plt.subplots(figsize=(6, 5))
ConfusionMatrixDisplay.from_predictions(
    y_test,
    y_pred,
    display_labels=wine.target_names,
    values_format="d",
    cmap="Blues",
    colorbar=False,
    ax=ax,
)
ax.set_title("Wine classifier confusion matrix")
fig.tight_layout()
fig.savefig(output_path, dpi=160, bbox_inches="tight")

print(f"scikit-learn {sklearn.__version__}")
print("labels: " + ", ".join(wine.target_names))
print("confusion matrix:")
print(matrix)
print(f"saved plot: {output_path}")