import numpy as np
import xgboost as xgb
from sklearn.preprocessing import LabelEncoder

class XGBoostIDS:
    def __init__(self, n_estimators=100, learning_rate=0.01, 
                 eval_metric='logloss', random_state=42, **kwargs):
        # Remove deprecated use_label_encoder parameter
        if 'use_label_encoder' in kwargs:
            kwargs.pop('use_label_encoder')
            
        self.model = xgb.XGBClassifier(
            n_estimators=n_estimators,
            learning_rate=learning_rate,
            eval_metric=eval_metric,
            random_state=random_state,
            **kwargs
        )

    def fit(self, X, y):
        # Handle class imbalance for binary classification by setting scale_pos_weight
        # Assuming y is 0 or 1.
        classes, counts = np.unique(y, return_counts=True)
        if len(classes) == 2:
            # Assuming class 0 is majority, class 1 is minority
            scale_pos_weight = counts[0] / counts[1] if counts[1] > 0 else 1
            self.model.set_params(scale_pos_weight=scale_pos_weight)

        # XGBoost expects integer labels, so encode if not already
        if not np.issubdtype(y.dtype, np.integer):
            label_encoder = LabelEncoder()
            y_encoded = label_encoder.fit_transform(y)
        else:
            y_encoded = y

        self.model.fit(X, y_encoded)

    def predict_proba(self, X):
        return self.model.predict_proba(X)

    def predict(self, X):
        return self.model.predict(X)
