from typing import Any
from typing_extensions import Self

import numpy as np
import pandas as pd
import xgboost as xgb
from sklearn.metrics import mean_squared_error

def to_category(
    df: pd.DataFrame, columns: list[str], inplace: bool = False
) -> pd.DataFrame:
    """
    Convert the specified columns in a pandas DataFrame to categorical data type.

    Parameters
    ----------
    df : pandas.DataFrame
        The input DataFrame to be converted.
    columns : List[str]
        The list of column names to be converted to categorical data type.
    inplace : bool, optional
        If True, modify the input DataFrame in place. Otherwise, return a new copy of the DataFrame.
        The default is False.

    Returns
    -------
    pandas.DataFrame
        The modified DataFrame with the specified columns converted to categorical data type.
        If inplace=True, the original DataFrame is modified.
    """

    if not inplace:
        df = df.copy()

    for c in columns:
        df[c] = df[c].astype("category")

    return df


class BoostingModel:
    """Implements a Boosting model to perform both regression and classification"""

    def __init__(
        self,
        params: dict[str, Any] = None,
        regression: bool = False,
        cat_features: list[str] = None,
    ):
        """
        Initializes an instance of the XGBModel class with the specified parameters.

        Parameters
        ----------
        params : dict[str, Any], optional
            A dictionary of XGBoost model parameters to use. If not specified, default values will be used.
        regression : bool, optional
            A boolean indicating whether the model is for regression (True) or classification (False). Default is False.
        cat_features : list[str], optional
            A list of the names of categorical features in the input data. Default is an empty list.

        Examples
        --------
        >>> model = XGBModel(regression=True, cat_features=['A', 'B'])
        """
        if params is None:
            params = {
                "objective": "reg:logistic",
                "eval_metric": "logloss",
                "verbosity": 0,
                "verbose_eval": False,
                "n_estimators": 1000,
                "early_stopping_rounds": 50,
                "tree_method": "hist",
                "n_jobs": -1,
                "enable_categorical": True,
                "booster": "gbtree",
                "max_depth": 3,
                "learning_rate": 0.03,
                "subsample": 0.8,
                "colsample_bytree": 0.8,
            }
            if regression:
                params["objective"] = "reg:squarederror"
                params["eval_metric"] = "rmse"
        self.params = params
        self.regression = regression
        if cat_features is None:
            cat_features = []
        self.cat_features = cat_features

    def fit(
        self,
        X_train: pd.DataFrame,
        y_train: pd.DataFrame,
        X_val: pd.DataFrame,
        y_val: pd.DataFrame,
    ) -> Self:
        """
        Trains the XGBoost model using the specified training and
        validation data (used for early stopping).

        Parameters
        ----------
        X_train : pd.DataFrame
            A pandas DataFrame of training input data.
        y_train : pd.DataFrame
            A pandas DataFrame of training target data.
        X_val : pd.DataFrame
            A pandas DataFrame of validation input data.
        y_val : pd.DataFrame
            A pandas DataFrame of validation target data.

        Returns
        -------
        XGBModel
            The trained XGBModel object.

        Examples
        --------
        >>> model = XGBModel()
        >>> model.fit(X_train, y_train, X_val, y_val)
        """
        self.clf = xgb.XGBModel(**self.params)
        X_train = to_category(X_train, self.cat_features)
        X_val = to_category(X_val, self.cat_features)
        self.clf.fit(X_train, y_train, eval_set=[(X_val, y_val)], verbose=False)

        self.refit(pd.concat([X_train, X_val]), pd.concat([y_train, y_val]))
        return self

    def refit(self, X_train: pd.DataFrame, y_train: pd.DataFrame) -> Self:
        params = self.params.copy()
        params["n_estimators"] = self.clf.best_iteration
        params["early_stopping_rounds"] = None

        X_train = to_category(X_train, self.cat_features)
        self.clf = xgb.XGBModel(**params).fit(X_train, y_train)

        return self

    def predict(self, X: pd.DataFrame) -> np.ndarray:
        X = to_category(X, self.cat_features)
        preds = self.clf.predict(X)
        if not self.regression:
            preds = preds / np.sum(preds, axis=1, keepdims=True)

        return preds

    def score(self, X: pd.DataFrame, y: pd.DataFrame) -> float:
        X = to_category(X, self.cat_features)
        if self.regression:
            return mean_squared_error(y, self.predict(X))
        eps = 1e-8
        y_pred = np.clip(self.predict(X), eps, 1 - eps)
        return -(y.values * np.log(y_pred)).sum(axis=1).mean()
