# 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
from fastcore.test import test_fail

# %% auto 0
__all__ = ['CFExplanationResults', 'METRICS_CALLABLE', 'METRICS', 'DEFAULT_METRICS', 'Explanation', 'BaseGenerationStrategy',
           'IterativeGenerationStrategy', 'VmapGenerationStrategy', 'PmapGenerationStrategy', 'StrategyFactory',
           '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
class BaseGenerationStrategy:
    """Base class for mapping strategy."""
    
    def __call__(
        self, 
        fn: Callable, # Function to generate cf for a single input
        X: jnp.ndarray, # Input instances to be explained
        pred_fn: Callable[[Array], Array],
        **kwargs
    ) -> Array: # Generated counterfactual explanations
        raise NotImplementedError

# %% ../nbs/06_evaluate.ipynb 9
class IterativeGenerationStrategy(BaseGenerationStrategy):
    """Iterativly generate counterfactuals."""

    def __call__(
        self, 
        fn: Callable, # Function to generate cf for a single input
        X: jnp.ndarray, # Input instances to be explained
        pred_fn: Callable[[Array], Array],
        **kwargs
    ) -> Array: # Generated counterfactual explanations
        
        assert X.ndim == 2
        cfs = jnp.stack([fn(X[i], pred_fn=pred_fn, **kwargs) for i in range(X.shape[0])])
        assert X.shape == cfs.shape
        return cfs

# %% ../nbs/06_evaluate.ipynb 10
class VmapGenerationStrategy(BaseGenerationStrategy):
    """Generate counterfactuals via `jax.vmap`."""

    def __call__(
        self, 
        fn: Callable, # Function to generate cf for a single input
        X: jnp.ndarray, # Input instances to be explained
        pred_fn: Callable[[Array], Array],
        **kwargs
    ) -> Array: # Generated counterfactual explanations
        
        assert X.ndim == 2
        partial_fn = partial(fn, pred_fn=pred_fn, **kwargs)
        cfs = jax.vmap(partial_fn)(X)
        return cfs

# %% ../nbs/06_evaluate.ipynb 11
def _pad_divisible_X(
    X: Array,
    n_devices: int
):
    """Pad `X` to be divisible by `n_devices`."""
    if X.shape[0] % n_devices != 0:
        pad_size = n_devices - X.shape[0] % n_devices
        X = jnp.concatenate([X, jnp.zeros((pad_size, *X.shape[1:]))])
    X_padded = X.reshape(n_devices, -1, *X.shape[1:])
    return X_padded

# %% ../nbs/06_evaluate.ipynb 13
class PmapGenerationStrategy(BaseGenerationStrategy):
    def __init__(
        self, 
        n_devices: int = None, # Number of devices. If None, use all available devices
        strategy: str = 'auto', # Strategy to generate counterfactuals
        **kwargs
    ):
        self.strategy = strategy
        self.n_devices = n_devices or jax.device_count()

    def __call__(
        self, 
        fn: Callable, # Function to generate cf for a single input
        X: jnp.ndarray, # Input instances to be explained
        pred_fn: Callable[[Array], Array],
        **kwargs
    ) -> Array: # Generated counterfactual explanations
        
        assert X.ndim == 2
        X_padded = _pad_divisible_X(X, self.n_devices)
        partial_fn = partial(fn, pred_fn=pred_fn, **kwargs)
        cfs = jax.pmap(jax.vmap(partial_fn))(X_padded)
        cfs = cfs.reshape(-1, *cfs.shape[2:])
        cfs = cfs[:X.shape[0]]
        return cfs

# %% ../nbs/06_evaluate.ipynb 19
class StrategyFactory(object):
    """Factory class for Parallelism Strategy."""

    __strategy_map = {
        'iter': IterativeGenerationStrategy(),
        'vmap': VmapGenerationStrategy(),
        'pmap': PmapGenerationStrategy(),
    }

    def __init__(self) -> None:
        raise ValueError("This class should not be instantiated.")
        
    @staticmethod
    def get_default_strategy() -> BaseGenerationStrategy:
        """Get default strategy."""
        return VmapGenerationStrategy()

    @classmethod
    def get_strategy(cls, strategy: str | BaseGenerationStrategy) -> BaseGenerationStrategy:
        """Get strategy."""
        if isinstance(strategy, BaseGenerationStrategy):
            return strategy
        elif isinstance(strategy, str) and strategy in cls.__strategy_map:
            return cls.__strategy_map[strategy]
        else:
            raise ValueError(f"Invalid strategy: {strategy}")

# %% ../nbs/06_evaluate.ipynb 22
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,
    pred_fn=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, pred_fn=pred_fn)
    return cf_module

