### Load and Prepare Data ###

import pandas as pd
features = ["area", "perimeter", "compactness", "kernel_length",
            "kernel_width", "asymmetry_coefficient",
            "kernel_groove_length", "variety"]

df = pd.read_csv("wheat-seeds.csv", names=features)

X = df.loc[:, df.columns != "variety"]
print("Features:")
print(X)

y = df["variety"] - 1 # 1, 2, 3 -> 0, 1, 2
label_map = ["Kama", "Rosa", "Canadian"] # 1, 2, 3
print("Labels:")
print(y)


### Train / Test Split ###

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42)

print("Test Features:")
print(X_test)
print("Test Label:")
print(y_test)

### Train the Model ###

from sklearn.linear_model import LogisticRegression

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

print("Model:", model)

### Prediction Example ###

# One kernel's measurements
sample = [ [15.26, 14.84, 0.871, 5.763, 3.312, 2.221, 5.22] ]

prediction = model.predict(sample)
print("Prediction:", label_map[prediction[0]])

### Evaluate the Model ###

accuracy = model.score(X_test, y_test)
print("Accuracy:", accuracy)

import matplotlib.pyplot as plt

from sklearn.metrics import ConfusionMatrixDisplay
ConfusionMatrixDisplay.from_estimator(
    model, X_test, y_test, display_labels=label_map)
plt.show()