import math
import operator
import random


class Truncation:
    def __init__(self, sample_from_percent=0.2, resample_if_not_in_percent=0.8):
        """
        A simple truncation mechanism as specified in (Jaderberg et al., 2017)
        :param sample_from_percent: Percent of best part of the population
        :param resample_if_not_in_percent: Percent of untouched agents
        """
        self.sample_from_percent = sample_from_percent
        self.resample_if_not_in_percent = resample_if_not_in_percent

    def __call__(self, own_name: str, scores: dict) -> str:
        """
        Find a better agent or return own_name, if the agent is good enough.
        :param own_name: The agent of the current agent
        :param scores: A dict with names and scores of all agents
        :return: The name of the chosen better agent
        """
        if len(scores) == 1:
            return own_name
        if own_name in self._get_best(self.resample_if_not_in_percent, scores):
            return own_name
        else:
            return random.choice(list(
                self._get_best(self.sample_from_percent, scores).keys()))

    def _get_best(self, percent, scores):
        sorted_scores = sorted(
            scores.items(), reverse=True, key=operator.itemgetter(1))
        last_index = math.ceil((len(scores) - 1) * percent)
        return {name: score for name, score in sorted_scores[:last_index]}
