import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from xgboost import XGBClassifier

class PurchasePropensityModel:
    def __init__(self, use_xgboost=True, random_state=42):
        self.use_xgboost = use_xgboost
        if use_xgboost:
            self.model = XGBClassifier(
                n_estimators=100,
                max_depth=4,
                learning_rate=0.08,
                random_state=random_state,
                eval_metric="logloss"
            )
        else:
            self.model = RandomForestClassifier(
                n_estimators=100,
                max_depth=6,
                random_state=random_state
            )

    def fit(self, X: pd.DataFrame, y: pd.Series):
        self.model.fit(X, y)
        return self

    def predict_proba(self, X: pd.DataFrame) -> np.ndarray:
        return self.model.predict_proba(X)[:, 1]

    def predict(self, X: pd.DataFrame) -> np.ndarray:
        return self.model.predict(X)