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 elect_cands_from_set_ranking, borda_scores
from votekit.graphs import PairwiseComparisonGraph


class CondoBorda(RankingElection):
    """
    Just like DominatingSets (Smith method), but user gets to choose the number of winners,
    :math:`m`.  Ties are broken with Borda scores.

    Args:
        profile (RankProfile): Profile to conduct election on.
        m (int, optional): Number of seats to elect. Defaults to 1.

    """

    def __init__(self, profile: RankProfile, m: int = 1):
        if len(profile.candidates_cast) < m:
            raise ValueError("Not enough candidates received votes to be elected.")
        self.m = m
        super().__init__(profile, score_function=borda_scores)

    def _is_finished(self):
        # single round election
        if len(self.election_states) == 2:
            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.
        Compute the dominating tiers, and return the top :math:`m` candidates by tier.
        Breaks ties using Borda scores.

        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.
        """
        pwc_graph = PairwiseComparisonGraph(profile)
        dominating_tiers = pwc_graph.get_dominating_tiers()

        dt_ranking = tuple([frozenset(s) for s in dominating_tiers])
        elected, remaining, tiebreak = elect_cands_from_set_ranking(
            dt_ranking, self.m, profile, tiebreak="borda"
        )

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

        new_profile = remove_and_condense_rank_profile(
            [c for s in elected for c in s], profile
        )

        if store_states:
            self.election_states.append(
                ElectionState(
                    round_number=prev_state.round_number + 1,
                    elected=elected,
                    remaining=remaining,
                    tiebreaks=tiebreaks,
                    scores=borda_scores(new_profile),
                )
            )

        return new_profile
