# type: ignore

"""Octo Bags."""

# import concurrent.futures
# import logging
# import gzip
import pickle
from pathlib import Path
from statistics import mean
from typing import Any

import lz4.frame
import numpy as np
import pandas as pd
from attrs import define, field, validators
from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin
from upath import UPath

# sklearn imports for compatibility
from octopus.logger import LogGroup, get_logger

# Adjust this import path as needed depending on your package layout
from octopus.manager.ray_parallel import (
    init_ray,
    run_parallel_inner,
)
from octopus.metrics.utils import get_performance_from_predictions
from octopus.modules.octo.training import Training

logger = get_logger()


class TrainingWithLogging:
    """Logging class for trainings."""

    def __init__(self, inner_training, idx, logger, log_group_cls, log_prefix="EXP"):
        self._inner = inner_training
        self._idx = idx
        self._logger = logger
        self._log_group_cls = log_group_cls
        self._log_prefix = log_prefix

    def fit(self):
        """Fit function."""
        # Your logging policy lives here
        self._logger.set_log_group(self._log_group_cls.PROCESSING, f"{self._log_prefix} {self._idx}")
        self._logger.info("Starting execution")
        try:
            result = self._inner.fit()
            self._logger.set_log_group(self._log_group_cls.PREPARE_EXECUTION, f"{self._log_prefix} {self._idx}")
            self._logger.info("Completed successfully")
            return result
        except Exception as e:
            self._logger.exception(f"Exception occurred while executing training {self._idx}: {e!s}")
            # Decide your policy: raise to fail-fast, or return a sentinel to continue.
            return None  # or: raise


class FeatureImportanceWithLogging:
    """Logging wrapper for feature importance calculations."""

    def __init__(self, training, idx, fi_type, partition, logger, log_group_cls, log_prefix="FI"):
        self._training = training
        self._idx = idx
        self._fi_type = fi_type
        self._partition = partition
        self._logger = logger
        self._log_group_cls = log_group_cls
        self._log_prefix = log_prefix

    def fit(self):
        """Calculate feature importance for the training.

        Uses fit() method for Ray compatibility.
        """
        training_id = getattr(self._training, "training_id", self._idx)
        self._logger.set_log_group(self._log_group_cls.PROCESSING, f"{self._log_prefix} {training_id}")
        self._logger.info(f"Starting {self._fi_type} feature importance calculation")

        try:
            if self._fi_type == "internal":
                self._training.calculate_fi_internal()
            elif self._fi_type == "shap":
                self._training.calculate_fi_shap(partition=self._partition)
            elif self._fi_type == "permutation":
                self._training.calculate_fi_group_permutation(partition=self._partition)
            elif self._fi_type == "lofo":
                self._training.calculate_fi_lofo()
            elif self._fi_type == "constant":
                self._training.calculate_fi_constant()
            else:
                raise ValueError(f"FI type {self._fi_type} not supported")

            self._logger.set_log_group(self._log_group_cls.PREPARE_EXECUTION, f"{self._log_prefix} {training_id}")
            self._logger.info(f"Completed {self._fi_type} feature importance calculation")
            return self._training
        except Exception as e:
            self._logger.exception(f"Exception in {self._fi_type} FI calculation for training {training_id}: {e!s}")
            # Return the training object even if FI calculation failed
            return self._training


