# -*- coding: utf-8 -*-
"""FLAVORS2

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1UVzHGy3VSrA2g2EeIygclh1_JeSITr1p
"""

# Core Python
import numpy as np
import pandas as pd
import random
import datetime
from collections import defaultdict
from heapq import heappop, heappush
from functools import wraps
import inspect

# Scikit-learn
from sklearn.utils import check_array
from sklearn.model_selection import cross_val_score, train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.feature_selection import mutual_info_classif, mutual_info_regression
from sklearn.base import BaseEstimator, TransformerMixin, clone
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import RobustScaler


# Joblib
from joblib import Parallel, delayed,parallel_backend

class RankOrderTree:
    def __init__(self):
        self.tree = []

    def insert(self, score):
        heappush(self.tree, -score)

    def remove_max(self):
        if self.tree:
            heappop(self.tree)

    def get_best_score(self):
        if self.tree:
            return -self.tree[0]
        return float('inf') # Or perhaps -float('inf') if maximizing? Depends on context.

    def get_second_best_score(self):
        # Note: self.tree[1] is not guaranteed to be the second smallest element in a min-heap.
        # This might be an approximation.
        if len(self.tree) >= 2:
             # Find the actual second smallest element
             first = heappop(self.tree)
             if not self.tree:
                 heappush(self.tree, first) # Put it back if it was the only one
                 return float('inf') # Or appropriate value indicating no second element
             second_smallest = self.tree[0]
             heappush(self.tree, first) # Put the smallest back
             return -second_smallest # Return the original second largest score
        return float('inf') # Or appropriate value

def ensure_metric_protocol(metric):
    """
    Wrap a user metric so it:
      - accepts signature (X, y, sample_weight=None)
      - returns a dict with at least {'score': float}
      - if the original returns a number, coerce to {'score': number}
    """
    sig = inspect.signature(metric)
    params = list(sig.parameters.values())
    accepts_sample_weight = any(p.name == "sample_weight" for p in params)

    @wraps(metric)
    def wrapped(X, y, sample_weight=None):
        # Call with or without sample_weight depending on the user's signature
        if accepts_sample_weight:
            out = metric(X, y, sample_weight=sample_weight)
        else:
            out = metric(X, y)

        # Normalize output to dict
        if isinstance(out, (int, float, np.number)):
            return {"score": float(out)}  # numeric -> dict
        if not isinstance(out, dict):
            raise TypeError("Custom metric must return a number or a dict with at least {'score': float}.")
        if "score" not in out:
            raise KeyError("Custom metric dict must contain a 'score' key.")
        # Make sure score is a float
        out = dict(out)
        out["score"] = float(out["score"])
        return out

    return wrapped


def test_metric_protocol(custom_metric):
    """Lightweight self-test on shape + protocol (dict with score)."""
    X = np.random.rand(20, 5)
    y = np.random.randint(0, 2, 20)
    res = custom_metric(X, y, sample_weight=None)
    if not isinstance(res, dict) or "score" not in res:
        raise RuntimeError("Metric protocol violation: must return dict with 'score'.")
    if not np.isfinite(res["score"]):
        raise RuntimeError("Metric 'score' must be finite.")


