from typing import Dict
import torch
from typing import Tuple


def output_transform_select_affinity_abs(
    output: Dict[str, torch.Tensor]
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Select predicted affinities (in absolute value) and experimental (target) affinities
    from output dictionary.

    Parameters
    ----------
    output: Dict[str, ignite.metrics.Metric]
        Engine output

    Returns
    -------
    Tuple[torch.Tensor, torch.Tensor]
        Predicted binding affinity (absolute value) and experimental binding affinity

    Notes
    -----
    This function is used as :code:`output_transform` in
    :class:`ignite.metrics.metric.Metric` and allow to select affinity predictions from
    the dictionary that the evaluator returns.

    Affinities can have negative values when they are associated to bad poses. The sign
    is used by :class:`AffinityLoss`, but in order to compute standard metrics the
    absolute value is needed, which is returned here.
    """
    # Return pose class probabilities and true labels
    return output["affinities_pred"], torch.abs(output["affinities"])