# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/06_evaluate.ipynb.

# %% ../nbs/06_evaluate.ipynb 3
from __future__ import annotations
from .import_essentials import *
from .trainer import TrainingConfigs
from .data import TabularDataModule
from .utils import accuracy, proximity
from .methods.base import BaseCFModule, BaseParametricCFModule, BasePredFnCFModule
from .methods.counternet import CounterNet
from copy import deepcopy
from sklearn.neighbors import NearestNeighbors

# %% auto 0
__all__ = ['CFExplanationResults', 'METRICS', 'DEFAULT_METRICS', 'Explanation', 'generate_cf_explanations', 'BaseEvalMetrics',
           'PredictiveAccuracy', 'Validity', 'Proximity', 'Sparsity', 'ManifoldDist', 'Runtime', 'compute_so_validity',
           'compute_so_proximity', 'compute_so_sparsity', 'evaluate_cfs', 'benchmark_cfs']

# %% ../nbs/06_evaluate.ipynb 4
@dataclass
class Explanation:
    """Generated CF Explanations class."""
    cf_name: str  # cf method's name
    data_module: TabularDataModule  # data module
    cfs: jnp.DeviceArray  # generated cf explanation of `X`
    total_time: float  # total runtime
    pred_fn: Callable[[jnp.DeviceArray], jnp.DeviceArray]  # predict function
    dataset_name: str = str()  # dataset name
    X: jnp.ndarray = None  # input
    y: jnp.ndarray = None  # label

    def __post_init__(self):
        if self.data_module:
            if self.dataset_name == str():
                self.dataset_name = self.data_module.data_name
            test_X, label = self.data_module.test_dataset[:]
            if self.X is None:
                self.X = test_X
            if self.y is None:
                self.y = label

CFExplanationResults = Explanation

# %% ../nbs/06_evaluate.ipynb 8
def _validate_configs(
    cf_module: BaseCFModule,
    datamodule: TabularDataModule,
    pred_fn: Callable[[jnp.DeviceArray], jnp.DeviceArray] = None,
    t_configs=None
):
    if (pred_fn is None) and (not isinstance(cf_module, BasePredFnCFModule)):
        warnings.warn(f"`{type(cf_module).__name__}` is not a subclass of `BasePredFnCFModule`."
            "This might cause problems as you set `pred_fn=None`, "
            f"which infers that `{type(cf_module).__name__}` has an attribute `pred_fn`.")


def _prepare_module(
    cf_module: BaseCFModule,
    datamodule: TabularDataModule
):
    cf_module.hook_data_module(datamodule)
    return cf_module

def _train_parametric_module(
    cf_module: BaseParametricCFModule,
    datamodule: TabularDataModule,
    t_configs=None
):
    if not cf_module._is_module_trained():
        print(f'{type(cf_module).__name__} contains parametric models. '
            'Starts training before generating explanations...')
        cf_module.train(datamodule, t_configs)
    return cf_module

# %% ../nbs/06_evaluate.ipynb 9
def _check_aux_pred_fn_args(pred_fn_args: dict | None):
    if pred_fn_args is None:
        return dict()
    elif isinstance(pred_fn_args, dict):
        return pred_fn_args
    else:
        raise ValueError(f'`pred_fn_args` should be a `dict`,',
            f'but got `{type(pred_fn_args).__name__}`')

class _AuxPredFn:
    def __init__(self, pred_fn, pred_fn_args: dict | None):
        self.pred_fn = pred_fn
        self.fn_args = deepcopy(_check_aux_pred_fn_args(pred_fn_args))

    def __call__(self, x: jnp.DeviceArray) -> jnp.DeviceArray:
        return self.pred_fn(x, **self.fn_args)


def _check_pred_fn(
    pred_fn: callable | None, 
    cf_module: BaseCFModule
) -> callable:
    if pred_fn is None:
        try:
            pred_fn = cf_module.pred_fn
        except AttributeError:
            raise AttributeError(
                    "`generate_cf_explanations` is incorrectly configured."
                    f"It is supposed to be `pred_fn != None`,"
                    f"or `{type(cf_module).__name__}` has attribute `pred_fn`."
                    f"However, we got `pred_fn={pred_fn}`, "
                    f"and `{type(cf_module).__name__}` has not attribute `pred_fn`."
            )
    elif isinstance(cf_module, BasePredFnCFModule):
        # override pred_fn if `cf_module` has `pred_fn`
        pred_fn = cf_module.pred_fn
    return pred_fn