class FLAVORS2:
    def __init__(self, budget, minimize=False, metrics=None, c_sub=4, feature_priors=None, iters=None,
                 n_jobs=1,boruta=False,random_state=42):

        self.leaderboard = []
        self.budget = budget
        self.c_sub = c_sub
        self.minimize = minimize
        self.fitted = False
        self.dup_dct = defaultdict(int)
        self.feature_cost_history = {}
        self.feature_addition_times = {}
        self.feature_removal_times = {}
        self.performance_history = []
        self.cost_history = []
        self.feature_priors = feature_priors
        self.n_jobs = n_jobs
        self.boruta = boruta
        self._rng = np.random.RandomState(random_state)

        if metrics is None:
            # Use a static method reference directly
            metrics = [FLAVORS2.default_metric]

        # Apply validation wrapper to each metric
        if metrics is None:
            metrics = [FLAVORS2.default_metric]

        # Wrap: (1) ensure protocol (dict+sample_weight), (2) validate arrays
        metrics = [FLAVORS2.validate_input_arrays(ensure_metric_protocol(m)) for m in metrics]

        for metric in metrics:
            # New protocol-aware test
            test_metric_protocol(metric)

        self.metrics = metrics
        if len(metrics) == 1:
            self.metric = metrics[0]


    def update_pareto_history(self, new_tuple):
        if not hasattr(self, 'pareto_history'):
            self.pareto_history = []
        self.pareto_history.append(new_tuple)
        # Optional: Keep pareto front updated (remove dominated points)
        # This can be complex to implement efficiently here.
        # Simple append keeps all evaluated points.

    def calculate_hypervolume(self):
        # Note: This hypervolume calculation seems non-standard.
        # Standard algorithms often use reference points and handle dominance.
        # Review this logic if exact hypervolume calculation is critical.
        # The provided logic might be a custom approximation.
        if not hasattr(self, 'pareto_history') or not self.pareto_history:
            return 0.0

        # Assuming higher values are better for all objectives for this calculation
        # If minimizing=True for some objectives, they need to be negated before this calculation.
        # This example assumes maximization or metrics are already adjusted.
        points = sorted(self.pareto_history, key=lambda x: x[0]) # Sort by first objective
        n = len(points)
        if n == 0:
            return 0.0

        # Example using a simple hypervolume calculation (requires a reference point)
        # This needs adaptation based on the specific hypervolume algorithm intended.
        # Let's assume a reference point (e.g., origin if all objectives are positive)
        ref_point = np.zeros(len(points[0])) # Adjust if objectives can be negative

        # A simplified placeholder - real hypervolume calculation is more complex.
        # This calculates sum of areas/volumes relative to origin, which isn't standard hypervolume.
        # Consider using a dedicated library (e.g., pygmo) if precise HV is needed.
        hypervolume = 0.0
        last_point = ref_point
        for i in range(n):
             # This part is conceptually incorrect for standard hypervolume.
             # It's summing individual point contributions without proper slicing.
             volume_contribution = 1.0
             for j in range(len(points[i])):
                 volume_contribution *= max(0.0, points[i][j] - ref_point[j])
             # This logic needs review based on the intended hypervolume algorithm
             # Example using a slicing approach (still simplified):
             if i == 0:
                 hypervolume += volume_contribution # Incorrect logic, placeholder
             else:
                 # Very simplified diff - incorrect for general case
                 diff_vol = 1.0
                 for j in range(len(points[i])):
                      diff_vol *= max(0.0, points[i][j] - points[i-1][j]) # Incorrect
                 hypervolume += diff_vol # Incorrect logic, placeholder


        # Return a placeholder - replace with correct HV calculation if needed
        # A simple proxy could be the number of non-dominated points or average objective values.
        # For now, returning count as a placeholder for change.
        return len(self.pareto_history) # Placeholder

    @staticmethod
    def _extract_importances(est):
        """
        Return a 1D non-negative importance vector if available, else None.
        Supports coef_ (uses abs mean across classes) and feature_importances_.
        """
        if est is None:
            return None
        if hasattr(est, "feature_importances_"):
            imp = np.asarray(est.feature_importances_, dtype=float).ravel()
            return np.abs(imp)
        if hasattr(est, "coef_"):
            coefs = np.asarray(est.coef_, dtype=float)
            if coefs.ndim == 2:
                imp = np.mean(np.abs(coefs), axis=0)
            else:
                imp = np.abs(coefs).ravel()
            return imp
        return None

    @staticmethod
    def _make_shadow_features(X_subset, n_shadow, rng):
        """
        Boruta-style shadows: pick columns (with replacement) from X_subset and permute rows.
        Returns array of shape (n_samples, n_shadow).
        """
        n_samples, n_feats = X_subset.shape
        if n_feats == 0 or n_shadow <= 0:
            return np.zeros((n_samples, 0), dtype=X_subset.dtype)
        pick = rng.randint(0, n_feats, size=int(n_shadow))
        shadows = []
        for j in pick:
            col = X_subset[:, j].copy()
            rng.shuffle(col)  # in-place permutation
            shadows.append(col.reshape(-1, 1))
        return np.hstack(shadows) if shadows else np.zeros((n_samples, 0), dtype=X_subset.dtype)


    @staticmethod
    def validate_input_arrays(func):
        @wraps(func)
        def wrapper(X, y, sample_weight=None):
            X = check_array(X, ensure_2d=True, force_all_finite=True)
            y = check_array(y, ensure_2d=False, force_all_finite=True, ensure_min_samples=1)
            if sample_weight is not None:
                sample_weight = check_array(sample_weight, ensure_2d=False, force_all_finite=True)
                if sample_weight.shape[0] != y.shape[0]:
                    raise ValueError("sample_weight must have shape (n_samples,).")
            return func(X, y, sample_weight=sample_weight)
        return wrapper

    @staticmethod
    def test_custom_function(custom_metric):
        # Generate slightly more realistic data
        X = np.random.rand(20, 5)
        y = np.random.randint(0, 2, 20)

        try:
            # Call the potentially wrapped function
            result = custom_metric(X, y)
        except Exception as e:
            raise RuntimeError(f"Error occurred while testing the custom metric function: {e}") from e

        assert isinstance(result, (int, float, np.number)), f'Metric must return a number, but got {type(result)}'
        assert np.isfinite(result), 'Metric must return a finite number'


    @staticmethod
    def default_metric(X, y, sample_weight=None):
        """
        Default classifier accuracy with optional sample_weight.
        Returns {'score': float, 'model': fitted_estimator}. Higher is better.
        """
        if X.shape[1] == 0:
            return {"score": 0.0, "model": None}

        try:
            # We need the fitted model, so we can't use cross_val_score directly.
            # We'll do a single split and train. For a more robust default,
            # one might train on the full dataset or implement a CV loop manually.
            X_train, X_test, y_train, y_test = train_test_split(
                X, y, test_size=0.25, random_state=42, stratify=y if len(np.unique(y)) > 1 else None
            )

            est = LogisticRegression(max_iter=1000, solver='liblinear', random_state=42)

            fit_params = {}
            if sample_weight is not None:
                # Need to split sample_weight accordingly if used
                sw_train, sw_test = train_test_split(sample_weight, test_size=0.25, random_state=42)
                fit_params["sample_weight"] = sw_train

            est.fit(X_train, y_train, **fit_params)
            score = est.score(X_test, y_test)

            return {"score": float(score), "model": est}

        except ValueError as e:
            print(f"Warning: Default metric failed for subset. Returning 0. Error: {e}")
            return {"score": 0.0, "model": None}

    def estimate_feature_cost_impact(self, current_subset, candidate_feature):
        # Ensure fastest_feat is valid before using it
        valid_fastest_feat = self.fastest_feat is not None and not np.isnan(self.fastest_feat) and int(self.fastest_feat) < len(self.unique_counts)

        n_current = len(current_subset)
        # Avoid division by zero for scaling factor
        scaling_factor = (n_current + 1) / n_current if n_current > 0 else 1.0

        # Use try-except for accessing unique_counts with potentially invalid index
        try:
            unique_value_factor = self.unique_counts[candidate_feature] / self.unique_counts[int(self.fastest_feat)] \
                                  if valid_fastest_feat and self.unique_counts[int(self.fastest_feat)] > 0 else 1.0
        except IndexError:
            unique_value_factor = 1.0 # Default if indices are out of bounds

        # Ensure candidate_feature exists in history if used
        historical_factor = self.feature_cost_history.get(candidate_feature, 1.0)

        # Use hasattr for optional model complexity function
        if hasattr(self, 'model_time_complexity') and callable(self.model_time_complexity):
            try:
                model_factor = self.model_time_complexity(n_current, n_current + 1)
                scaling_factor *= model_factor
            except Exception as e:
                print(f"Warning: model_time_complexity function failed: {e}")
                 # Proceed without model_factor if it fails

        # Combine factors
        # Ensure weights are reasonable (e.g., sum to 1)
        final_cost = scaling_factor * (0.3 * unique_value_factor + 0.7 * historical_factor)

        # Ensure cost is positive and finite
        return max(1e-9, final_cost) if np.isfinite(final_cost) else 1.0


    def update_feature_performance(self, feature_scores):
        """
        Update feature performance history based on feature-specific scores.

        Args:
            feature_scores (dict): A dictionary mapping feature indices to their
                                   specific performance scores for an evaluation.
        """
        # Ensure performance tracking attributes are initialized
        if not hasattr(self, 'feature_performance') or self.feature_performance is None:
            self.feature_performance = np.zeros(self.n_feats)
        if not hasattr(self, 'feature_counts') or self.feature_counts is None:
            self.feature_counts = np.zeros(self.n_feats)

        # Iterate over each feature and its specific score from the latest evaluation
        for feat, performance in feature_scores.items():
            # Ensure performance is a finite number BEFORE calculating performance_score
            if performance is None or not np.isfinite(performance):
                 continue

            # Determine score based on optimization direction (lower is better internally)
            performance_score = -performance if self.minimize else performance
            if not np.isfinite(performance_score):
                 continue

            # --- Initialize or update global min/max of per-feature scores ---
            if not hasattr(self, 'min_performance') or self.min_performance is None or \
               not hasattr(self, 'max_performance') or self.max_performance is None:
                self.min_performance = performance_score
                self.max_performance = performance_score
            else:
                self.min_performance = min(self.min_performance, performance_score)
                self.max_performance = max(self.max_performance, performance_score)

            # --- Normalize the feature's performance score ---
            if not isinstance(self.min_performance, (int, float, np.number)) or \
               not isinstance(self.max_performance, (int, float, np.number)):
                 norm_performance = 0.5 # Default normalization if state is unexpected
            else:
                 range_perf = self.max_performance - self.min_performance
                 # Avoid division by zero if range is too small
                 norm_performance = (performance_score - self.min_performance) / max(range_perf, 1e-9) if range_perf > 1e-9 else 0.5

            # --- Update performance for the feature using Exponential Moving Average (EMA) ---
            if 0 <= feat < self.n_feats:
                 self.feature_counts[feat] += 1
                 alpha = 0.1 # Smoothing factor (adjust as needed)
                 current_perf = float(self.feature_performance[feat]) if np.isfinite(self.feature_performance[feat]) else 0.0
                 self.feature_performance[feat] = alpha * norm_performance + (1 - alpha) * current_perf

        # Update feature stability based on leaderboard changes
        self.update_feature_stability()


    def update_feature_cost_history(self, previous_subset, new_subset, evaluation_time):
        # Ensure evaluation_time is valid
        if not np.isfinite(evaluation_time) or evaluation_time <= 0:
            evaluation_time = np.mean(self.cost_history) if self.cost_history else 1e-3 # Use average or small default
            evaluation_time = max(1e-9, evaluation_time) # Ensure positive

        prev_set = set(previous_subset) if previous_subset is not None else set()
        new_set = set(new_subset) if new_subset is not None else set()

        added_features = new_set - prev_set
        removed_features = prev_set - new_set

        # Initialize history tracking if first time
        if not hasattr(self, 'feature_addition_times'): self.feature_addition_times = {}
        if not hasattr(self, 'feature_removal_times'): self.feature_removal_times = {}
        if not hasattr(self, 'feature_cost_history'): self.feature_cost_history = {}

        for feat in added_features:
            if feat not in self.feature_addition_times:
                self.feature_addition_times[feat] = []
                self.feature_removal_times[feat] = []
                self.feature_cost_history[feat] = 1.0 # Initial default cost factor

            self.feature_addition_times[feat].append(evaluation_time)
            # Keep limited history to avoid memory issues and adapt to recent changes
            self.feature_addition_times[feat] = self.feature_addition_times[feat][-10:]

            # Update cost history based on comparison with removal times
            if self.feature_addition_times[feat] and self.feature_removal_times.get(feat):
                avg_time_with = np.mean(self.feature_addition_times[feat])
                avg_time_without = np.mean(self.feature_removal_times[feat])
                # Avoid division by zero, ensure cost is reasonable
                self.feature_cost_history[feat] = max(0.1, min(10.0, avg_time_with / max(1e-9, avg_time_without)))


        for feat in removed_features:
            if feat not in self.feature_removal_times:
                self.feature_addition_times[feat] = []
                self.feature_removal_times[feat] = []
                self.feature_cost_history[feat] = 1.0

            self.feature_removal_times[feat].append(evaluation_time)
            self.feature_removal_times[feat] = self.feature_removal_times[feat][-10:]

            if self.feature_removal_times[feat] and self.feature_addition_times.get(feat):
                 avg_time_with = np.mean(self.feature_addition_times[feat])
                 avg_time_without = np.mean(self.feature_removal_times[feat])
                 self.feature_cost_history[feat] = max(0.1, min(10.0, avg_time_with / max(1e-9, avg_time_without)))


    def initial_ECI(self, feat):
        # ECI: Estimated Cost per unit of Improvement (lower is better)
        # This estimates cost impact of adding the feature to an empty set.
        # Improvement estimate could be added here (e.g., using feature_priors)
        est_cost = self.estimate_feature_cost_impact([], feat)
        # Placeholder for improvement (e.g., mutual info score)
        est_improvement = self.feature_priors[feat] if hasattr(self, 'feature_priors') and feat < len(self.feature_priors) else 1e-3
        # Avoid division by zero
        return est_cost / max(1e-9, est_improvement)


    def calculate_ECI(self, K0, K1, K2, delta, best_error, feature_error, tau):
         # This seems like a custom ECI calculation, maybe related to FLAML's CFEI/CS.
         # K0, K1, K2 likely relate to scores/errors of related configurations.
         # delta: change in score/error
         # best_error: current best error
         # feature_error: error with the feature change
         # tau: iteration count or time?
         # self.c, self.kl: exploration/cost parameters

         # Ensure inputs are valid
         delta = max(abs(delta), 1e-9) # Avoid division by zero
         time_cost = self.kl if hasattr(self, 'kl') and self.kl > 0 else 1e-3

         # Error gain (or loss if minimizing and feature_error is higher)
         error_diff = feature_error - best_error # Assumes lower error is better
         if self.minimize: error_diff = -error_diff # Adjust if minimizing

         # Improvement part of ECI (higher is better)
         improvement_term = error_diff * tau / delta

         # Exploration/Uncertainty part (related to score variance/range)
         exploration_term1 = max(abs(K0 - K1), abs(K1 - K2)) # Score difference/range
         # Exploration related to cost/time?
         exploration_term2 = self.c * time_cost if hasattr(self, 'c') else 1.0 * time_cost

         # Combine terms - this logic needs clarification based on the source algorithm
         # If feature_error is the same as best_error, use exploration term only?
         if abs(feature_error - best_error) < 1e-9:
             eci = min(exploration_term1, exploration_term2)
         else:
             # Combine improvement and exploration - MAX suggests taking the more promising signal?
             eci = max(improvement_term, min(exploration_term1, exploration_term2))

         # ECI typically measures Cost / Improvement, so lower is better.
         # The calculation above seems to yield a value where higher might be better? Recheck formula.
         # Assuming higher value from calc means "more promising", return it.
         return eci

    @staticmethod
    def _compute_subset_score(subset, X, y, sample_weight, metrics, minimize, n_feats):
        start = datetime.datetime.now()
        # Ensure subset is a list of unique integers
        subset = sorted(list(set(map(int, subset))))

        # Handle empty subset
        if not subset:
             eval_time = (datetime.datetime.now() - start).total_seconds()
             # Return default poor score for empty set
             default_score = float('inf') if minimize else -float('inf')
             if len(metrics) > 1:
                 error_marker = tuple([default_score] * len(metrics))
                 current_error = float('inf') # Or appropriate indicator for multi-objective
             else:
                 current_error = default_score
                 error_marker = default_score
             return {
                 'current_scores': [default_score] * len(metrics),
                 'any_model': None,
                 'eval_time': eval_time,
                 'current_error': current_error,
                 'error_marker': error_marker,
                 'subset_key': tuple(subset)
             }

        # Ensure subset indices are valid
        if not all(0 <= i < n_feats for i in subset):
            print(f"Warning: Invalid feature indices in subset {subset}. Skipping evaluation.")
            eval_time = (datetime.datetime.now() - start).total_seconds()
            # Return default poor score
            default_score = float('inf') if minimize else -float('inf')
            if len(metrics) > 1:
                 current_error = float('inf')
                 error_marker = tuple([default_score] * len(metrics))
            else:
                 current_error = default_score
                 error_marker = default_score
            return {
                 'current_scores': [default_score] * len(metrics),
                 'any_model': None,
                 'eval_time': eval_time,
                 'current_error': current_error,
                 'error_marker': error_marker,
                 'subset_key': tuple(subset)
             }

        X_subset = X[:, subset]

        current_scores = []
        any_model = None  # Holder for a model returned by a metric

        if len(metrics) > 1:
            tuple_metric = []
            for metric in metrics:
                res = metric(X_subset, y, sample_weight=sample_weight)
                score = float(res["score"])
                current_scores.append(score)
                if "model" in res and any_model is None:
                    any_model = res["model"]
                # Adjust for direction: higher is better after this line
                adjusted = -score if (isinstance(minimize, bool) and minimize) else score
                tuple_metric.append(adjusted)

            tuple_metric = tuple(tuple_metric)

            # NEW: use a simple scalarization proxy (mean of adjusted objectives)
            # -> higher mean = better; we negate to keep "lower is better" internal error
            current_error = -float(np.mean(tuple_metric))
            error_marker = tuple_metric
        else: # Single metric
            res = metrics[0](X_subset, y, sample_weight=sample_weight)
            score = float(res["score"])
            current_scores.append(score)
            any_model = res.get("model", None)
            current_error = -score if not minimize else score
            error_marker = current_error

        eval_time = (datetime.datetime.now() - start).total_seconds()

        return {
            'current_scores': current_scores,
            'any_model': any_model,
            'eval_time': eval_time,
            'current_error': current_error,
            'error_marker': error_marker,
            'subset_key': tuple(subset)
        }



    @staticmethod
    def calculate_hypervolume_static(points_or_point):
        """
        Very simple/tolerant placeholder:
        - Accepts either a list/iterable of points (tuples of floats) OR a single point (tuple of floats).
        - Returns a cheap proxy (count of non-dominated points) to avoid complex HV logic.
        """

        # Normalize input to a list of points
        pts = points_or_point
        # If given a single point tuple like (a, b, ...), wrap it
        if isinstance(pts, tuple) and all(isinstance(v, (int, float, np.number)) for v in pts):
            pts = [pts]
        # If given a single list/array of numbers, wrap it
        if isinstance(pts, (list, np.ndarray)) and pts and isinstance(pts[0], (int, float, np.number)):
            pts = [tuple(pts)]

        if not pts:
            return 0.0

        # Ensure list of tuples
        clean = []
        for p in pts:
            p = tuple(map(float, p))
            clean.append(p)

        # Compute number of non-dominated points as a cheap stable proxy
        def dominates(a, b):
            # a dominates b if a is >= b on all objectives and > on at least one (maximization view)
            ge_all = all(x >= y for x, y in zip(a, b))
            gt_any = any(x > y for x, y in zip(a, b))
            return ge_all and gt_any

        non_dominated = []
        for i, p in enumerate(clean):
            if not any(dominates(q, p) for j, q in enumerate(clean) if j != i):
                non_dominated.append(p)

        return float(len(non_dominated))


    def _update_after_evaluation(self, result, subset, previous_subset):
      """
      Update selector state after evaluating a subset.

      - Propagates per-feature scores using the primary metric.
      - Optionally applies a Boruta-style penalty per iteration:
        build shadow features (2*sqrt(k)), fit an augmented model on [X_sub | X_shadow],
        compare real feature importances (from that augmented fit) to the max shadow importance,
        and subtract (shadow_max * |primary_score|) for any feature that loses.
      """
      from sklearn.base import clone as _sk_clone

      # ---- Unpack evaluation result
      current_scores   = result.get('current_scores')
      any_model        = result.get('any_model')
      eval_time        = result.get('eval_time', 0.0)
      self.current_error = result.get('current_error')
      self.error_marker  = result.get('error_marker')
      subset_key       = result.get('subset_key')

      # ---- Bookkeeping
      if subset_key is not None:
          self.dup_dct[subset_key] += 1
      self.performance_history.append(self.current_error)
      self.cost_history.append(eval_time)

      # ---- Build feature-specific scores if we have a score payload
      if current_scores:
          # Primary score can be list/tuple or dict (prefers 'score')
          if isinstance(current_scores, dict):
              primary_score = current_scores.get('score', next(iter(current_scores.values())))
          else:
              primary_score = current_scores[0]
          primary_score = float(primary_score)

          # Extract importances from the model we just evaluated (non-augmented)
          model_importances = self._extract_importances(any_model) if any_model is not None else None

          # Base propagation: importance * primary_score, else uniform primary_score
          if (model_importances is not None) and (len(model_importances) == len(subset)):
              weighted_importances = np.asarray(model_importances, float).ravel() * primary_score
              feature_scores = dict(zip(subset, weighted_importances))
          else:
              feature_scores = {feat: primary_score for feat in subset}

          # ---- Boruta-style penalty (optional)
          if self.boruta and len(subset) > 0:
              try:
                  # k = number of real features under test this iteration
                  n_shadow = int(max(1, 2 * np.sqrt(len(subset))))

                  X_sub = self.X[:, subset]
                  X_shadow = self._make_shadow_features(X_sub, n_shadow, self._rng)

                  if X_shadow.shape[1] > 0 and any_model is not None:
                      est_shadow = _sk_clone(any_model)
                      X_aug = np.hstack([X_sub, X_shadow])

                      # Thread sample_weight if available
                      fit_kwargs = {}
                      sw = getattr(self, "sample_weight", None)
                      if sw is not None:
                          fit_kwargs["sample_weight"] = sw

                      est_shadow.fit(X_aug, self.y, **fit_kwargs)

                      imp_aug = self._extract_importances(est_shadow)
                      if imp_aug is not None and imp_aug.size == X_aug.shape[1]:
                          real_imps_aug = np.asarray(imp_aug[:len(subset)], float)
                          shadow_imps   = np.asarray(imp_aug[-n_shadow:], float)

                          if shadow_imps.size > 0 and np.isfinite(shadow_imps).any():
                              shadow_max = float(np.nanmax(shadow_imps))
                              penalty = shadow_max * abs(primary_score)

                              weak_mask = (real_imps_aug < shadow_max)
                              if np.any(weak_mask):
                                  for local_idx, is_weak in enumerate(weak_mask):
                                      if is_weak:
                                          feat = subset[local_idx]
                                          feature_scores[feat] = feature_scores.get(feat, 0.0) - penalty
              except Exception:
                  # Fail-open: skip Boruta for this iteration if anything goes sideways
                  pass

          # Commit per-feature updates
          self.update_feature_performance(feature_scores)

      # ---- Multi-metric Pareto frontier (if applicable)
      if len(self.metrics) > 1:
          self.update_pareto_history(self.error_marker)

      # ---- Keep last produced model for inspection
      if any_model is not None:
          self.last_metric_model_ = any_model

      # ---- Drive the search state machine and housekeeping
      self.update_search(previous_subset, subset, eval_time)
      for f in subset:
          self.unevaluated.discard(f)


    def update_search(self, previous_subset, current_subset, time_cost):
        # Ensure cost history and performance history are updated before calling this
        # (They are updated in evaluate_subset)

        # Update feature cost dynamics
        self.update_feature_cost_history(previous_subset, current_subset, time_cost)

        # Update fastest evaluation time and corresponding feature (if subset not empty)
        # Ensure placeholder_coefficient is initialized properly in fit()
        if time_cost < self.placeholder_coefficient:
             self.placeholder_coefficient = time_cost
             # Store the first feature of the fastest subset as a heuristic
             self.fastest_feat = current_subset[0] if current_subset else self.fastest_feat # Keep old if current is empty

        # Increment iteration counter
        self.iters += 1

        # Check for improvement (using self.current_error where lower is better)
        # Ensure self.best_error is initialized properly in fit()
        if self.current_error < self.best_error:
            self.iters_best = self.iters # Record iteration of best found
             # Use self.error_marker which stores the actual performance value(s)
            self.leaderboard.append((self.error_marker, current_subset))
             # Sort leaderboard if needed (e.g., keep top K) - simple append for now
            self.leaderboard.sort(key=lambda x: x[0] if isinstance(x[0], (int, float, np.number)) else sum(x[0])) # Simple sort key
            self.leaderboard = self.leaderboard[:50] # Limit leaderboard size
            self.best_error = self.current_error # Update best comparable error

             # Update feature stability after leaderboard change
            self.update_feature_stability()


        # Update exploration parameters (self.c, self.kl) - based on FLAML/related algorithms?
        # self.kl often relates to cost or time
        self.kl = max(1e-9, time_cost) # Use current evaluation time/cost

        # self.c balances exploration/exploitation, might depend on subset size
        current_subset_size = len(current_subset) if current_subset else 1
        # Ensure c_sub and n_feats are valid
        c_sub = self.c_sub if hasattr(self, 'c_sub') else 4
        n_feats = self.n_feats if hasattr(self, 'n_feats') else 1
        # Avoid division by zero, ensure c is reasonable
        self.c = max(1.0, min(c_sub, n_feats / max(1, current_subset_size)))


    def adjust_step_size(self, current_features_len, remaining_budget_fraction):
        # Ensure current_features_len is at least 1 for sqrt
        base_step = max(1, int(np.sqrt(max(1, current_features_len))))

        # Ensure iters and iters_best are initialized
        iters = self.iters if hasattr(self, 'iters') else 0
        iters_best = self.iters_best if hasattr(self, 'iters_best') else 0

        no_improvement_steps = iters - iters_best
        total_steps = iters # Corrected from 'self yüzde'

        # Avoid division by zero for improvement rate
        improvement_rate = iters_best / total_steps if total_steps > 0 else 0.0 # Rate of finding best

        # Use recent cost history for adaptation
        cost_hist = self.cost_history if hasattr(self, 'cost_history') else []
        recent_costs = cost_hist[-10:] if len(cost_hist) > 10 else cost_hist
        avg_recent_cost = np.mean(recent_costs) if recent_costs else 1e-3 # Use small default if no history
        avg_recent_cost = max(1e-9, avg_recent_cost) # Ensure positive

        # Factors influencing step size
        exploitation_factor = 1.0 - improvement_rate # Higher if stuck
        budget_factor = max(0.1, min(1.0, remaining_budget_fraction)) # Scale between 0.1 and 1

        # Dynamic step size calculation
        if no_improvement_steps > 5: # Increase exploration if stuck for a while
             exploration_term = np.log1p(no_improvement_steps) / 2 # Log scale increase
             # Increase step size more aggressively when stuck, tempered by budget
             step_size = base_step * (1 + exploration_term * exploitation_factor * budget_factor)
        else: # Focus more on exploitation/refinement if recently improved
             # Smaller steps, influenced slightly by budget
             step_size = base_step * (0.8 + 0.2 * budget_factor)

        # Ensure step size is at least 1 and not excessively large
        n_feats = self.n_feats if hasattr(self, 'n_feats') else 10 # Default if not set
        step_size = max(1, min(int(np.ceil(step_size)), max(1, n_feats // 4))) # Limit to 1/4th of features

        # Reduce step size drastically when budget is very low for fine-tuning
        if remaining_budget_fraction < 0.1:
             step_size = min(step_size, max(1, int(n_feats * 0.05))) # Limit to 5% of features or 1

        return step_size

    def search_strategy(self, n_target_feats):
        # Ensure feature weights/priors are initialized
        if not hasattr(self, 'feature_priors') or self.feature_priors is None:
            # Initialize with uniform priors if not provided or calculated
            self.feature_priors = np.ones(self.n_feats) / self.n_feats

        if not hasattr(self, 'feature_performance') or self.feature_performance is None:
            self.feature_performance = np.zeros(self.n_feats)

        if not hasattr(self, 'feature_stability'):
            self.feature_stability = np.zeros(self.n_feats)


        weights = np.ones(self.n_feats) * 1e-9 # Start with small positive weights

        # Combine signals: priors, performance, stability, cost
        iters = self.iters if hasattr(self, 'iters') else 0

        # Performance weight increases with iterations (exploitation)
        # Ensure feature_performance is normalized (0 to 1 ideally)
        perf_norm = self.feature_performance # Assume already normalized (0=worst, 1=best)
        # Simple linear weight increase for exploitation, capped at 0.7
        exploit_weight = min(0.7, iters / 100.0) if iters > 10 else 0.0

        # Stability weight (use if available and stable)
        stab_weight = 0.1 if np.sum(self.feature_stability) > 0 else 0.0

        # Prior weight (exploration, decreases initially)
        prior_weight = max(0.1, 1.0 - exploit_weight - stab_weight)

        # Combine base weights
        base_weights = (prior_weight * self.feature_priors +
                        exploit_weight * perf_norm +
                        stab_weight * self.feature_stability)


        # Factor in feature cost (lower cost = higher weight)
        if hasattr(self, 'feature_cost_history') and self.feature_cost_history:
            costs = np.array([self.feature_cost_history.get(feat, 1.0) for feat in range(self.n_feats)])
             # Add small epsilon to avoid division by zero, dampen effect with sqrt
            cost_factor = 1.0 / (np.sqrt(np.maximum(costs, 1e-9)) + 0.1)
            weights = base_weights * cost_factor
        else:
            weights = base_weights

        # Ensure weights are positive and finite
        weights = np.maximum(weights, 1e-9)
        weights[~np.isfinite(weights)] = 1e-9


        # Normalize weights to get probabilities
        total_weight = np.sum(weights)
        if total_weight > 1e-9:
             probabilities = weights / total_weight
        else:
             probabilities = np.ones(self.n_feats) / self.n_feats # Fallback to uniform


        n_target_feats = max(1, min(n_target_feats, self.n_feats))

        if self.unevaluated:                       # coverage mode
            k_unseen   = int(np.ceil(0.20 * n_target_feats))           # 20 % of the subset
            k_unseen   = min(k_unseen, len(self.unevaluated))
            must_pick  = random.sample(list(self.unevaluated), k_unseen)
            remaining  = n_target_feats - k_unseen


            if remaining > 0:
              pool = np.setdiff1d(np.arange(self.n_feats), must_pick)

              # 1️⃣ slice
              probs_pool = probabilities[pool]
              # 2️⃣ renormalise
              probs_pool = probs_pool / probs_pool.sum()

              others = np.random.choice(pool,
                                        size=remaining,
                                        replace=False,
                                        p=probs_pool)
              selected_feats = np.concatenate([must_pick, others])
            else:
              selected_feats = np.array(must_pick, dtype=int)

        else:                                      # normal mode
            selected_feats = np.random.choice(self.n_feats,
                                              size=n_target_feats,
                                              replace=False,
                                              p=probabilities)

        return sorted(selected_feats.tolist())


    def select_best_feature_subset(self, budget):
        start_time = datetime.datetime.now()
        end_time = start_time + datetime.timedelta(seconds=budget)
        # Allocate small portion of budget for refinement at the end
        refinement_duration = budget * 0.1
        main_search_end = end_time - datetime.timedelta(seconds=refinement_duration)

        # Initialize performance/counts if not already done (e.g., in fit)
        if not hasattr(self, 'feature_performance'): self.feature_performance = np.zeros(self.n_feats)
        if not hasattr(self, 'feature_counts'): self.feature_counts = np.zeros(self.n_feats)

        # Initial subset size heuristic
        if not hasattr(self, 'num_features') or not self.fitted:
            self.num_features = max(1, min(int(np.sqrt(self.n_feats)), self.n_feats // 2, 10)) # Start small/medium

        # Generate and evaluate initial subset
        current_subset = self.search_strategy(self.num_features)
        result = self._compute_subset_score(current_subset, self.X, self.y, self.sample_weight, self.metrics, self.minimize, self.n_feats)
        self._update_after_evaluation(result, current_subset, [])

        # Initialize best error if first run (using lower-is-better convention)
        if not self.fitted:
            self.best_error = self.current_error if hasattr(self, 'current_error') else float('inf')
            self.iters_best = 0 # Initialize best iteration tracker

        no_improvement_counter = 0
        batch_size = max(1, self.n_jobs)

        # Main search loop
        with parallel_backend('loky'):
            while datetime.datetime.now() < main_search_end:
                remaining_time = (main_search_end - datetime.datetime.now()).total_seconds()
                if remaining_time <= 0: break # Exit if time is up
                remaining_budget_fraction = max(0.0, remaining_time / max(1e-9, (main_search_end - start_time).total_seconds()))

                # Generate batch of candidates
                batch_subsets = []
                for _ in range(batch_size):
                    step_size = self.adjust_step_size(self.num_features, remaining_budget_fraction)

                    # Determine change in subset size (add/remove features)
                    # More random exploration early or if stuck
                    if self.iters < 10 or no_improvement_counter > 5:
                        # Larger random steps initially or when stuck
                        step_change = random.choice([-2, -1, 1, 2]) * step_size
                    else:
                        # Smaller steps, allow staying same size
                        step_change = random.choice([-1, 0, 1]) * step_size

                    # Calculate new target size, ensuring it's within bounds
                    new_num_features = max(1, min(self.n_feats, self.num_features + int(np.round(step_change))))

                    # Generate new candidate subset
                    new_subset = self.search_strategy(new_num_features)
                    batch_subsets.append(new_subset)

                # Evaluate batch in parallel if n_jobs > 1
                if self.n_jobs > 1:
                    parallel = Parallel(n_jobs=self.n_jobs)
                    batch_results = parallel(delayed(self._compute_subset_score)(
                        sub, self.X, self.y, self.sample_weight, self.metrics, self.minimize, self.n_feats
                    ) for sub in batch_subsets)
                else:
                    batch_results = [self._compute_subset_score(
                        sub, self.X, self.y, self.sample_weight, self.metrics, self.minimize, self.n_feats
                    ) for sub in batch_subsets]

                # Process results sequentially
                improved = False
                best_in_batch_error = float('inf')
                best_in_batch_subset = None
                prev_error = self.current_error

                for idx, result in enumerate(batch_results):
                    subset = batch_subsets[idx]
                    self._update_after_evaluation(result, subset, current_subset)

                    if self.current_error < prev_error:
                        improved = True
                        if self.current_error < best_in_batch_error:
                            best_in_batch_error = self.current_error
                            best_in_batch_subset = subset

                if improved:
                    current_subset = best_in_batch_subset[:]
                    self.num_features = len(current_subset)
                    no_improvement_counter = 0
                else:
                    no_improvement_counter += batch_size  # Increment by batch size

                # Probabilistic restart if stuck for too long
                restart_prob = min(0.1, no_improvement_counter / 50.0) # Low probability restart
                if random.random() < restart_prob and self.iters > 20: # Avoid restarting too early
                    print("Performing probabilistic restart...")
                    # Restart near a good known size or explore a different size randomly
                    if self.leaderboard:
                         # Choose size from top 5 leaderboard entries
                         good_sizes = [len(s) for _, s in self.leaderboard[:min(5, len(self.leaderboard))]]
                         restart_features = random.choice(good_sizes) if good_sizes else self.num_features
                    else:
                         # Randomly perturb current size
                         restart_features = max(1, min(self.n_feats,
                                                  int(np.random.normal(self.num_features, self.n_feats / 5))))

                    # Generate and evaluate restart subset
                    current_subset = self.search_strategy(restart_features)
                    self.num_features = len(current_subset)
                    result = self._compute_subset_score(current_subset, self.X, self.y, self.sample_weight, self.metrics, self.minimize, self.n_feats)
                    self._update_after_evaluation(result, current_subset, current_subset)  # Previous is itself for restart
                    no_improvement_counter = 0

        # --- Refinement Phase ---
        # Focus search around the best solution found so far using local moves
        refine_start_time = datetime.datetime.now()
        with parallel_backend('loky'):
            while datetime.datetime.now() < end_time:
                 if not self.leaderboard: break # Exit if no solution found

                 # Get the current best subset from leaderboard
                 # Sort leaderboard by error (lower is better)
                 self.leaderboard.sort(key=lambda x: x[0] if isinstance(x[0], (int, float, np.number)) else sum(x[0]))
                 best_error_refine, best_subset_refine = self.leaderboard[0]
                 best_subset_refine = list(best_subset_refine) # Ensure it's a mutable list

                 # Rank features not in the best subset (candidates to add)
                 perf_norm = self.feature_performance # Assumed normalized (higher=better)
                 candidates_add = sorted([f for f in range(self.n_feats) if f not in best_subset_refine],
                                         key=lambda f: perf_norm[f], reverse=True)

                 # Rank features in the best subset (candidates to remove)
                 candidates_remove = sorted(best_subset_refine, key=lambda f: perf_norm[f])

                 # Generate batch of refinement actions
                 batch_subsets_refine = []
                 for _ in range(batch_size):
                     action = random.choice(['add', 'remove', 'swap'])

                     new_subset_refine = None

                     if action == 'add' and candidates_add and len(best_subset_refine) < self.n_feats:
                         # Add one of the top-k candidates
                         k = max(1, min(3, len(candidates_add)//3))
                         feat_to_add = random.choice(candidates_add[:k])
                         new_subset_refine = best_subset_refine + [feat_to_add]

                     elif action == 'remove' and len(best_subset_refine) > 1:
                         # Remove one of the worst-k candidates
                         k = max(1, min(3, len(candidates_remove)//3))
                         feat_to_remove = random.choice(candidates_remove[:k])
                         new_subset_refine = [f for f in best_subset_refine if f != feat_to_remove]

                     elif action == 'swap' and candidates_add and len(best_subset_refine) > 0:
                         # Swap one of the worst-k out with one of the best-k in
                         k_out = max(1, min(3, len(candidates_remove)//3))
                         feat_to_remove = random.choice(candidates_remove[:k_out])
                         k_in = max(1, min(3, len(candidates_add)//3))
                         feat_to_add = random.choice(candidates_add[:k_in])
                         temp_subset = [f for f in best_subset_refine if f != feat_to_remove]
                         new_subset_refine = temp_subset + [feat_to_add]

                     if new_subset_refine is not None:
                         batch_subsets_refine.append(sorted(list(set(new_subset_refine))))
                     else:
                         batch_subsets_refine.append(best_subset_refine)  # Fallback

                 # Evaluate batch in parallel
                 if self.n_jobs > 1:
                     parallel = Parallel(n_jobs=self.n_jobs)
                     batch_results = parallel(delayed(self._compute_subset_score)(
                         sub, self.X, self.y, self.sample_weight, self.metrics, self.minimize, self.n_feats
                     ) for sub in batch_subsets_refine)
                 else:
                     batch_results = [self._compute_subset_score(
                         sub, self.X, self.y, self.sample_weight, self.metrics, self.minimize, self.n_feats
                     ) for sub in batch_subsets_refine]

                 # Process results
                 for idx, result in enumerate(batch_results):
                     subset = batch_subsets_refine[idx]
                     self._update_after_evaluation(result, subset, best_subset_refine)


        # Final check to ensure leaderboard has at least one entry
        if not self.leaderboard:
            # If somehow leaderboard is empty, add the last evaluated subset
            last_subset = current_subset if 'current_subset' in locals() else []
            last_error = self.current_error if hasattr(self, 'current_error') else float('inf')
            last_marker = self.error_marker if hasattr(self, 'error_marker') else last_error
            if last_subset: # Only add if a subset was actually evaluated
                self.leaderboard.append((last_marker, last_subset))
                self.best_error = last_error


    def update_feature_stability(self):
        if not hasattr(self, 'feature_stability'):
            self.feature_stability = np.zeros(self.n_feats)

        # Update stability based on top K solutions in the leaderboard
        if self.leaderboard and self.n_feats > 0:
            stability_scores = np.zeros(self.n_feats)
            # Sort leaderboard by error (lower is better)
            # Use try-except for potentially mixed types in error marker during sorting
            try:
                sorted_leaderboard = sorted(self.leaderboard, key=lambda x: x[0] if isinstance(x[0], (int, float, np.number)) else sum(x[0]))
            except TypeError:
                 # Fallback if comparison fails (e.g., tuple vs float)
                 sorted_leaderboard = self.leaderboard # Use unsorted

            top_k = min(10, len(sorted_leaderboard)) # Consider top 10 solutions

            for i, (_, subset) in enumerate(sorted_leaderboard[:top_k]):
                weight = (top_k - i) / top_k # Linear weight decay (1.0 for best, lower for others)
                for feat in subset:
                    if 0 <= feat < self.n_feats: # Check index validity
                        stability_scores[feat] += weight

            # Normalize stability scores to be between 0 and 1
            max_stability = np.max(stability_scores)
            if max_stability > 0:
                self.feature_stability = stability_scores / max_stability
            else:
                self.feature_stability = np.zeros(self.n_feats) # Reset if all scores are zero

    def fit(self, X, y=None, sample_weight=None,budget=None):
        self.budget = budget if budget is not None else self.budget
        if self.budget is None:
             raise ValueError("Budget must be provided either during initialization or fit.")

        # Validate inputs using sklearn utilities
        X = check_array(X, ensure_2d=True, force_all_finite=True)
        y = check_array(y, ensure_2d=False, force_all_finite=True, ensure_min_samples=1)
        if X.shape[0] != y.shape[0]:
            raise ValueError("X and y must have the same number of samples.")

        self.X = X
        self.y = y
        self.sample_weight=sample_weight
        new_n_feats = self.X.shape[1]

        # --- State Initialization or Warm Start Check ---
        if not self.fitted or not hasattr(self, 'n_feats') or self.n_feats != new_n_feats:
            print("Initializing FLAVORS2 state...")
            self.n_feats = new_n_feats
            self.unevaluated = set(range(self.n_feats))
            self.leaderboard = []
            self.performance_history = []
            self.cost_history = []
            self.feature_cost_history = {}
            self.feature_addition_times = {}
            self.feature_removal_times = {}
            self.dup_dct = defaultdict(int)
            # Reset performance/stability tracking
            self.feature_performance = np.zeros(self.n_feats)
            self.feature_counts = np.zeros(self.n_feats)
            self.feature_stability = np.zeros(self.n_feats)
            self.min_performance = None # Reset min/max tracking
            self.max_performance = None

            # Initialize error tracking (lower is better convention)
            self.current_error = float('inf')
            self.best_error = float('inf')
            self.error_marker = None # Stores actual metric value(s)

            # Initialize counters and parameters
            self.kl = 1e-3 # Initial small cost estimate
            self.c = self.c_sub
            self.placeholder_coefficient = float('inf') # Tracks min eval time
            self.iters = 0
            self.iters_best = 0
            self.fastest_feat = None

            # Calculate unique counts per feature (used in cost estimation)
            # Handle potential memory issues for very wide datasets if needed
            self.unique_counts = np.array([len(np.unique(self.X[:, i])) for i in range(self.n_feats)])


            # Calculate feature priors (e.g., using mutual information)
            if self.feature_priors is None or len(self.feature_priors) != self.n_feats:
                print("...")
                # Check if classification or regression task based on y
                # Heuristic: if y has few unique values and is integer type
                if len(np.unique(self.y)) <= max(10, self.X.shape[0] * 0.1) and np.issubdtype(self.y.dtype, np.integer):
                    print("(Assuming classifiCalculating initial feature priors using mutual informationcation task for MI)")
                    # Added random_state for reproducibility
                    self.feature_priors = mutual_info_classif(self.X, self.y, discrete_features='auto', random_state=42, n_jobs=self.n_jobs)
                else:
                    print("(Assuming regression task for MI)")
                    self.feature_priors = mutual_info_regression(self.X, self.y, discrete_features='auto', random_state=42, n_jobs=self.n_jobs)

                # Normalize priors to be between 0 and 1 (optional, but good practice)
                min_prior, max_prior = np.min(self.feature_priors), np.max(self.feature_priors)
                if max_prior > min_prior:
                     self.feature_priors = (self.feature_priors - min_prior) / (max_prior - min_prior)
                else:
                     self.feature_priors = np.ones(self.n_feats) / self.n_feats # Uniform if MI is constant

            # Initialize multi-objective tracking if needed
            if len(self.metrics) > 1:
                print("Initializing Pareto front history for multi-objective optimization.")
                self.pareto_history = []
                # self.hypervolume_values = [] # Could track hypervolume over time

            self.fitted = False # Mark as not fully fitted until search completes

        else: # Warm start
             print("Performing warm start - continuing search...")
             # Assert shape compatibility for warm start
             assert new_n_feats == self.n_feats, \
             f"Input X has {new_n_feats} features, but model was previously fitted with {self.n_feats} features. Re-initialize for new data."
             # Reset counters? Or continue counting? Continuing for now.
             # Maybe reset no_improvement_counter?
             # self.iters_best = self.iters # Consider current point as 'best' for warm start phase?


        # --- Run the main search ---
        print(f"Starting feature selection search with budget: {self.budget} seconds...")
        self.select_best_feature_subset(self.budget)
        print("Search finished.")

        if not self.leaderboard:
             print("Warning: No feature subsets evaluated successfully or leaderboard is empty.")
             # Optionally: add a default subset (e.g., all features) if search fails completely
        else:
             # Sort final leaderboard
             try:
                  self.leaderboard.sort(key=lambda x: x[0] if isinstance(x[0], (int, float, np.number)) else sum(x[0]))
             except TypeError:
                  pass # Keep unsorted if comparison fails
             best_marker, best_subset_final = self.leaderboard[0]
             print(f"Best subset found: {best_subset_final} with score marker: {best_marker}")


        self.fitted = True # Mark as fitted after search completes
        return self


    def transform(self, X):
        if not self.fitted:
            raise RuntimeError("FLAVORS2 model has not been fitted yet. Call fit() first.")

        X = check_array(X, ensure_2d=True) # Validate input

        # Check for consistent number of features
        if X.shape[1] != self.n_feats:
             raise ValueError(f"Input X has {X.shape[1]} features, but model was fitted with {self.n_feats} features.")

        if not self.leaderboard:
             print("Warning: Leaderboard is empty after fitting. Returning original features.")
             return X

        # Get the best subset from the leaderboard
        # Ensure leaderboard is sorted (should be by end of fit)
        # best_marker, best_subset = min(self.leaderboard, key=lambda x: x[0] if isinstance(x[0], (int, float, np.number)) else sum(x[0]))
        best_marker, best_subset = self.leaderboard[0] # Assume sorted
        best_subset_indices = list(best_subset) # Ensure list of indices

        # Check if the best subset is empty (shouldn't happen with guards, but check)
        if not best_subset_indices:
            print("Warning: Best subset found is empty. Returning array with zero columns.")
            return np.zeros((X.shape[0], 0))

        # Return selected features
        return X[:, best_subset_indices]


from sklearn.base import BaseEstimator, TransformerMixin
import numpy as np
import pandas as pd

class FLAVORS2FeatureSelector(BaseEstimator, TransformerMixin):
    def __init__(self, budget=30, minimize=False, metrics=None, c_sub=4, feature_priors=None, n_jobs=1,
                 boruta=False, random_state=42):
        self.budget = budget
        self.minimize = minimize
        self.metrics = metrics
        self.c_sub = c_sub
        self.feature_priors = feature_priors
        self.n_jobs = n_jobs
        self.boruta = boruta
        self.random_state = random_state

        # sklearn-style fitted attributes (initialized here for clarity)
        self.selector = None
        self.selected_indices_ = None
        self.selected_indices = None
        self.n_features_in_ = None
        self.feature_names_in_ = None
        self._sample_weight = None

    def fit(self, X, y=None, sample_weight=None):
        # record input schema (sklearn convention)
        self.n_features_in_ = X.shape[1]
        if isinstance(X, pd.DataFrame):
            self.feature_names_in_ = X.columns.astype(str).to_numpy()

        self._sample_weight = sample_weight

        self.selector = FLAVORS2(
            budget=self.budget,
            minimize=self.minimize,
            metrics=self.metrics,
            c_sub=self.c_sub,
            feature_priors=self.feature_priors,
            n_jobs=self.n_jobs,
            boruta=self.boruta,
            random_state=self.random_state
        )
        # Fit underlying searcher (passes sample_weight through)
        self.selector.fit(X, y, sample_weight=sample_weight)

        # Persist the chosen subset in sklearn style
        if self.selector.leaderboard and len(self.selector.leaderboard) > 0:
            _, best_subset = self.selector.leaderboard[0]
            self.selected_indices_ = list(map(int, best_subset))
        else:
            self.selected_indices_ = list(range(X.shape[1]))  # fallback: all features

        # Back-compat alias (your older code referenced this)
        self.selected_indices = self.selected_indices_

        # (Optional) quick telemetry
        if hasattr(self.selector, "feature_counts"):
            unevaluated_indices = np.where(np.asarray(self.selector.feature_counts) == 0)[0]
            print(f"Number of feats that were never evaluated: {len(unevaluated_indices)}")

        return self

    def transform(self, X):
        # Use the underlying selector’s transform (keeps logic in one place),
        # but this also works if you prefer local slicing using selected_indices_
        return self.selector.transform(X)

    # sklearn helpers for compatibility
    def get_support(self, indices: bool = False):
        if self.selected_indices_ is None:
            raise AttributeError("The selector is not fitted yet.")
        if indices:
            return np.array(self.selected_indices_, dtype=int)
        mask = np.zeros(self.n_features_in_, dtype=bool)
        mask[self.selected_indices_] = True
        return mask

    def get_feature_names_out(self, input_features=None):
        if self.selected_indices_ is None:
            raise AttributeError("The selector is not fitted yet.")
        # If caller provides names, respect them; otherwise use feature_names_in_ if available
        if input_features is None:
            if self.feature_names_in_ is not None:
                input_features = list(self.feature_names_in_)
            else:
                input_features = [f"x{i}" for i in range(self.n_features_in_)]
        return [str(input_features[i]) for i in self.selected_indices_]