@define
class BagBase(BaseEstimator):
    """Base Container for Trainings.

    Supports:
    - execution of trainings, sequential/parallel
    - saving/loading
    - sklearn compatibility for inference tools like SHAP and permutation importance
    """

    bag_id: str = field(validator=[validators.instance_of(str)])
    trainings: list[Training] = field(validator=[validators.instance_of(list)])
    # same config parameters (execution type, num_workers) also used for
    # parallelization of optuna optimizations of individual inner loop trainings
    parallel_execution: bool = field(validator=[validators.instance_of(bool)])
    num_workers: int = field(validator=[validators.instance_of(int)])
    target_metric: str = field(validator=[validators.instance_of(str)])
    target_assignments: dict = field(validator=[validators.instance_of(dict)])
    row_column: str = field(validator=[validators.instance_of(str)])
    ml_type: str = field(validator=[validators.instance_of(str)])
    log_dir: UPath = field(validator=[validators.instance_of(UPath)])
    train_status: bool = field(default=False)

    # bag training outputs, initialized in post_init
    feature_importances: dict[str, pd.DataFrame] = field(init=False, validator=[validators.instance_of(dict)])
    n_features_used_mean: float = field(init=False, validator=[validators.instance_of(float)])

    @property
    def feature_groups(self) -> dict[str, list[str]]:
        """Experiment wide feature groups."""
        # assuming that there is at least one training
        return self.trainings[0].feature_groups

    @property
    def classes_(self):
        """Get unique classes from all trainings in the bag.

        Returns:
            numpy.ndarray or None: Array of unique class labels sorted in ascending order
                                  for classification/multiclass tasks, None otherwise.
        """
        # Return None for non-classification tasks
        if self.ml_type not in ["classification", "multiclass"]:
            return None

        # Check if we have any trainings
        if not self.trainings:
            return None

        # Collect all unique classes from all trainings
        all_classes = set()

        for training in self.trainings:
            # Check if the training has a fitted model with classes_ attribute
            if hasattr(training, "model") and hasattr(training.model, "classes_"):
                # Add classes from this training to the set
                training_classes = training.model.classes_
                if training_classes is not None:
                    all_classes.update(training_classes)

        # If no classes found, return None
        if not all_classes:
            return None

        # Convert to numpy array and sort (sklearn compatible)
        return np.array(sorted(all_classes))

    @property
    def positive_class(self):
        """Get the positive class for binary classification from trainings.

        Returns:
            The positive class value, or None if not applicable/determinable.

        Raises:
            ValueError: If training is missing config_training, missing positive_class
                       in config_training, or if trainings have inconsistent positive_class values.

        For binary classification, this method requires that each training has
        positive_class explicitly set in config_training. No fallback solutions are provided.
        """
        # Only applicable for binary classification
        if self.ml_type != "classification":
            return None

        # Initialize _positive_class if it doesn't exist (e.g., when loaded from pickle)
        if not hasattr(self, "_positive_class"):
            self._positive_class = None

        # Return cached value if already computed
        if self._positive_class is not None:
            return self._positive_class

        # Check if we have any trainings
        if not self.trainings:
            return None

        # Require explicit positive_class in training configurations
        # All trainings must have this property set
        for training in self.trainings:
            if not (hasattr(training, "config_training") and isinstance(training.config_training, dict)):
                raise ValueError(f"Training {getattr(training, 'training_id', 'unknown')} missing config_training")

            positive_class = training.config_training.get("positive_class")
            if positive_class is None:
                raise ValueError(
                    f"Training {getattr(training, 'training_id', 'unknown')} missing positive_class in config_training"
                )

            # Cache the first valid positive_class found
            if self._positive_class is None:
                self._positive_class = positive_class

            # Verify all trainings have the same positive_class
            elif self._positive_class != positive_class:
                raise ValueError(
                    f"Inconsistent positive_class values across trainings: {self._positive_class} vs {positive_class}"
                )

        return self._positive_class

    def __attrs_post_init__(self):
        # initialization here due to "Python immutable default"
        self.feature_importances = {}
        self.n_features_used_mean = 0.0
        self._positive_class = None  # Will be inferred when needed

    def get_params(self, deep=True):
        """Get parameters for this estimator (sklearn requirement)."""
        return {
            "bag_id": self.bag_id,
            "parallel_execution": self.parallel_execution,
            "num_workers": self.num_workers,
            "target_metric": self.target_metric,
            "ml_type": self.ml_type,
        }

    def set_params(self, **params):
        """Set parameters for this estimator (sklearn requirement)."""
        for key, value in params.items():
            if hasattr(self, key):
                setattr(self, key, value)
        return self

    def _more_tags(self):
        """Provide additional tags for sklearn compatibility."""
        return {
            "requires_fit": True,
            "no_validation": True,  # We handle our own validation
            "_xfail_checks": {
                "check_estimators_dtypes": "Custom data handling",
                "check_fit_score_takes_y": "Custom scoring method",
            },
        }

    def _train_parallel(self):
        """Run trainings in parallel using Ray (delegated to ray_parallel)."""
        # Ensure Ray is initialized before parallel execution
        init_ray(start_local_if_missing=True)

        # Prepare wrapped trainings with logging
        wrapped = [
            TrainingWithLogging(t, idx, logger, LogGroup, log_prefix="EXP") for idx, t in enumerate(self.trainings)
        ]
        # Orchestrate with Ray; exceptions propagate only if your wrapper re-raises.
        results = run_parallel_inner(wrapped, log_dir=self.log_dir)
        self.trainings = results

    def _train_sequential(self):
        """Run trainings sequentially in the current process."""
        successful_trainings = []
        failed_trainings = []

        for idx, training in enumerate(self.trainings):
            training_id = getattr(training, "training_id", idx)
            try:
                training.fit()
                successful_trainings.append(training)
                logger.info(
                    f"Inner sequential training completed for bag_id {self.bag_id} and training id {training_id}"
                )
            except Exception as e:  # pylint: disable=broad-except
                failed_trainings.append((training_id, str(e), type(e).__name__))
                logger.error(
                    f"Training failed for bag_id {self.bag_id}, training_id {training_id}: {e}, type: {type(e).__name__}"
                )

        # Log summary of training results
        total_trainings = len(self.trainings)
        successful_count = len(successful_trainings)
        failed_count = len(failed_trainings)

        logger.info(
            f"Bag {self.bag_id} training summary: {successful_count}/{total_trainings} successful, "
            f"{failed_count} failed"
        )

        if failed_trainings:
            logger.warning(f"Failed trainings in bag {self.bag_id}:")
            for training_id, error_msg, error_type in failed_trainings:
                logger.warning(f"  - Training {training_id}: {error_type} - {error_msg}")

        self.trainings = successful_trainings

    def fit(self):
        """Run all available trainings."""
        if self.parallel_execution is True:
            self._train_parallel()
        else:
            self._train_sequential()

        self.train_status = True

        # get used features in bag
        n_feat_lst = []
        for training in self.trainings:
            n_feat_lst.append(float(len(training.features_used)))

        if not n_feat_lst:
            raise ValueError(f"Empty feature list in bag: '{self.bag_id}'.")

        self.n_features_used_mean = mean(n_feat_lst)

    def get_predictions(self):
        """Extract bag predictions for train, dev, and test.

        Returns:
            dict: Dictionary containing predictions for each training and ensemble.
                 Keys are training IDs plus 'ensemble', values are dicts with 'train', 'dev', 'test' keys.
        """
        if not self.train_status:
            logger.set_log_group(LogGroup.TRAINING)
            logger.info("Running trainings first to be able to get scores")
            self.fit()

        predictions = {}
        pool: dict[str, list[pd.DataFrame]] = {key: [] for key in ["train", "dev", "test"]}

        for training in self.trainings:
            predictions[training.training_id] = training.predictions
            for part, pool_value in pool.items():
                pool_value.append(training.predictions[part])

        # Create ensemble predictions for each partition
        predictions["ensemble"] = {}
        for part, pool_value in pool.items():
            ensemble = pd.concat(pool_value, axis=0).groupby(by=self.row_column).mean().reset_index()
            for column in list(self.target_assignments.values()):
                ensemble[column] = ensemble[column].astype(self.trainings[0].data_train[column].dtype)
            predictions["ensemble"][part] = ensemble

        return predictions

    def get_performance(self):
        """Get performance using get_performance_from_predictions utility.

        This is a simpler alternative to get_performance() that:
        1. Gets predictions from bag.get_predictions()
        2. Calculates performance using get_performance_from_predictions()
        3. Restructures output to match expected format

        Returns:
            dict: Dictionary with performance values in the same format as get_performance()
        """
        # Get predictions from the bag
        predictions = self.get_predictions()

        # Calculate performance using the utility function
        performance = get_performance_from_predictions(
            predictions=predictions,
            target_metric=self.target_metric,
            target_assignments=self.target_assignments,
            positive_class=self.positive_class,
        )

        # Create performance_output dictionary with restructured data
        performance_output: dict[str, list[float] | float] = {}

        # Collect lists for train, dev, test from all non-ensemble trainings
        train_lst = []
        dev_lst = []
        test_lst = []

        for training_id, partitions in performance.items():
            if training_id != "ensemble":
                train_lst.append(partitions["train"])
                dev_lst.append(partitions["dev"])
                test_lst.append(partitions["test"])

        # Calculate averages
        performance_output["train_avg"] = mean(train_lst)
        performance_output["train_lst"] = train_lst
        performance_output["dev_avg"] = mean(dev_lst)
        performance_output["dev_lst"] = dev_lst
        performance_output["test_avg"] = mean(test_lst)
        performance_output["test_lst"] = test_lst

        # Add ensemble performance with renamed keys
        if "ensemble" in performance:
            performance_output["train_pool"] = performance["ensemble"]["train"]
            performance_output["dev_pool"] = performance["ensemble"]["dev"]
            performance_output["test_pool"] = performance["ensemble"]["test"]

        return performance_output

    def _calculate_fi_parallel(self, fi_type="internal", partition="dev"):
        """Calculate feature importance in parallel using Ray."""
        # Ensure Ray is initialized before parallel execution
        init_ray(start_local_if_missing=True)

        # Prepare wrapped trainings with logging for feature importance calculation
        wrapped = [
            FeatureImportanceWithLogging(
                training=t,
                idx=idx,
                fi_type=fi_type,
                partition=partition,
                logger=logger,
                log_group_cls=LogGroup,
                log_prefix="FI",
            )
            for idx, t in enumerate(self.trainings)
        ]

        # Execute feature importance calculations in parallel
        # Use the same pattern as training execution
        results = run_parallel_inner(wrapped, log_dir=self.log_dir, num_cpus=1)

        # Update trainings with results (should be the same objects with FI calculated)
        self.trainings = results

    def _calculate_fi_sequential(self, fi_type="internal", partition="dev"):
        """Calculate feature importance sequentially."""
        successful_calculations = []
        failed_calculations = []

        for idx, training in enumerate(self.trainings):
            training_id = getattr(training, "training_id", idx)
            try:
                if fi_type == "internal":
                    training.calculate_fi_internal()
                elif fi_type == "shap":
                    training.calculate_fi_shap(partition=partition)
                elif fi_type == "permutation":
                    training.calculate_fi_group_permutation(partition=partition)
                elif fi_type == "lofo":
                    training.calculate_fi_lofo()
                elif fi_type == "constant":
                    training.calculate_fi_constant()
                else:
                    raise ValueError(f"FI type {fi_type} not supported")

                successful_calculations.append(training)
                logger.info(
                    f"Feature importance ({fi_type}) calculation completed for bag_id {self.bag_id} "
                    f"and training id {training_id}"
                )
            except Exception as e:  # pylint: disable=broad-except
                failed_calculations.append((training_id, str(e), type(e).__name__))
                logger.error(
                    f"Feature importance ({fi_type}) calculation failed for bag_id {self.bag_id}, "
                    f"training_id {training_id}: {e}, type: {type(e).__name__}"
                )
                # Still include the training even if FI calculation failed
                successful_calculations.append(training)

        # Log summary of FI calculation results
        total_trainings = len(self.trainings)
        successful_count = len(successful_calculations)
        failed_count = len(failed_calculations)

        logger.info(
            f"Bag {self.bag_id} feature importance ({fi_type}) summary: "
            f"{successful_count}/{total_trainings} successful, {failed_count} failed"
        )

        if failed_calculations:
            logger.warning(f"Failed feature importance ({fi_type}) calculations in bag {self.bag_id}:")
            for training_id, error_msg, error_type in failed_calculations:
                logger.warning(f"  - Training {training_id}: {error_type} - {error_msg}")

        self.trainings = successful_calculations

    def _calculate_fi(self, fi_type="internal", partition="dev"):
        """Calculate feature importance using parallel or sequential execution."""
        if self.parallel_execution:
            self._calculate_fi_parallel(fi_type=fi_type, partition=partition)
        else:
            self._calculate_fi_sequential(fi_type=fi_type, partition=partition)

    def get_selected_features(self, fi_methods=None):
        """Get features selected by model, depending on fi method.

        The list of selected features will be derived only from one feature
        importance method out of the ones specified in fi_methods,
        with the following ranking: (1) permutation (2) shap (3) internal.
        """
        # we assume that feature_importances were previously calculated
        if fi_methods is None:
            fi_methods = []

        if "permutation" in fi_methods:
            fi_df = self.feature_importances["permutation_dev_mean"]
        elif "shap" in fi_methods:
            fi_df = self.feature_importances["shap_dev_mean"]
        elif "internal" in fi_methods:
            fi_df = self.feature_importances["internal_mean"]
        elif "constant" in fi_methods:
            fi_df = self.feature_importances["constant_mean"]
        else:
            logger.set_log_group(LogGroup.RESULTS)
            logger.info("No features selected, return empty list")
            return []

        # only keep nonzero features
        fi_df = fi_df[fi_df["importance"] != 0]

        # store group features
        groups_df = fi_df[fi_df["feature"].str.startswith("group")].copy()

        # remove all group features -> single features
        fi_df = fi_df[~fi_df["feature"].str.startswith("group")]
        feat_single = fi_df["feature"].tolist()

        # For each feature group with positive importance (only),
        # check if any feature is in feat_single. In not, add the
        # one with the largest feature importance
        groups = groups_df[groups_df["importance"] > 0]["feature"].tolist()
        feat_additional = []
        for key in groups:
            features = self.feature_groups.get(key, [])
            if features and not any(feature in feat_single for feature in features):
                # Find the feature with the highest importance in fi_df
                feature_importances = fi_df[fi_df["feature"].isin(features)]
                if not feature_importances.empty:
                    best_feature = feature_importances.loc[feature_importances["importance"].idxmax(), "feature"]
                    feat_additional.append(best_feature)

        # Add the additional features to feat_single and remove duplicates
        feat_all = list(set(feat_single + feat_additional))

        logger.set_log_group(LogGroup.RESULTS, f"BAG {self.bag_id}")
        logger.info(f"Number of selected features: {len(feat_all)}")
        logger.info(f"Number of single features: {len(feat_single)}")
        logger.info(f"Number of features from groups: {len(feat_additional)}")

        return sorted(feat_all, key=lambda x: (len(x), sorted(x)))

    def calculate_feature_importances(self, fi_methods=None, partitions=None):
        """Extract feature importances of all models in bag."""
        # we always extract internal feature importances, if available
        if fi_methods is None:
            fi_methods = []
        if partitions is None:
            partitions = ["dev", "test"]

        self._calculate_fi(fi_type="internal")

        for method in fi_methods:
            if method == "internal":
                pass  # already done
            elif method == "shap":
                for partition in partitions:
                    self._calculate_fi(fi_type="shap", partition=partition)
            elif method == "permutation":
                for partition in partitions:
                    self._calculate_fi(fi_type="permutation", partition=partition)
            elif method == "lofo":
                self._calculate_fi(fi_type="lofo")
            elif method == "constant":
                self._calculate_fi(fi_type="constant")
            else:
                raise ValueError(f"Feature importance method {method} not supported.")

        # save feature importances for every training in bag
        for training in self.trainings:
            self.feature_importances[training.training_id] = training.feature_importances

        # summary feature importances for all trainings (mean + count)
        # internal, permutation_dev, shap_dev only
        # save in bag
        for method in fi_methods:
            if method == "internal":
                method_str = "internal"
            elif method == "constant":
                method_str = "constant"
            else:
                method_str = method + "_dev"
            fi_pool = []
            for training in self.trainings:
                fi_pool.append(training.feature_importances[method_str])
            fi = pd.concat(fi_pool, axis=0)

            # calculate mean feature importances, keep zero entries
            self.feature_importances[method_str + "_mean"] = (
                fi[["feature", "importance"]]
                .groupby(by="feature")
                .sum()
                .div(len(self.trainings))  # not all features in each fi-table
                .sort_values(by="importance", ascending=False)
                .reset_index()
            )

            # calculate count feature importances, keep zero entries
            non_zero_importances = fi[fi["importance"] != 0][["feature", "importance"]].groupby(by="feature").count()
            # Create a DataFrame with all features, init importance counts to zero
            all_features = pd.DataFrame(fi["feature"].unique(), columns=["feature"])
            all_features["importance"] = 0
            # Update the importance counts for non-zero importances
            all_features = all_features.set_index("feature")
            all_features.update(non_zero_importances)
            all_features = all_features.reset_index()
            # Sort and reset index
            self.feature_importances[method_str + "_count"] = all_features.sort_values(
                by="importance", ascending=False
            ).reset_index(drop=True)

        return self.feature_importances

    def predict(self, x):
        """Predict with sklearn compatibility."""
        # Import sklearn validation here to avoid auto-formatter issues

        # Check if the bag has fitted trainings
        if not self.trainings:
            raise ValueError("No trainings available in bag")

        # Check if all trainings are fitted
        for training in self.trainings:
            if not getattr(training, "is_fitted", False):
                raise ValueError(f"Training {training.training_id} is not fitted")

        # Convert pandas DataFrame to numpy if needed for sklearn compatibility
        if hasattr(x, "values"):
            x = x.values

        preds_lst = []
        weights_lst = []
        for training in self.trainings:
            train_w = training.training_weight
            weights_lst.append(train_w)
            preds_lst.append(train_w * training.predict(x))

        # return mean of weighted predictions
        return np.sum(np.array(preds_lst), axis=0) / sum(weights_lst)

    def predict_proba(self, x):
        """Predict_proba with sklearn compatibility."""
        # Only available for classification tasks
        if self.ml_type not in ["classification", "multiclass"]:
            raise AttributeError(f"predict_proba is not available for ml_type '{self.ml_type}'")

        # Check if the bag has fitted trainings
        if not self.trainings:
            raise ValueError("No trainings available in bag")

        # Check if all trainings are fitted
        for training in self.trainings:
            if not getattr(training, "is_fitted", False):
                raise ValueError(f"Training {training.training_id} is not fitted")

        # Convert pandas DataFrame to numpy if needed for sklearn compatibility
        if hasattr(x, "values"):
            x = x.values

        preds_lst = []
        weights_lst = []
        for training in self.trainings:
            train_w = training.training_weight
            weights_lst.append(train_w)
            preds_lst.append(train_w * training.predict_proba(x))

        # return mean of weighted predictions
        return np.sum(np.array(preds_lst), axis=0) / sum(weights_lst)

    def _estimator_type(self):
        """Return the estimator type for sklearn compatibility."""
        if self.ml_type in ["classification", "multiclass"]:
            return "classifier"
        else:
            return "regressor"

    def to_pickle(self, file_path: str | Path | UPath):
        """Save object to a compressed pickle file.

        Args:
            file_path: The name of the file to save the pickle data to.
        """
        # Gzip was too slow, switched to lz4
        # with file_path.open("wb") as file, gzip.GzipFile(fileobj=file, mode="wb") as gzip_file:
        #    pickle.dump(self, gzip_file)
        with file_path.open("wb") as file, lz4.frame.open(file, mode="wb") as lz4_file:
            pickle.dump(self, lz4_file)

    @classmethod
    def from_pickle(cls, file_path: str | Path | UPath) -> "BagBase":
        """Load object to a compressed pickle file.

        Args:
            file_path: The path to the file to load the pickle data from.

        Returns:
            BagBase: The loaded instance of BagBase.

        Raises:
            TypeError: If the file does not contain an BagBase instance.
        """
        # Gzip was too slow, switched to lz4
        # with file_path.open("rb") as file, gzip.GzipFile(fileobj=file, mode="rb") as gzip_file:
        #    return pickle.load(gzip_file)
        with file_path.open("rb") as file, lz4.frame.open(file, mode="rb") as lz4_file:
            data = pickle.load(lz4_file)

        if not isinstance(data, cls):
            raise TypeError(f"Loaded object is not of type {cls.__name__}")

        return data


