from .abstract_ranking import RankingElection
from ....pref_profile import PreferenceProfile
from ...election_state import ElectionState
from ....utils import (
    first_place_votes,
    remove_cand,
    score_dict_to_ranking,
    tiebreak_set,
)
import random


class RandomDictator(RankingElection):
    """
    Choose a winner randomly from the distribution of first place votes. For multi-winner elections
    repeat this process for every winner, removing that candidate from every voter's ballot
    once they have been elected.

    Args:
      profile (PreferenceProfile): PreferenceProfile to run election on.
      m (int): Number of seats to elect.
    """

    def __init__(self, profile: PreferenceProfile, m: int):
        if m <= 0:
            raise ValueError("m must be positive.")
        elif m > len(profile.candidates):
            raise ValueError(
                "m must be less than or equal to the number of candidates."
            )

        self.m = m
        super().__init__(profile, score_function=first_place_votes)

    def _is_finished(self) -> bool:
        cands_elected = [len(s) for s in self.get_elected()]
        return sum(cands_elected) >= self.m

    def _run_step(
        self, profile: PreferenceProfile, prev_state: ElectionState, store_states=False
    ) -> PreferenceProfile:
        """
        Run one step of an election from the given profile and previous state.
        If m candidates have not yet been elected:
        finds a single winning candidate to add to the list of elected
        candidates by sampling from the distribution of first place votes.
        Removes that candidate from all ballots in the preference profile.

        Args:
            profile (PreferenceProfile): Profile of ballots.
            prev_state (ElectionState): The previous ElectionState.
            store_states (bool, optional): True if `self.election_states` should be updated with the
                ElectionState generated by this round. This should only be True when used by
                `self._run_election()`. Defaults to False.

        Returns:
            PreferenceProfile: The profile of ballots after the round is completed.
        """

        ballots = profile.ballots
        weights = [b.weight for b in ballots]
        random_ballot = random.choices(ballots, weights=weights, k=1)[0]

        if not random_ballot.ranking:
            return PreferenceProfile()

        if len(random_ballot.ranking[0]) > 1:
            tiebroken_ranking = tiebreak_set(
                random_ballot.ranking[0], tiebreak="random"
            )
            tiebreaks = {random_ballot.ranking[0]: tiebroken_ranking}

        else:
            tiebroken_ranking = (random_ballot.ranking[0],)
            tiebreaks = {}

        winning_cand = list(tiebroken_ranking[0])[0]
        elected = (frozenset({winning_cand}),)

        new_profile = remove_cand(winning_cand, profile)

        if store_states:
            if self.score_function:
                scores = self.score_function(new_profile)
            remaining = score_dict_to_ranking(scores)

            new_state = ElectionState(
                round_number=prev_state.round_number + 1,
                elected=elected,
                remaining=remaining,
                scores=scores,
                tiebreaks=tiebreaks,
            )

            self.election_states.append(new_state)
        return new_profile
