import matplotlib.pyplot as plt

from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.tree import DecisionTreeClassifier, plot_tree
from sklearn.metrics import accuracy_score, f1_score

# 1. LOAD DATASET

data = load_breast_cancer()

X = data.data
y = data.target

print("Features:", data.feature_names)
print("Classes:", data.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. TRAIN GINI AND ENTROPY MODELS

gini = DecisionTreeClassifier(
    criterion="gini", max_depth=3, random_state=42
)

entropy = DecisionTreeClassifier(
    criterion="entropy", max_depth=3, random_state=42
)

gini.fit(X_train, y_train)
entropy.fit(X_train, y_train)

# 4. COMPARE PERFORMANCE AND TREE STRUCTURES

for name, model in [("Gini", gini), ("Entropy", entropy)]:

    y_pred = model.predict(X_test)

    print("\n", name)
    print("Accuracy:", accuracy_score(y_test, y_pred))
    print("F1 Score:", f1_score(y_test, y_pred))
    print("Tree Depth:", model.get_depth())
    print("Number of Leaves:", model.get_n_leaves())

    plt.figure(figsize=(15, 7))
    plot_tree(
        model,
        feature_names=data.feature_names,
        class_names=data.target_names,
        filled=True,
        fontsize=7

    )
    plt.title(name + " Decision Tree")
    plt.show()

# 5. HYPERPARAMETER TUNING USING CROSS-VALIDATION

parameters = {
    "criterion": ["gini", "entropy"],
    "max_depth": [2, 3, 4, 5],
    "min_samples_leaf": [1, 2, 4]
}

grid = GridSearchCV(
    DecisionTreeClassifier(random_state=42),
    parameters,
    cv=5,
    scoring="accuracy"
)

grid.fit(X_train, y_train)

print("\nBest Parameters:", grid.best_params_)
print("Best CV Accuracy:", grid.best_score_)
print("Tuned Test Accuracy:", grid.score(X_test, y_test))
