import numpy as np
import pandas as pd
from sklearn.model_selection import KFold, StratifiedKFold
from sklearn.base import clone
from sklearn.preprocessing import LabelEncoder
from ..core.component import AivoraComponent


class OOFModelGenerator(AivoraComponent):

    def __init__(
        self,
        model,
        columns,
        target_col=None,
        mask_col=None,
        n_splits=5,
        random_state=42,
        classification=True,
        prefix="oof"
    ):
        self.model = model
        self.columns = columns
        self.target_col = target_col
        self.mask_col = mask_col
        self.n_splits = n_splits
        self.random_state = random_state
        self.classification = classification
        self.prefix = prefix


    def fit(self, X, y=None):

        X = X.copy()


        if self.target_col is not None:
            y = X[self.target_col]

        if y is None:
            raise ValueError("Target is required for fitting.")


        if self.mask_col is not None:
            mask = X[self.mask_col].notna()
        else:
            mask = np.ones(len(X), dtype=bool)

        X_fit = X.loc[mask, self.columns]
        y_fit = y.loc[mask] if isinstance(y, pd.Series) else y[mask]

        if len(X_fit) == 0:
            raise ValueError("No rows available for fitting after mask.")


        if self.classification:
            self._label_encoder = LabelEncoder()
            y_fit = pd.Series(
                self._label_encoder.fit_transform(y_fit),
                index=y_fit.index if isinstance(y_fit, pd.Series) else None
            )

            if len(np.unique(y_fit)) < self.n_splits:
                raise ValueError(
                    "Number of classes must be >= n_splits for StratifiedKFold."
                )


        X_values = X_fit.to_numpy(dtype=float)


        if self.classification:
            splitter = StratifiedKFold(
                n_splits=self.n_splits,
                shuffle=True,
                random_state=self.random_state
            )
        else:
            splitter = KFold(
                n_splits=self.n_splits,
                shuffle=True,
                random_state=self.random_state
            )

        self.models_ = []


        for train_idx, val_idx in splitter.split(X_values, y_fit):

            model_clone = clone(self.model)

            y_train_fold = (
                y_fit.iloc[train_idx]
                if isinstance(y_fit, pd.Series)
                else y_fit[train_idx]
            )

            model_clone.fit(
                X_values[train_idx],
                y_train_fold
            )

            self.models_.append(model_clone)

        # ---- store metadata
        if self.classification:
            self.classes_ = self._label_encoder.classes_
        else:
            self.classes_ = None

        self.columns_ = self.columns

        return self


    def transform(self, X):

        if not hasattr(self, "models_"):
            raise ValueError("Model not fitted. Call fit() first.")

        X_values = X[self.columns_].to_numpy(dtype=float)

        preds = []

        for model in self.models_:
            if self.classification:
                preds.append(model.predict_proba(X_values))
            else:
                preds.append(model.predict(X_values).reshape(-1, 1))

        preds = np.mean(preds, axis=0)

        if self.classification:
            cols = [f"{self.prefix}_prob_{c}" for c in self.classes_]
        else:
            cols = [f"{self.prefix}_pred"]

        return pd.DataFrame(preds, columns=cols, index=X.index)


    def fit_transform(self, X, y=None):
        self.fit(X, y)
        return self.transform(X)
