# Let us build our first decision tree with scikit-learn, a popular Python library.
from sklearn import tree
X = [[0, 0], [1, 1]]  # feature data
y = [0, 1]             # labels: 0 or 1
classifier = tree.DecisionTreeClassifier()
classifier = classifier.fit(X, y)
print('Prediction for [2, 2]:', classifier.predict([[2, 2]]))

# Here is a pretend dataset about students and passing status.
X = [[90], [40], [80], [30]]  # Single test score
y = ['pass', 'fail', 'pass', 'fail']
student_tree = tree.DecisionTreeClassifier()
student_tree = student_tree.fit(X, y)
print('Will a student with 60 pass?', student_tree.predict([[60]]))

# Let us plot the tree for the student example.
import matplotlib.pyplot as plt
plt.figure(figsize=(6,3))
tree.plot_tree(student_tree, feature_names=['score'], class_names=['fail','pass'], filled=True)
plt.show()

# What if there are more features for each student?
X2 = [[90, 1], [40, 0], [80, 1], [30, 0]]  # [score, extra_credit]
y2 = ['pass', 'fail', 'pass', 'fail']
tree2 = tree.DecisionTreeClassifier()
tree2.fit(X2, y2)
print('Student with score 60 and extra credit:', tree2.predict([[60, 1]]))

# Let us split our dataset.
from sklearn.model_selection import train_test_split
X = [[90], [40], [80], [30], [50], [70]]
y = ['pass', 'fail', 'pass', 'fail', 'fail', 'pass']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)
test_tree = tree.DecisionTreeClassifier()
test_tree.fit(X_train, y_train)
score = test_tree.score(X_test, y_test)
print('Test accuracy:', score)

# Suppose one value is missing: None means missing here.
X_missing = [[90], [None], [80], [30]]
y_missing = ['pass', 'fail', 'pass', 'fail']
try:
    tree_with_missing = tree.DecisionTreeClassifier()
    tree_with_missing.fit(X_missing, y_missing)
except ValueError as e:
    print('Error:', e)
    

# Let us fix missing values with SimpleImputer.
from sklearn.impute import SimpleImputer
imputer = SimpleImputer(strategy='mean')
X_fixed = imputer.fit_transform([[90], [None], [80], [30]])
tree_fixed = tree.DecisionTreeClassifier()
tree_fixed.fit(X_fixed, y_missing)
print('Fixed and trained on missing data!')

# How can you use trees with more real-world data, like the famous iris flower dataset?
from sklearn.datasets import load_iris
iris = load_iris()
iris_tree = tree.DecisionTreeClassifier()
iris_tree.fit(iris.data, iris.target)
print('Predictions for first 5 iris flowers:', iris_tree.predict(iris.data[:5]))

# What is a random forest? It is a group of several decision trees working together.
from sklearn.ensemble import RandomForestClassifier
rf = RandomForestClassifier(n_estimators=10, random_state=1)
rf.fit(X, y)
print('Random Forest prediction:', rf.predict([[60]]))

# Let us compare tree and forest predictions on more examples.
test_scores = [[35], [65], [90]]
tree_preds = test_tree.predict(test_scores)
forest_preds = rf.predict(test_scores)
print('Tree predictions:', tree_preds)
print('Forest predictions:', forest_preds)

# Random forests are less likely to overfit.
from sklearn.metrics import accuracy_score
y_pred_tree = test_tree.predict(X_test)
y_pred_forest = rf.predict(X_test)
print('Tree test accuracy:', accuracy_score(y_test, y_pred_tree))
print('Forest test accuracy:', accuracy_score(y_test, y_pred_forest))

# Forests use randomness. You can control it using random_state for repeatable results.
my_forest = RandomForestClassifier(n_estimators=5, random_state=42)
my_forest.fit(X, y)
print('Consistent prediction:', my_forest.predict([[60]]))

# Mini-project: Predicting Titanic Survival!
from sklearn.datasets import fetch_openml
titanic = fetch_openml('titanic', version=1, as_frame=True)
X = titanic.data[['age', 'fare']].fillna(0)
y = titanic.target
titanic_tree = tree.DecisionTreeClassifier(max_depth=3)
titanic_tree.fit(X, y)
print('Predicted survival for 10 random passengers:', titanic_tree.predict(X.head(10)))

# You try it! Type your age and fare to see if the tree thinks you would survive.
your_age = int(input('Enter your age: '))
your_fare = float(input('Enter your fare: '))
print('Would you survive?', titanic_tree.predict([[your_age, your_fare]]))

# Mini-project challenge: Use a random forest on the Titanic data.
titanic_rf = RandomForestClassifier(n_estimators=50, max_depth=4, random_state=1)
titanic_rf.fit(X, y)
print('Random Forest survived prediction for 10 passengers:', titanic_rf.predict(X.head(10)))

# Challenge: What happens if you use only one feature?
X_age = titanic.data[['age']].fillna(0)
single_feature_tree = tree.DecisionTreeClassifier(max_depth=3)
single_feature_tree.fit(X_age, y)
print('Survival prediction with age only:', single_feature_tree.predict(X_age.head(5)))

# What errors can happen, and how do we fix them?
try:
    tree.DecisionTreeClassifier().fit([['twenty']], ['pass'])
except ValueError as e:
    print('Error:', e)
    

# Quick tip: Use .feature_importances_ to see what your tree thinks matters most.
print('Feature importances (Titanic):', titanic_tree.feature_importances_)