# %% ../nbs/06_evaluate.ipynb 10
def generate_cf_explanations(
    cf_module: BaseCFModule, # CF Explanation Module
    datamodule: TabularDataModule, # Data Module
    pred_fn: callable = None, # Predictive function
    t_configs: TrainingConfigs = None, # training configs for `BaseParametricCFModule`
    pred_fn_args: dict = None # auxiliary arguments for `pred_fn` 
) -> Explanation:
    """Generate CF explanations."""

    _validate_configs(cf_module, datamodule, pred_fn, t_configs)
    cf_module = _prepare_module(cf_module, datamodule)

    if isinstance(cf_module, BaseParametricCFModule):
        cf_module = _train_parametric_module(
            cf_module, datamodule, t_configs=t_configs
        )
    X, _ = datamodule.test_dataset[:]
    
    # create `pred_fn` which only takes `x` as an input
    if pred_fn is not None:
        pred_fn = _AuxPredFn(pred_fn, pred_fn_args=pred_fn_args)

    # generate cfs
    current_time = time.time()
    cfs = cf_module.generate_cfs(X, pred_fn=pred_fn)
    total_time = time.time() - current_time
    # check pred_fn
    pred_fn = _check_pred_fn(pred_fn, cf_module)

    return Explanation(
        cf_name=cf_module.name,
        data_module=datamodule,
        cfs=cfs,
        total_time=total_time,
        pred_fn=pred_fn,
    )


# %% ../nbs/06_evaluate.ipynb 18
class BaseEvalMetrics(ABC):
    """Base evaluation metrics class."""
    def __str__(self) -> str:
        raise NotImplementedError

    def __call__(self, cf_explanations: Explanation) -> Any:
        raise NotImplementedError

# %% ../nbs/06_evaluate.ipynb 19
def _compute_acc(
    input: jnp.DeviceArray, # input dim: [N, k]
    label: jnp.DeviceArray, # label dim: [N] or [N, 1]
    pred_fn: Callable[[jnp.DeviceArray], jnp.DeviceArray]
) -> float:
    y_pred = pred_fn(input).reshape(-1, 1).round()
    label = label.reshape(-1, 1)
    return accuracy(y_pred, label).item()

# %% ../nbs/06_evaluate.ipynb 20
class PredictiveAccuracy(BaseEvalMetrics):
    """Compute the accuracy of the predict function."""
    def __str__(self) -> str:
        return "accuracy"

    def __call__(self, cf_explanations: Explanation) -> float:
        X, y = cf_explanations.data_module.test_dataset[:]
        return _compute_acc(X, y, cf_explanations.pred_fn)

# %% ../nbs/06_evaluate.ipynb 21
def _compute_val(
    input: jnp.DeviceArray, # input dim: [N, k]
    cfs: jnp.DeviceArray, # cfs dim: [N, k]
    pred_fn: Callable[[jnp.DeviceArray], jnp.DeviceArray]
):
    y_pred = pred_fn(input).reshape(-1, 1).round()
    y_prime = jnp.ones_like(y_pred) - y_pred
    cf_y = pred_fn(cfs).reshape(-1, 1).round()
    return accuracy(y_prime, cf_y).item()

# %% ../nbs/06_evaluate.ipynb 22
class Validity(BaseEvalMetrics):
    """Compute fraction of input instances on which CF explanation methods output valid CF examples."""
    def __str__(self) -> str:
        return "validity"
    
    def __call__(self, cf_explanations: Explanation) -> float:
        X, _ = cf_explanations.data_module.test_dataset[:]
        return _compute_val(
            X, cf_explanations.cfs, cf_explanations.pred_fn
        )

# %% ../nbs/06_evaluate.ipynb 23
class Proximity(BaseEvalMetrics):
    """Compute L1 norm distance between input datasets and CF examples divided by the number of features."""
    def __str__(self) -> str:
        return "proximity"
    
    def __call__(self, cf_explanations: Explanation) -> float:
        X, _ = cf_explanations.data_module.test_dataset[:]
        return proximity(X, cf_explanations.cfs).item()

# %% ../nbs/06_evaluate.ipynb 24
def _compute_spar(
    input: jnp.DeviceArray,
    cfs: jnp.DeviceArray,
    cat_idx: int
):
    # calculate sparsity
    cat_sparsity = proximity(input[:, cat_idx:], cfs[:, cat_idx:]) / 2
    cont_sparsity = jnp.linalg.norm(
        jnp.abs(input[:, :cat_idx] - cfs[:, :cat_idx]), ord=0, axis=1
    ).mean()
    return (cont_sparsity + cat_sparsity).item()


# %% ../nbs/06_evaluate.ipynb 25
class Sparsity(BaseEvalMetrics):
    """Compute the number of feature changes between input datasets and CF examples."""
    def __str__(self) -> str:
        return "sparsity"
    
    def __call__(self, cf_explanations: Explanation) -> float:
        X, _ = cf_explanations.data_module.test_dataset[:]
        return _compute_spar(X, cf_explanations.cfs, cf_explanations.cat_idx)

# %% ../nbs/06_evaluate.ipynb 26
def _compute_manifold_dist(
    input: jnp.DeviceArray,
    cfs: jnp.DeviceArray,
    n_neighbors: int = 1,
    p: int = 2
):
    knn = NearestNeighbors(n_neighbors=n_neighbors, p=p)
    knn.fit(input)
    nearest_dist, nearest_points = knn.kneighbors(cfs, 1, return_distance=True)
    return jnp.mean(nearest_dist).item()