@define
class BagClassifier(BagBase, ClassifierMixin):
    """Bag for classification tasks with sklearn ClassifierMixin."""

    def _estimator_type(self):
        """Return the estimator type for sklearn compatibility."""
        return "classifier"


@define
class BagRegressor(BagBase, RegressorMixin):
    """Bag for regression tasks with sklearn RegressorMixin."""

    def _estimator_type(self):
        """Return the estimator type for sklearn compatibility."""
        return "regressor"

    def predict_proba(self, x):
        """Predict_proba not available for regression tasks."""
        raise AttributeError("predict_proba is not available for regression tasks")


class _BagFunction:
    """Class to create appropriate Bag instance based on ml_type (factory function).

    Args:
        **kwargs: Arguments to pass to the Bag constructor

    Returns:
        BagClassifier or BagRegressor: Appropriate Bag instance based on ml_type
    """

    def __call__(self, **kwargs: Any) -> BagClassifier | BagRegressor:
        ml_type = kwargs.get("ml_type", "regression")

        if ml_type in ["classification", "multiclass"]:
            return BagClassifier(**kwargs)
        else:
            return BagRegressor(**kwargs)

    @staticmethod
    def from_pickle(file_path: str | Path | UPath) -> BagClassifier | BagRegressor:
        """Load a Bag object from a compressed pickle file.

        Args:
            file_path: The path to the file to load the pickle data from.

        Returns:
            BagClassifier or BagRegressor: The loaded instance.

        Raises:
            TypeError: If the file does not contain an BagClassifier | BagRegressor instance.
        """
        with file_path.open("rb") as file, lz4.frame.open(file, mode="rb") as lz4_file:
            data = pickle.load(lz4_file)

        if not isinstance(data, BagClassifier | BagRegressor):
            raise TypeError(f"Loaded object is not of type {BagClassifier | BagRegressor}")

        return data


Bag = _BagFunction()
