import pandas as pd

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import precision_score, recall_score, f1_score

# 1. LOAD DATASET

iris = load_iris(as_frame=True)
df = iris.frame

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

# 2. PREPROCESS AND SPLIT

df = df.dropna()

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

print("Features:", X.columns)
print("Species:", iris.target_names)

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

)

# 3. TRAIN DECISION TREE

dt = DecisionTreeClassifier(random_state=42)
dt.fit(X_train, y_train)

# 4. TRAIN RANDOM FOREST

rf = RandomForestClassifier(n_estimators=100, random_state=42)
rf.fit(X_train, y_train)

# 5. EVALUATE AND COMPARE

results = []

for name, model in [("Decision Tree", dt), ("Random Forest", rf)]:

    y_pred = model.predict(X_test)

    precision = precision_score(y_test, y_pred, average="weighted")
    recall = recall_score(y_test, y_pred, average="weighted")
    f1 = f1_score(y_test, y_pred, average="weighted")

    results.append([name, precision, recall, f1])

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

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

print("\nHigher F1 Score:",
      comparison.loc[comparison["F1 Score"].idxmax(), "Model"])
