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.utils import first_place_votes
from votekit.cleaning import remove_and_condense_rank_profile
from votekit.elections.election_types.ranking import Plurality
from typing import Optional, Literal
from functools import partial


class TopTwo(RankingElection):
    """
    Eliminates all but the top two plurality vote-getters, and then conducts a runoff between them,
    reallocating other ballots.

    Args:
        profile (RankProfile): Profile to conduct election on.
        tiebreak (str, optional): Tiebreak method to use. Options are None, 'random', and 'borda'.
            Defaults to None, in which case a tie raises a ValueError.
        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.
            Only used by ``score_function`` parameter.

    """

    def __init__(
        self,
        profile: RankProfile,
        tiebreak: Optional[str] = None,
        fpv_tie_convention: Literal["high", "low", "average"] = "average",
    ):
        self.tiebreak = tiebreak
        super().__init__(
            profile,
            score_function=partial(
                first_place_votes, tie_convention=fpv_tie_convention
            ),
            sort_high_low=True,
        )

    def _is_finished(self):
        # two round election
        if len(self.election_states) == 3:
            return True
        return False

    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.
        Choose top two candidates by first-place votes, then remove all other candidates
        from ballots and re-run to elect single winner by first-place votes.

        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.
        """
        if prev_state.round_number == 0:
            plurality = Plurality(profile, 2, self.tiebreak)

            remaining = plurality.get_elected()
            eliminated = plurality.get_remaining()
            tiebreaks = plurality.election_states[-1].tiebreaks
            new_profile: RankProfile = remove_and_condense_rank_profile(
                [c for s in eliminated for c in s], profile
            )

            if self.score_function is None:
                raise ValueError(
                    "score_function must be set for TopTwo elections when recording states"
                )

            scores = self.score_function(new_profile)

            if store_states:
                new_state = ElectionState(
                    round_number=prev_state.round_number + 1,
                    remaining=remaining,
                    eliminated=eliminated,
                    scores=scores,
                    tiebreaks=tiebreaks,
                )

                self.election_states.append(new_state)

        else:
            plurality = Plurality(profile, 1, self.tiebreak)
            new_profile = plurality.get_profile()

            if store_states:
                # first state was already stored by round 1 plurality
                # need to update round numbers
                plurality.election_states[1].round_number = 2
                self.election_states.append(plurality.election_states[1])

        return new_profile
