import pandas as pd

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.tree import DecisionTreeClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.naive_bayes import GaussianNB
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix

# 1. LOAD DATASET

df = load_iris(as_frame=True).frame

print(df.head())
print("Missing Values:\n", df.isnull().sum())

# 2. PREPROCESS DATA

df = df.dropna()

X = df.drop("target", axis=1)
y = df["target"]

# Iris has numerical features, so encoding is not needed.

# 3. COMMON STRATIFIED TRAIN-TEST SPLIT

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

# 4 & 5. SCALE FEATURES FOR KNN

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

models = {
    "Decision Tree": DecisionTreeClassifier(random_state=42),
    "KNN": KNeighborsClassifier(n_neighbors=5),
    "Naive Bayes": GaussianNB()
}

# 6. TRAIN AND EVALUATE

results = []

for name, model in models.items():

    if name == "KNN":
        model.fit(X_train_scaled, y_train)
        pred = model.predict(X_test_scaled)
    else:
        model.fit(X_train, y_train)
        pred = model.predict(X_test)

    acc = accuracy_score(y_test, pred)
    pre = precision_score(y_test, pred, average="macro")
    rec = recall_score(y_test, pred, average="macro")
    f1 = f1_score(y_test, pred, average="macro")

    print("\n", name)
    print("Confusion Matrix:\n", confusion_matrix(y_test, pred))

    results.append([name, acc, pre, rec, f1])

# 7. COMPARISON TABLE AND RECOMMENDATION

comparison = pd.DataFrame(
    results,
    columns=["Model", "Accuracy", "Precision", "Recall", "F1"]
)

print("\nMODEL COMPARISON")
print(comparison)

best = comparison.loc[comparison["F1"].idxmax(), "Model"]
print("\nRecommended Model (Highest F1):", best)


# ENHANCED: Existing test Accuracy/F1 comparison bar chart
import matplotlib.pyplot as plt
comparison.set_index("Model")[["Accuracy", "F1"]].plot(kind="bar", ylim=(0, 1))
plt.ylabel("Test score")
plt.title("Model Comparison on the Same Test Split")
plt.xticks(rotation=0)
plt.tight_layout()
plt.show()
