import numpy as np 

from hep_ml.reweight import BinsReweighter


def select_feature(feature: np.ndarray, limits: tuple[float, float]) -> tuple[np.ndarray, list]:
    selection_indices = []
    for index, value in enumerate(feature):
        if value > limits[0] and value < limits[1]:
            selection_indices.append(index)

    return feature[selection_indices], selection_indices


def reweight_feature(original_feature: list, target_feature: list, n_bins: int, n_neighs: int = 2):
    original_weights = np.ones(len(original_feature))
    bin_reweighter   = BinsReweighter(n_bins = n_bins, n_neighs = n_neighs)
    bin_reweighter.fit(original = original_feature, target = target_feature, original_weight = original_weights)

    return bin_reweighter.predict_weights(original = original_feature, original_weight = original_weights)