from .abstract_ranking import RankingElection
from ....pref_profile import PreferenceProfile
from ...election_state import ElectionState
from ....utils import first_place_votes, remove_cand
from ..ranking import Plurality
from typing import Optional


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

    Args:
        profile (PreferenceProfile): 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.

    """

    def __init__(self, profile: PreferenceProfile, tiebreak: Optional[str] = None):
        self.tiebreak = tiebreak
        super().__init__(profile, score_function=first_place_votes, 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: PreferenceProfile, prev_state: ElectionState, store_states=False
    ) -> PreferenceProfile:
        """
        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 (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.
        """
        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 = remove_cand([c for s in eliminated for c in s], profile)
            if self.score_function:
                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
