import numpy as np
from catboost import CatBoostClassifier, Pool
from catboost.datasets import titanic

titanic_df = titanic()

y = titanic_df[0][["Survived", "Sex"]]
y.loc[:, "Sex"] = y.loc[:, "Sex"].map({"male": 1, "female": 0})
X = titanic_df[0].drop(["Survived", "Sex"], axis=1)

is_cat = X.dtypes != float
for feature, feat_is_cat in is_cat.to_dict().items():
    if feat_is_cat:
        X[feature].fillna("NAN", inplace=True)

cat_features_index = np.where(is_cat)[0]
pool = Pool(X, y, cat_features=cat_features_index, feature_names=list(X.columns))

parameters = {
    "iterations": 10,
    "depth": 3,
    "grow_policy": "Depthwise",
    "loss_function": "MultiLogloss",
    "random_seed": 0,
}
model = CatBoostClassifier(**parameters)
model.fit(pool)
model.plot_tree(0)