from votekit.elections.election_types.ranking.abstract_ranking import RankingElection
from votekit.pref_profile import RankProfile
from votekit.elections.election_state import ElectionState
from votekit.cleaning import remove_and_condense_rank_profile
from votekit.utils import (
    first_place_votes,
    score_dict_to_ranking,
)
import random
from typing import Literal
from functools import partial


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 (RankProfile): RankProfile to run election on.
      m (int): Number of seats to elect.
      fpv_tie_convention (Literal["high", "average", "low"], optional): How to award points
            for tied first place votes. Defaults to "average", where if n candidates are tied for
            first, each receives 1/n points. "high" would award them each one point, and "low" 0.
    """

    def __init__(
        self,
        profile: RankProfile,
        m: int,
        fpv_tie_convention: Literal["high", "average", "low"] = "average",
    ):
        if m <= 0:
            raise ValueError("m must be positive.")
        elif len(profile.candidates_cast) < m:
            raise ValueError("Not enough candidates received votes to be elected.")
        self.m = m
        super().__init__(
            profile,
            score_function=partial(
                first_place_votes, tie_convention=fpv_tie_convention
            ),
        )

    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: RankProfile, prev_state: ElectionState, store_states=False
    ) -> RankProfile:
        """
        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 (RankProfile): 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:
            RankProfile: The profile of ballots after the round is completed.
        """
        fpv = prev_state.scores
        candidates = list(fpv.keys())
        weights: list[float] = list(fpv.values())
        winning_cand = random.choices(candidates, weights=weights, k=1)[0]
        elected = (frozenset({winning_cand}),)

        new_profile = remove_and_condense_rank_profile(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={},
            )

            self.election_states.append(new_state)
        return new_profile
