# AUTOGENERATED! DO NOT EDIT! File to edit: ../001_module_stats.ipynb.

# %% auto 0
__all__ = ['bootstrap_sampling', 'compute_evaluation_estimates']

# %% ../001_module_stats.ipynb 4
from typing import Callable, List, Optional
import numpy as np
import pandas as pd

# %% ../001_module_stats.ipynb 6
def bootstrap_sampling(
    data: pd.DataFrame,  # Data containing the columns we want to generate bootstrap estimates from. 
    estimator: Callable = np.mean,  # estimator function that accepts an array-like argument. 
    n_boot: int = 1000,  # Number of bootstrap estimates to compute.     
    columns_to_exclude: List[str] = None  # Column names to exclude.
):
    "Compute bootstrap estimates of the data distribution"
    if not columns_to_exclude:
        columns_to_exclude = []
    data = data[[x for x in data.columns if x not in columns_to_exclude]]
    boot_dist = []
    for i in range(int(n_boot)):
        sample = data.sample(n=data.shape[0], replace=True)
        boot_dist.append(estimator(sample, axis=0))
    return pd.DataFrame(boot_dist)

# %% ../001_module_stats.ipynb 22
def compute_evaluation_estimates(
    df: pd.DataFrame,  # Evaluations per query data, usually obtained pyvespa evaluate method.
    n_boot: int = 1000,  # Number of bootstrap samples.  
    estimator: Callable = np.mean,  # estimator function that accepts an array-like argument. 
    quantile_low: float = 0.025,  # lower quantile to compute confidence interval
    quantile_high = 0.975  # upper quantile to compute confidence interval
):
    "Compute estimate and confidence interval for evaluation per query metrics."
    estimates = (
        df
        .groupby("model")
        .apply(bootstrap_sampling, 
               estimator = estimator, 
               n_boot = n_boot, 
               columns_to_exclude = ["query_id", "model"]
              )
        .reset_index(level="model")
        .groupby("model")
        .agg(
            [
                lambda x: np.quantile(x, q=quantile_low), 
                lambda x: np.quantile(x, q=0.5), 
                lambda x: np.quantile(x, q=quantile_high)
            ]
        )
        .rename(
            columns={
                "<lambda_0>": "low", 
                "<lambda_1>": "median", 
                "<lambda_2>": "high"
            }
        )
        .T
        .reset_index()
        .rename_axis(None, axis=1)
        .rename(
            columns={
                "level_0": "metric", "level_1": "quantile"
            }
        )
    )
    estimates = pd.melt(estimates, id_vars=["metric", "quantile"])    
    estimates = (
        pd.pivot(
            estimates,
            index=[
                x for x in estimates.columns if x not in ["quantile", "value"]
            ], 
            columns="quantile", values="value"
        )
        .reset_index()
        .rename_axis(None, axis=1)
        .rename(columns={"variable":"model"})[
            ["metric", "model", "low", "median", "high"]
        ]
    )
    return estimates
