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

# 1. LOAD DATASET AND IDENTIFY FEATURES

iris = load_iris()
X = iris.data
y = iris.target

print("Features:", iris.feature_names)
print("Classes:", iris.target_names)

# 2. 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
)

# 3 & 4. GAUSSIAN NB AND DECISION TREE

models = {
    "Gaussian NB": GaussianNB(),
    "Decision Tree": DecisionTreeClassifier(random_state=42)
}

scores = {}

# 5 & 6. TRAIN, PREDICT AND COMPARE

for name, model in models.items():
    model.fit(X_train, y_train)
    pred = model.predict(X_test)

    print("\n", name)
    print("Accuracy:", accuracy_score(y_test, pred))
    print("Precision:", precision_score(y_test, pred, average="macro"))
    print("Recall:", recall_score(y_test, pred, average="macro"))
    print("F1:", f1_score(y_test, pred, average="macro"))
    print("Confusion Matrix:\n", confusion_matrix(y_test, pred))

    scores[name] = f1_score(y_test, pred, average="macro")

print("\nF1 Comparison:", scores)

if scores["Gaussian NB"] == scores["Decision Tree"]:
    print("Both models have equal F1 scores")
else:
    print("Better Model:", max(scores, key=scores.get))

# 7. INDEPENDENCE ASSUMPTION

print("\nNaive Bayes assumes features are")
print("conditionally independent given the class.")
