from typing import Dict
from typing import Any


def extremise_distribution(distribution: Dict[Any, float],
                           tau: float = 1) -> Dict[Any, float]:
    """Calculate an extremised probability distribution parametrised by
    the temperature parameter `tau`.

    The extremised values are proportional to the values in
    `distribution` exponentiated to 1 / `tau`, where `tau` is a number
    between 1 and 0. Here, a value of 1 leaves the relative values
    unchanged and merely normalises the `distribution` whereas a value
    tending to 0 maximally extremises the `distribution`. In this case,
    the entry corresponding to the the highest value in the input
    distribution tends to 1 and the all the others tend to 0.

    Parameters
    ----------
    distribution
        A dictionary with values equal to positive floats.
    tau
        A parameter between 1 and 0 defining the reciprocal exponent or
        'temperature' of the extremised distribution.

    Returns
    -------
    dict:
        A probability distribution whose keys are equal to those
        of the input distribution and values are proportional to
        extremised (exponentiated to 1 / tau) values of the input
        distribution.
    """
    assert tau > 0
    assert min(distribution.values()) >= 0

    max_value = max(distribution.values())
    rescaled_distribution = {k: v / max_value for k, v in distribution.items()}

    total = sum(v ** (1 / tau) for v in rescaled_distribution.values())
    extremised_distribution = {k: (v ** (1 / tau) / total)
                               for k, v in rescaled_distribution.items()}

    return extremised_distribution