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


class Plurality(RankingElection):
    """
    Plurality election. Winners are the m candidates with the most first-place votes.

    Args:
        profile (PreferenceProfile): Profile to conduct election on.
        m (int, optional): Number of seats to elect. Defaults to 1.
        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, m: int = 1, tiebreak: Optional[str] = None
    ):
        self.m = m
        self.tiebreak = tiebreak
        super().__init__(profile, score_function=first_place_votes, sort_high_low=True)

    def _is_finished(self):
        # single round election
        if len(self.election_states) == 2:
            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.
        In a Plurality election, find the :math:`m` candidates with the highest first-place vote
        totals.

        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.
        """
        # since fpv is the score function, the remaining cands from round 0
        # are ranked by fpv
        # raises a ValueError is tiebreak is None and a tie occurs.
        elected, remaining, tie_resolution = elect_cands_from_set_ranking(
            prev_state.remaining, self.m, profile=profile, tiebreak=self.tiebreak
        )

        new_profile = remove_cand([c for s in elected for c in s], profile)
        if store_states:
            if self.score_function:
                scores = self.score_function(new_profile)
            else:
                raise ValueError()

            if tie_resolution:
                tiebreaks = {tie_resolution[0]: tie_resolution[1]}
            else:
                tiebreaks = {}

            new_state = ElectionState(
                round_number=1,  # single shot election
                remaining=remaining,
                elected=elected,
                scores=scores,
                tiebreaks=tiebreaks,
            )

            self.election_states.append(new_state)

        return new_profile


class SNTV(Plurality):
    """
    Wrapper for Plurality election. Winners are the m candidates with the most first-place votes.

    Args:
        profile (PreferenceProfile): Profile to conduct election on.
        m (int, optional): Number of seats to elect. Defaults to 1.
        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, m: int = 1, tiebreak: Optional[str] = None
    ):
        super().__init__(profile, m, tiebreak)