# %% ../nbs/06_evaluate.ipynb 23
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 24
def generate_cf_explanations(
    cf_module: BaseCFModule, # CF Explanation Module
    datamodule: TabularDataModule, # Data Module
    pred_fn: callable = None, # Predictive function
    strategy: str | BaseGenerationStrategy = 'vmap', # Parallelism Strategy for generating CFs
    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)

    # 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)

    if isinstance(cf_module, BaseParametricCFModule):
        cf_module = _train_parametric_module(
            cf_module, datamodule, t_configs=t_configs, pred_fn=pred_fn
        )
    X, _ = datamodule.test_dataset[:]

    strategy = StrategyFactory.get_strategy(strategy)
    current_time = time.time()
    cfs = strategy(cf_module.generate_cf, 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 32
class BaseEvalMetrics(ABC):
    """Base evaluation metrics class."""

    def __init__(self, name: str = None):
        if name is None: name = type(self).__name__
        self.name = name

    def __str__(self) -> str:
        has_name = hasattr(self, 'name')
        if not has_name:
            raise ValidationError(
                "EvalMetrics must have a name. Add the following as the first line in your "
                f"__init__ method:\n\nsuper({self.__name__}, self).__init__()")
        return self.name

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

# %% ../nbs/06_evaluate.ipynb 33
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 35
class PredictiveAccuracy(BaseEvalMetrics):
    """Compute the accuracy of the predict function."""
    
    def __init__(self, name: str = "accuracy"):
        super().__init__(name=name)

    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 37
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 39
class Validity(BaseEvalMetrics):
    """Compute fraction of input instances on which CF explanation methods output valid CF examples."""
    
    def __init__(self, name: str = "validity"):
        super().__init__(name=name)
    
    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 40
def _compute_proximity(
    inputs: jnp.DeviceArray, # input dim: [N, k]
    cfs: jnp.DeviceArray, # cfs dim: [N, k]
):
    prox = jnp.linalg.norm(inputs - cfs, ord=1, axis=1).mean()
    return prox.item()

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

# %% ../nbs/06_evaluate.ipynb 43
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 44
class Sparsity(BaseEvalMetrics):
    """Compute the number of feature changes between input datasets and CF examples."""

    def __init__(self, name: str = "sparsity"):
        super().__init__(name=name)
    
    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 45
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 46
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, name: str = "manifold_dist"):
        super().__init__(name=name)
        self.n_neighbors = n_neighbors
        self.p = p
        
    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 47
class Runtime(BaseEvalMetrics):
    """Get the running time to generate CF examples."""
    def __init__(self, name: str = "runtime"):
        super().__init__(name=name)
    
    def __call__(self, cf_explanations: Explanation) -> float:
        return cf_explanations.total_time

# %% ../nbs/06_evaluate.ipynb 49
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 51
def fake_explanations():
    """Generate sudo explanations for testing."""
    from relax.data import load_data

    dm = load_data("adult")
    X, y = dm.test_dataset[:]
    cfs = X
    dn = dm.data_name
    pred_fn = lambda x: jax.random.bernoulli(jax.random.PRNGKey(0), 0.5, (x.shape[0], 1)).astype(float)
    assert y.shape == pred_fn(X).shape
    return Explanation(
        cf_name='sudo', data_module=dm, cfs=cfs, pred_fn=pred_fn, total_time=0.0, dataset_name=dn
    )


# %% ../nbs/06_evaluate.ipynb 52
# 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
# )

METRICS_CALLABLE = [
    PredictiveAccuracy('acc'),
    PredictiveAccuracy('accuracy'),
    Validity(),
    Proximity(),
    Runtime(),
    ManifoldDist(),
]

METRICS = { m.name: m for m in METRICS_CALLABLE }

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

# %% ../nbs/06_evaluate.ipynb 54
def _get_metric(metric: str | BaseEvalMetrics, cf_exp: Explanation):
    if isinstance(metric, str):
        if metric not in METRICS.keys():
            raise ValueError(f"'{metric}' is not supported. Must be one of {METRICS.keys()}")
        res = METRICS[metric](cf_exp)
    elif callable(metric):
        # f(cf_exp) not supported for now
        if not isinstance(metric, BaseEvalMetrics):
            raise ValueError(f"metric needs to be a subclass of `BaseEvalMetrics`.")
        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 56
def evaluate_cfs(
    cf_exp: Explanation, # CF Explanations
    metrics: Iterable[Union[str, BaseEvalMetrics]] = 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:
        metric_name = str(metric)
        result_dict[(data_name, cf_name)][metric_name] = _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 58
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)
