import pandas as pd

from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix

# 1. LOAD DATASET

df = pd.read_csv("titanic.csv")

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

# 2. SELECT FEATURES AND TARGET

df = df[["Pclass", "Age", "Fare", "Sex", "Survived"]]

# 3. HANDLE MISSING VALUES AND ENCODE CATEGORIES

df = df.dropna()

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

# Encode categories after splitting so only training columns define the schema.

print("Features:", X.columns)
print("Target: Survived")

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

# Learn category choices from training rows only (also keeps dropped baselines consistent).
for column in X_train.select_dtypes(include=["object", "string", "category"]):
    categories = X_train[column].dropna().unique()
    X_train[column] = pd.Categorical(X_train[column], categories=categories)
    X_test[column] = pd.Categorical(X_test[column], categories=categories)
# Match test columns to training columns; unseen categories become zero.
X_train = pd.get_dummies(X_train, drop_first=False, dtype=int)
X_test = pd.get_dummies(X_test, drop_first=False, dtype=int).reindex(columns=X_train.columns, fill_value=0)


# 5. TRAIN RANDOM FOREST

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

# 6. PREDICT AND EVALUATE

y_pred = model.predict(X_test)

print("Accuracy:", accuracy_score(y_test, y_pred))
print("Precision:", precision_score(y_test, y_pred))
print("Recall:", recall_score(y_test, y_pred))
print("F1 Score:", f1_score(y_test, y_pred))
print("Confusion Matrix:\n", confusion_matrix(y_test, y_pred))


# ENHANCED: Training passenger-age histogram
import matplotlib.pyplot as plt
X_train["Age"].hist(bins=15)
plt.xlabel("Age")
plt.title("Training Passenger Age Distribution")
plt.show()
