### INCOMPLETE ###
import numpy as np
import pandas as pd


def generate_rashomon_set(models: np.array,
                          data: pd.DataFrame,
                          metric: callable,
                          epsilon: float = 0.2) -> set:
    """
    Genearte a Rashomon set using trained models by first finding the optimal model given a metric and data, 
    then adding every model performing within parameter epsilon worse than optimal on the same metric.

    Parameters:
        models: Array of fitted models that could be eligible for the Rashomon set.
        data: Data frame which will be used to evalute models in a metric.
        metric: Callable function that accepts (model, data) as parameters and returns a float score tending to 0.
        epsilon: Hyperparameter to determine acceptable error.
    
    Returns:
        Set containing models that fit the criteria of a Rashomon set.

    Example use:
        TO-DO!
    """
    scores = np.array([metric(model, data) for model in models])
    optimal_score = np.min(scores)
    model_in_set = scores <= optimal_score * (1 + epsilon)
    rashomon_set = set(models[model_in_set])
    return rashomon_set


def pdp(model, 
        data: pd.DataFrame, 
        feature: str, 
        feature_value: float) -> float:
    """
    Given a model for a predictive function, modify a specified feature to a given value in a copy of data to 
    calculate the mean over all data point variations.

    Parameters:
        model: Model that will be used for the prediction function.
        data: Data frame which will be copied and have a specific feature uniformally modified for PDP.
        feature: Specific feature column name in the data to modify.
        feature_value: Desired feature value for PDP calculation.

    Returns:
        Float value for the Partial Dependence Profile (PDP)

    Example use:
        TO-DO!
    """
    # Prepare calculation
    if feature not in data.columns:
        raise KeyError("PDP feature does not exist in given data.")
    modified_data = data.copy()
    modified_data[feature] = feature_value
    results = np.zeros(len(data))

    # Calculate PDP of given feature value
    for idx, sample in enumerate(modified_data):
        results[idx] = model.predict(sample)
    return np.mean(results)


def rashomon_pdp(models: np.array,
                 data: pd.DataFrame,
                 feature: str,
                 feature_value: float) -> float:
    """
    !!! INCOMPLETE !!!
    Description to be added based on the functionality and paper.
    Need to include Confidence Interval calculation here or in next function.

    Parameters:
        models: Rashomon set to be used in calculating the unified central trend.
        data: Data frame which be used to generate the PDPs in the models.
        feature: Specific feature column name in the data to modify.
        feature_value: Desired feature value for PDP calculation.
    
    Returns:
        Mean of PDPs from the Rashomon set models.
    
    Example use:
        TO-DO!
    """
    pdp_values = np.array([pdp(model, data, feature, feature_value) for model in models])
    return np.mean(pdp_values)  # Requires bootstrapping before completion
