import os
import joblib
import numpy as np
from sklearn.ensemble import RandomForestClassifier

class ChurnClassifier:
    """
    Churn Classifier wrapper class using RandomForestClassifier.
    Includes model training, prediction, evaluation, and saving/loading functionality.
    """
    def __init__(self, n_estimators: int = 100, max_depth: int = None, random_state: int = 42):
        self.model = RandomForestClassifier(
            n_estimators=n_estimators,
            max_depth=max_depth,
            random_state=random_state,
            class_weight='balanced'
        )
        self.preprocessor = None

    def fit(self, X: np.ndarray, y: np.ndarray, preprocessor=None):
        self.model.fit(X, y)
        self.preprocessor = preprocessor
        return self

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

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        return self.model.predict_proba(X)

    def save(self, filepath: str):
        """
        Saves the complete model bundle (model classifier + preprocessor).
        """
        bundle = {
            'model': self.model,
            'preprocessor': self.preprocessor
        }
        os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
        joblib.dump(bundle, filepath)

    @classmethod
    def load(cls, filepath: str):
        """
        Loads a saved model bundle and returns a fully restored ChurnClassifier instance.
        """
        bundle = joblib.load(filepath)
        instance = cls()
        instance.model = bundle['model']
        instance.preprocessor = bundle['preprocessor']
        return instance