# %% ../nbs/06_evaluate.ipynb 27
class ManifoldDist(BaseEvalMetrics):
    """Compute the L1 distance to the n-nearest neighbor for all CF examples."""
    def __init__(self, n_neighbors: int = 1, p: int = 2):
        self.n_neighbors = n_neighbors
        self.p = p
    
    def __str__(self) -> str:
        return "manifold_dist"
    
    def __call__(self, cf_explanations: Explanation) -> float:
        X, _ = cf_explanations.data_module.test_dataset[:]
        return _compute_manifold_dist(
            X, cf_explanations.cfs, self.n_neighbors, self.p
        )

# %% ../nbs/06_evaluate.ipynb 28
class Runtime(BaseEvalMetrics):
    """Get the running time to generate CF examples."""
    def __call__(self, cf_explanations: Explanation) -> float:
        return cf_explanations.total_time

# %% ../nbs/06_evaluate.ipynb 30
def _create_second_order_cfs(cf_results: CFExplanationResults, threshold: float = 2.0):
    X, y = cf_results.data_module.test_dataset[:]
    cfs = cf_results.cfs
    scaler = cf_results.data_module.normalizer
    cat_idx = cf_results.data_module.cat_idx

    # get normalized threshold = threshold / (max - min)
    data_range = scaler.data_range_
    thredshold_normed = threshold / data_range

    # select continous features
    x_cont = X[:, :cat_idx]
    cf_cont = cfs[:, :cat_idx]
    # calculate the diff between x and c
    cont_diff = jnp.abs(x_cont - cf_cont) <= thredshold_normed
    # new cfs
    cfs_cont_hat = jnp.where(cont_diff, x_cont, cf_cont)

    cfs_hat = jnp.concatenate((cfs_cont_hat, cfs[:, cat_idx:]), axis=-1)
    return cfs_hat


def compute_so_validity(cf_results: CFExplanationResults, threshold: float = 2.0):
    cfs_hat = _create_second_order_cfs(cf_results, threshold)
    cf_results_so = deepcopy(cf_results)
    cf_results_so.cfs = cfs_hat
    compute_validity = Validity()
    return compute_validity(cf_results_so)


def compute_so_proximity(cf_results: CFExplanationResults, threshold: float = 2.0):
    cfs_hat = _create_second_order_cfs(cf_results, threshold)
    cf_results_so = deepcopy(cf_results)
    cf_results_so.cfs = cfs_hat
    compute_proximity = Proximity()
    return compute_proximity(cf_results_so)


def compute_so_sparsity(cf_results: CFExplanationResults, threshold: float = 2.0):
    cfs_hat = _create_second_order_cfs(cf_results, threshold)
    cf_results_so = deepcopy(cf_results)
    cf_results_so.cfs = cfs_hat
    compute_sparsity = Sparsity()
    return compute_sparsity(cf_results_so)


# %% ../nbs/06_evaluate.ipynb 32
METRICS = dict(
    acc=PredictiveAccuracy(),
    accuracy=PredictiveAccuracy(),
    validity=Validity(),
    proximity=Proximity(),
    runtime=Runtime(),
    manifold_dist=ManifoldDist(),
    # validity=compute_so_validity,
    # so_proximity=compute_so_proximity,
    # so_sparsity=compute_so_sparsity
)

DEFAULT_METRICS = ["acc", "validity", "proximity"]

# %% ../nbs/06_evaluate.ipynb 33
def _get_metric(metric: str | callable, cf_exp: Explanation):
    if isinstance(metric, str):
        try:
            res = METRICS[metric](cf_exp)
        except KeyError:
            raise ValueError(f"'{metric}' is not supported. Must be one of {METRICS.keys()}")
    elif callable(metric):
        res = metric(cf_exp)
    else:
        raise ValueError(f"{type(metric).__name__} is not supported as a metric.")
    
    if isinstance(res, jnp.ndarray) and res.shape == (1,):
        res = res.item()
    return res

# %% ../nbs/06_evaluate.ipynb 34
def evaluate_cfs(
    cf_exp: Explanation, # CF Explanations
    metrics: Iterable[Union[str, Callable]] = None, # A list of Metrics. Can be `str` or a subclass of `BaseEvalMetrics`
    return_dict: bool = True, # return a dictionary or not (default: True)
    return_df: bool = False # return a pandas Dataframe or not (default: False)
):
    cf_name = cf_exp.cf_name
    data_name = cf_exp.data_module.data_name
    result_dict = { (data_name, cf_name): dict() }

    if metrics is None:
        metrics = DEFAULT_METRICS

    for metric in metrics:
        result_dict[(data_name, cf_name)][metric] = _get_metric(metric, cf_exp)
    result_df = pd.DataFrame.from_dict(result_dict, orient="index")
    
    if return_dict and return_df:
        return (result_dict, result_df)
    elif return_dict or return_df:
        return result_df if return_df else result_dict


# %% ../nbs/06_evaluate.ipynb 35
def benchmark_cfs(
    cf_results_list: Iterable[CFExplanationResults],
    metrics: Optional[Iterable[str]] = None,
):
    dfs = [
        evaluate_cfs(
            cf_exp=cf_results, metrics=metrics, return_dict=False, return_df=True
        )
        for cf_results in cf_results_list
    ]
    return pd.concat(dfs)
